ruststream 0.7.5

//! The reply wiring a mount site's chain builds, driven end to end on the in-memory broker: the
//! codec the chain names encodes the reply, the transform it composes stamps it, the publisher's
//! own base headers reach what leaves through it, and `.transactional()` puts a batch's replies
//! in one broker transaction.
#![cfg(all(
    feature = "testing",
    feature = "macros",
    feature = "memory",
    feature = "json",
    feature = "cbor"
))]

use std::future::{Future, ready};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use ruststream::codec::{CborCodec, Codec};
use ruststream::memory::prelude::*;
use ruststream::memory::{ConnectedMemoryBroker, LogMode, MemoryPosition, MemoryPublisher};
use ruststream::runtime::{ContextKind, Outgoing, PublishTransform, Reads, for_batch};
use ruststream::testing::TestApp;
use ruststream::{BytesMut, HeaderMap, OutgoingMessage, PairError, PublishPolicy, Take};
use serde::{Deserialize, Serialize};

#[derive(Debug, Deserialize, Outgoing, Serialize, PartialEq)]
struct Order {
    id: u64,
}

/// The reply, which the slot case below also publishes by hand: it declares no destination of its
/// own, so the mount site's `reply("..")` names one and the slot's builder names the other.
#[derive(Debug, Deserialize, Outgoing, Serialize, PartialEq)]
struct Receipt {
    id: u64,
}

/// Stamps a provenance header, so a reply that went through the chain's transform is
/// distinguishable from one that did not.
struct Stamp;

impl<K: ContextKind, Options> PublishTransform<K, Options> for Stamp {
    type Destination = Reads;

    fn apply(&self, out: &mut Outgoing<'_>, _options: &mut Option<Options>, _cx: &K::View<'_>) {
        out.headers_mut().insert("x-stamped", b"1".to_vec());
    }
}

#[subscriber("codec.in", reply("codec.out"))]
async fn encode_reply(order: &Order) -> Receipt {
    Receipt { id: order.id }
}

/// `.codec(..)` names the reply codec: the reply leaves in CBOR while the request still arrives
/// under the scope's default (JSON).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_named_codec_encodes_the_reply() {
    let app = RustStream::new(AppInfo::new("reply-wiring", "0.1.0")).with_broker(
        MemoryBroker::new(),
        |b| {
            b.include(encode_reply).out(Reply, Publish).codec(CborCodec);
        },
    );
    let tb = TestApp::start(app).await.expect("harness start");

    tb.message(&Order { id: 7 })
        .to("codec.in")
        .publish()
        .await
        .expect("publish");

    tb.broker::<MemoryBroker>()
        .subscriber("codec.in")
        .assert_called_once()
        .with(&Order { id: 7 })
        .settled(HandlerOutcome::ack());

    let published = tb
        .broker::<MemoryBroker>()
        .published::<Receipt>("codec.out");
    let published = published.assert_called_once();
    let payload = published.messages()[0].payload();
    let decoded: Receipt = CborCodec
        .decode(payload)
        .expect("the reply must decode with the codec the chain named");
    assert_eq!(decoded, Receipt { id: 7 });
    assert!(
        serde_json::from_slice::<Receipt>(payload).is_err(),
        "the default codec must not have encoded this reply",
    );
}

/// The header a broker publisher carrying a delivery option for a run of messages contributes to
/// everything that leaves through it: a lane, a partition key, a tenant.
const LANE: &str = "x-lane";

/// A publish policy whose live publisher carries an argument for every message it sends: the
/// shape a broker takes when its option rides the headers rather than the message body.
#[derive(Debug, Clone, Copy)]
struct LanePublish;

impl PublishPolicy<ConnectedMemoryBroker> for LanePublish {
    type Live = Laned;

    fn pair(
        self,
        connected: &ConnectedMemoryBroker,
    ) -> impl Future<Output = Result<Laned, PairError>> {
        let mut base = HeaderMap::new();
        base.insert(LANE, "west");
        ready(Ok(Laned {
            inner: connected.publisher(),
            base,
        }))
    }
}

/// The live half of [`LanePublish`]: the broker's own publisher plus the base it contributes.
struct Laned {
    inner: MemoryPublisher,
    base: HeaderMap,
}

impl Publisher for Laned {
    type Payload = Take;
    type Error = MemoryError;
    type Options = ();

    async fn publish(
        &self,
        msg: OutgoingMessage<'_, BytesMut>,
        options: Option<&Self::Options>,
    ) -> Result<(), Self::Error> {
        self.inner.publish(msg, options).await
    }

    fn base_headers(&self) -> Option<&HeaderMap> {
        Some(&self.base)
    }
}

