ruststream-rdkafka 0.7.1

//! Live end-to-end for the Avro registry middleware: a publishing handler's plain-JSON
//! replies get framed as Avro by the app's `SchemaFrame` publish layer, and a plain JSON
//! handler consumes them back through the broker's transcoding subscription.

#![cfg(feature = "avro")]

use std::convert::Infallible;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use ruststream::runtime::{
    App, AppInfo, HandlerOutcome, Reply, RustStream, State, SubscriberSettings as _,
};
use ruststream::{
    Broker, ConnectedBroker, FromRef, Outgoing, OutgoingMessage, Publisher, subscriber,
};
use ruststream_rdkafka::{
    KafkaBroker, KafkaPublish, KafkaTopic, SchemaFrame, SchemaRegistry, SchemaType, StartOffset,
};
use serde::{Deserialize, Serialize};
use tokio::sync::Notify;

mod live;

/// The fixed reply topic (the macro's `reply(..)` takes a string literal). Runs share it:
/// each pins its own uniquely-named subject on the `SchemaFrame`, so its messages carry a
/// schema id no other run has, and the probe filters deliveries by a marker id range.
const FRAMED_TOPIC: &str = "avro-mw-frames-placeholder";

// The relay republishes the order it read, so the type declares no topic of its own: the run's
// reply topic stays the mount site's.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Outgoing)]
struct Order {
    id: i64,
    item: String,
}

#[derive(Clone)]
struct AvroProbe {
    base: i64,
    expected: usize,
    seen: Arc<Mutex<Vec<Order>>>,
    done: Arc<Notify>,
}

#[derive(FromRef)]
struct AvroApp {
    probe: AvroProbe,
}

// The producing side: plain JSON through the pipeline; the SchemaFrame layer transcodes the
// reply to an Avro datum under this run's subject.
#[subscriber(KafkaTopic, reply("avro-mw-frames-placeholder"))]
async fn avro_relay(order: &Order) -> Order {
    order.clone()
}

// The consuming side: a plain handler on the default JSON codec; the subscription already
// converted the Avro datum back to JSON.
#[subscriber(KafkaTopic)]
async fn avro_mw(order: &Order, State(probe): State<AvroProbe>) -> HandlerOutcome {
    let marker_range = probe.base..probe.base + i64::try_from(probe.expected).expect("small");
    if !marker_range.contains(&order.id) {
        return HandlerOutcome::ack(); // an earlier run's message on the shared topic
    }
    {
        let mut seen = probe.seen.lock().expect("seen mutex poisoned");
        seen.push(order.clone());
        if seen.len() < probe.expected {
            return HandlerOutcome::ack();
        }
    }
    probe.done.notify_one();
    HandlerOutcome::ack()
}

/// A subscription on this run's `topic` in `group`, from the start of the log.
fn from_earliest(topic: &str, group: String) -> KafkaTopic {
    KafkaTopic::new(topic)
        .group(group)
        .start(StartOffset::Earliest)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn live_avro_middleware_end_to_end() {
    const COUNT: usize = 3;
    static NEXT: AtomicU64 = AtomicU64::new(0);

    let Some(registry) = live::url("SCHEMA_REGISTRY_TEST_URL") else {
        return;
    };
    let Some(kafka) = live::url("KAFKA_TEST_URL") else {
        return;
    };
    let run = format!(
        "{}x{}",
        std::process::id(),
        NEXT.fetch_add(1, Ordering::Relaxed)
    );
    let base = i64::from(std::process::id()) * 100_000;
    let trigger = format!("avro-mw-trigger-{run}");

    // A run-unique subject with a run-unique record name: its schema id belongs to this run
    // alone, so seeing it back proves the layer framed (a fresh subject also sidesteps the
    // registry's compatibility check between runs).
    let subject = format!("avro-mw-{run}");
    let schema = format!(
        r#"{{"type":"record","name":"Order{base}","fields":[
            {{"name":"id","type":"long"}},{{"name":"item","type":"string"}}]}}"#,
    );
    let sr = SchemaRegistry::new(&registry);
    let id = sr
        .register(&subject, SchemaType::Avro, schema)
        .await
        .expect("register");

    // Seed the trigger topic with plain JSON before the app starts.
    let seed_broker = KafkaBroker::new([kafka.clone()])
        .connect()
        .await
        .expect("connect seed");
    let count = i64::try_from(COUNT).expect("small count");
    for n in 0..count {
        let json = format!(r#"{{"id":{},"item":"item-{n}"}}"#, base + n);
        seed_broker
            .publisher(KafkaPublish::default())
            .publish(OutgoingMessage::new(&trigger, json.as_bytes()), None)
            .await
            .expect("seed trigger");
    }
    seed_broker.shutdown().await.expect("seed shutdown");

    // One app, both edges cold: the SchemaFrame resolves the pinned subject lazily on the
    // first reply, and the consuming subscription fetches the schema by wire id.
    let consumer_sr = SchemaRegistry::new(&registry);
    let probe = AvroProbe {
        base,
        expected: COUNT,
        seen: Arc::new(Mutex::new(Vec::new())),
        done: Arc::new(Notify::new()),
    };
    let app_probe = probe.clone();
    let broker = KafkaBroker::new([kafka]).schema_registry(consumer_sr.clone());
    let app = RustStream::new(AppInfo::new("avro-mw", "0.0.0"))
        .publish_layer(
            SchemaFrame::new(SchemaRegistry::new(&registry)).subject(FRAMED_TOPIC, &subject),
        )
        .on_startup(async move |()| Ok::<_, Infallible>(AvroApp { probe: app_probe }))
        .with_broker(broker, |b| {
            b.include(avro_mw.map_source(|_| from_earliest(FRAMED_TOPIC, format!("group-{run}"))));
            b.include(
                avro_relay.map_source(|_| from_earliest(&trigger, format!("trigger-group-{run}"))),
            )
            .out(Reply, KafkaPublish::default());
        });

    let done = Arc::clone(&probe.done);
    let wait = async move {
        tokio::time::timeout(std::time::Duration::from_secs(20), done.notified())
            .await
            .expect("all deliveries within timeout");
    };
    App::run_until(app, wait).await.expect("run");

    let mut seen = probe.seen.lock().expect("seen mutex poisoned").clone();
    seen.sort_by_key(|order| order.id);
    let expected: Vec<Order> = (0..count)
        .map(|n| Order {
            id: base + n,
            item: format!("item-{n}"),
        })
        .collect();
    assert_eq!(
        seen, expected,
        "Avro on the wire must arrive as JSON in plain handlers",
    );
    assert!(
        consumer_sr.cached_schema(id).is_some(),
        "the consumer transcoded through this run's schema id, so the layer really framed",
    );
}