#[subscriber("lane.in", reply("lane.out"))]
async fn laned_reply(order: &Order) -> Receipt {
    Receipt { id: order.id }
}

/// A reply leaves through the publisher the chain named, so the base that publisher contributes
/// is on it - and the chain's own transform still writes over that base.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_publishers_base_headers_reach_the_reply() {
    let app = RustStream::new(AppInfo::new("reply-wiring", "0.1.0")).with_broker(
        MemoryBroker::new(),
        |b| {
            b.include(laned_reply)
                .out(Reply, LanePublish)
                .transform(Stamp);
        },
    );
    let tb = TestApp::start(app).await.expect("harness start");

    tb.message(&Order { id: 5 })
        .to("lane.in")
        .publish()
        .await
        .expect("publish");

    tb.broker::<MemoryBroker>()
        .published::<Receipt>("lane.out")
        .assert_called_once()
        .with(&Receipt { id: 5 })
        .with_header(LANE, "west")
        .with_header("x-stamped", b"1");
}

#[subscriber("lane.batch.in", reply("lane.batch.out"))]
async fn laned_batch(orders: &[Order]) -> Vec<Receipt> {
    orders
        .iter()
        .map(|order| Receipt { id: order.id })
        .collect()
}

/// The batch reply path builds its own outgoing message too, so the base has to be on those
/// replies as well.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_publishers_base_headers_reach_a_batch_reply() {
    let app = RustStream::new(AppInfo::new("reply-wiring", "0.1.0")).with_broker(
        MemoryBroker::new(),
        |b| {
            b.include(laned_batch.batch(nonzero!(8)))
                .out(Reply, LanePublish);
        },
    );
    let tb = TestApp::start(app).await.expect("harness start");

    tb.message(&Order { id: 6 })
        .to("lane.batch.in")
        .publish()
        .await
        .expect("publish");

    tb.broker::<MemoryBroker>()
        .published::<Receipt>("lane.batch.out")
        .assert_called_once()
        .with(&Receipt { id: 6 })
        .with_header(LANE, "west");
}

#[subscriber("lane.slot.in")]
async fn laned_slot(order: &Order, Out(out): Out<impl Publisher>) -> HandlerOutcome {
    if out
        .message(&Receipt { id: order.id })
        .to("lane.slot.out")
        .publish()
        .await
        .is_err()
    {
        return HandlerOutcome::retry();
    }
    HandlerOutcome::ack()
}

/// A slot publish reaches the same base through the publish builder the body writes, so what
/// leaves a slot carries the bound publisher's argument like a reply does.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_publishers_base_headers_reach_a_slot_publish() {
    let app = RustStream::new(AppInfo::new("reply-wiring", "0.1.0")).with_broker(
        MemoryBroker::new(),
        |b| {
            b.include(laned_slot).out(DefaultSlot, LanePublish).build();
        },
    );
    let tb = TestApp::start(app).await.expect("harness start");

    tb.message(&Order { id: 8 })
        .to("lane.slot.in")
        .publish()
        .await
        .expect("publish");

    tb.broker::<MemoryBroker>()
        .published::<Receipt>("lane.slot.out")
        .assert_called_once()
        .with(&Receipt { id: 8 })
        .with_header(LANE, "west");
}

#[subscriber("batch.in", reply("batch.out"))]
async fn confirm_batch(orders: &[Order]) -> Vec<Receipt> {
    orders
        .iter()
        .map(|order| Receipt { id: order.id })
        .collect()
}

/// A transactional publisher that marks every message with whether a transaction was open when
/// it was published, so a test tells a reply sent inside the batch's transaction from one sent
/// outside it.
struct Journaled {
    inner: MemoryPublisher,
    open: AtomicBool,
    /// How many transactions it began, shared with the test through its policy.
    begins: Arc<AtomicUsize>,
}

impl Publisher for Journaled {
    type Payload = Take;
    type Error = MemoryError;
    type Options = ();

    async fn publish(
        &self,
        msg: OutgoingMessage<'_, BytesMut>,
        options: Option<&Self::Options>,
    ) -> Result<(), Self::Error> {
        let open = if self.open.load(Ordering::SeqCst) {
            "open"
        } else {
            "none"
        };
        let mut headers = msg.headers().clone();
        headers.insert("x-transaction", open);
        let marked = OutgoingMessage::new(msg.name(), msg.payload()).with_headers(headers);
        self.inner.publish(marked, options).await
    }
}

impl TransactionalPublisher for Journaled {
    async fn begin_transaction(&self) -> Result<(), Self::Error> {
        self.inner.begin_transaction().await?;
        self.begins.fetch_add(1, Ordering::SeqCst);
        self.open.store(true, Ordering::SeqCst);
        Ok(())
    }

    async fn commit(&self) -> Result<(), Self::Error> {
        self.open.store(false, Ordering::SeqCst);
        self.inner.commit().await
    }

    async fn abort(&self) -> Result<(), Self::Error> {
        self.open.store(false, Ordering::SeqCst);
        self.inner.abort().await
    }
}

/// The policy half of [`Journaled`].
#[derive(Default)]
struct JournaledPublish {
    begins: Arc<AtomicUsize>,
}

impl<Log: LogMode> PublishPolicy<ConnectedMemoryBroker<Log>> for JournaledPublish {
    type Live = Journaled;

    fn pair(
        self,
        connected: &ConnectedMemoryBroker<Log>,
    ) -> impl Future<Output = Result<Journaled, PairError>> {
        ready(Ok(Journaled {
            inner: connected.publisher(),
            open: AtomicBool::new(false),
            begins: self.begins,
        }))
    }
}

/// `.transactional()` publishes a batch's replies inside one broker transaction, committed once
/// they are all sent: a batch of three replies opens one transaction, and every reply leaves
/// inside it.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_batch_reply_commits_its_transaction() {
    let broker = MemoryBroker::retaining(Retention::Messages(nonzero!(8)));
    let publisher = broker.publisher();
    for id in [11u64, 12, 13] {
        publisher
            .message(&Order { id })
            .to("batch.in")
            .publish()
            .await
            .expect("publish");
    }
    let policy = JournaledPublish::default();
    let begins = Arc::clone(&policy.begins);
    let app = RustStream::new(AppInfo::new("reply-wiring", "0.1.0")).with_broker(broker, |b| {
        b.include(
            confirm_batch
                .batch(nonzero!(8))
                .start_at(MemoryPosition::start()),
        )
        .out(Reply, policy)
        .transactional();
    });
    let tb = TestApp::start(app).await.expect("harness start");
    tb.settle().await.expect("the replayed batch settles");

    let replies = tb
        .broker::<MemoryBroker<Retaining>>()
        .published::<Receipt>("batch.out")
        .assert_called(3);
    let ids: Vec<u64> = replies.decoded().iter().map(|reply| reply.id).collect();
    assert_eq!(ids, vec![11, 12, 13]);
    assert!(
        replies
            .messages()
            .iter()
            .all(|reply| reply.headers().get_str("x-transaction") == Some("open")),
        "every reply must leave inside the batch's transaction"
    );
    assert_eq!(
        begins.load(Ordering::SeqCst),
        1,
        "the batch's replies must share one transaction"
    );
}

#[derive(OutSlot)]
#[publishes(Receipt)]
struct Journal;

#[subscriber("both.in", reply("both.out"))]
async fn confirm_and_journal(
    orders: &[Order],
    Out(journal): Out<impl Publisher, Journal>,
) -> Vec<Receipt> {
    for order in orders {
        let _ = journal
            .message(&Receipt { id: order.id })
            .to("both.journal")
            .publish()
            .await;
    }
    orders
        .iter()
        .map(|order| Receipt { id: order.id })
        .collect()
}

/// A registration that carries both positions commits with `.build()`, and the reply-only steps
/// still ride the reply: the batch transform stamps every reply and the transaction wraps them,
/// while the slot beside it takes its own policy.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_reply_and_slot_chain_keeps_the_reply_only_steps() {
    let app = RustStream::new(AppInfo::new("reply-wiring", "0.1.0")).with_broker(
        MemoryBroker::new(),
        |b| {
            b.include(confirm_and_journal.batch(nonzero!(8)))
                .out(Reply, TransactionalPublish)
                .batch_transform(for_batch(Stamp))
                .transactional()
                .out(Journal, Publish)
                .build();
        },
    );
    let tb = TestApp::start(app).await.expect("harness start");

    tb.message(&Order { id: 12 })
        .to("both.in")
        .publish()
        .await
        .expect("publish");

    let handle = tb.broker::<MemoryBroker>();
    handle
        .published::<Receipt>("both.out")
        .assert_called_once()
        .with(&Receipt { id: 12 })
        .with_header("x-stamped", b"1");
    // The slot publish is the body's own, so the reply's batch transform never touches it.
    let journalled = handle.published::<Receipt>("both.journal");
    let journalled = journalled.assert_called_once().with(&Receipt { id: 12 });
    assert!(
        journalled.messages()[0]
            .headers()
            .get("x-stamped")
            .is_none(),
        "a batch transform rides the reply position, not a slot",
    );
}