hyperhive/swarm-queue-client/src/router_todo.rs

//! The router's todo queue from both ends: the router creates the stream, and
//! each agent's harness creates and drains its own consumer. The names,
//! subjects and op are `swarm_wire::router_todo`.

use std::time::Duration;

use swarm_wire::router_todo::{
    DUPLICATE_WINDOW, MAX_AGE, MAX_MSGS_PER_SUBJECT, STREAM, consumer_name, subject,
    subject_wildcard,
};

/// How long a delivered op may stay unacked before the server delivers it
/// again.
const ACK_WAIT: Duration = Duration::from_secs(30);

/// The config [`open_or_create`] creates the stream with.
///
/// Work-queue retention: an op leaves the stream once its agent acks it. The
/// server then requires every consumer's filter to be disjoint from every
/// other's, which one consumer per agent subject is.
fn stream_config() -> async_nats::jetstream::stream::Config {
    use async_nats::jetstream::stream::{Config, DiscardPolicy, RetentionPolicy};

    Config {
        name: STREAM.to_owned(),
        description: Some("Notification router todo ops, one subject per agent".to_owned()),
        subjects: vec![subject_wildcard()],
        retention: RetentionPolicy::WorkQueue,
        max_age: MAX_AGE,
        max_messages_per_subject: MAX_MSGS_PER_SUBJECT,
        discard: DiscardPolicy::Old,
        duplicate_window: DUPLICATE_WINDOW,
        ..Default::default()
    }
}

/// Open the stream, creating it if it does not exist yet. The router is its
/// one writer and its one creator.
pub async fn open_or_create(
    client: &async_nats::Client,
) -> Result<async_nats::jetstream::stream::Stream, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    match js.get_stream(STREAM).await {
        Ok(stream) => Ok(stream),
        Err(e) => {
            tracing::info!(stream = STREAM, reason = %e, "router todo stream not available, creating it");
            js.create_stream(stream_config())
                .await
                .map_err(|source| crate::Error::CreateStream {
                    stream: STREAM,
                    source,
                })
        }
    }
}

/// The durable pull consumer `agent`'s harness drains: its own subject, from
/// the oldest op still queued, each acked explicitly.
fn consumer_config(agent: &str) -> async_nats::jetstream::consumer::pull::Config {
    use async_nats::jetstream::consumer::{AckPolicy, DeliverPolicy, pull::Config};

    Config {
        durable_name: Some(consumer_name(agent)),
        description: Some(format!("Router todo ops for {agent}")),
        filter_subject: subject(agent),
        deliver_policy: DeliverPolicy::All,
        ack_policy: AckPolicy::Explicit,
        ack_wait: ACK_WAIT,
        ..Default::default()
    }
}

/// Create `agent`'s consumer, or return it as it is when it already exists
/// with the same config.
///
/// Created on the stream by name, without looking the stream up first: the
/// agent's grant has no `STREAM.INFO`. The create subject is the extended
/// form ending in the consumer's filter, the one subject the grant names.
pub async fn open_own(
    client: &async_nats::Client,
    agent: &str,
) -> Result<
    async_nats::jetstream::consumer::Consumer<async_nats::jetstream::consumer::pull::Config>,
    async_nats::jetstream::stream::ConsumerError,
> {
    async_nats::jetstream::new(client.clone())
        .create_consumer_on_stream(consumer_config(agent), STREAM)
        .await
}

#[cfg(test)]
mod tests {
    use futures_util::StreamExt as _;
    use swarm_wire::router_todo::{STREAM, consumer_name, subject};

    use super::{consumer_config, open_or_create, open_own};
    use crate::test_nats::TestNats;

    async fn queue(js: &async_nats::jetstream::Context, agent: &str, id: &str) {
        js.send_publish(
            subject(agent),
            async_nats::jetstream::message::PublishMessage::build()
                .payload(id.as_bytes().to_vec().into())
                .message_id(id),
        )
        .await
        .expect("published")
        .await
        .expect("acked by the stream");
    }

    async fn queued(js: &async_nats::jetstream::Context) -> u64 {
        js.get_stream(STREAM)
            .await
            .expect("the stream")
            .info()
            .await
            .expect("its info")
            .state
            .messages
    }

    /// The server takes the stream and both agents' consumers, delivers each
    /// agent its own ops only, drops a republished op by its id, and removes an
    /// op once its agent acks it.
    #[tokio::test]
    async fn an_agents_consumer_drains_its_own_ops_and_an_ack_removes_them() {
        let nats = TestNats::start();
        let client = nats.client().await;
        let js = async_nats::jetstream::new(client.clone());
        open_or_create(&client).await.expect("created");
        open_or_create(&client).await.expect("opened as it is");

        queue(&js, "atlas", "1:assigned:sink:atlas").await;
        queue(&js, "atlas", "1:assigned:sink:atlas").await;
        queue(&js, "argus", "1:assigned:sink:argus").await;
        assert_eq!(queued(&js).await, 2, "the republished op is dropped");

        let atlas = open_own(&client, "atlas").await.expect("atlas's consumer");
        open_own(&client, "atlas")
            .await
            .expect("created again, unchanged");
        open_own(&client, "argus")
            .await
            .expect("argus's, disjoint from atlas's");

        let mut batch = atlas
            .fetch()
            .max_messages(10)
            .messages()
            .await
            .expect("a batch");
        let mut got = Vec::new();
        while let Some(message) = batch.next().await {
            let message = message.expect("a message");
            got.push(message.subject.to_string());
            message.double_ack().await.expect("acked");
        }
        assert_eq!(got, [subject("atlas")]);
        assert_eq!(queued(&js).await, 1, "argus's op waits for argus");
    }

    /// The grant's premise: in the extended create form the server refuses a
    /// payload filter that is not the subject's, so atlas's create subject
    /// cannot make a consumer that delivers argus's ops.
    #[tokio::test]
    async fn the_server_refuses_a_payload_filter_other_than_the_create_subjects() {
        let nats = TestNats::start();
        let client = nats.client().await;
        open_or_create(&client).await.expect("created");
        let create = |filter: String| {
            let mut config = consumer_config("atlas");
            config.filter_subject = filter;
            serde_json::to_vec(&serde_json::json!({
                "stream_name": STREAM,
                "config": async_nats::jetstream::consumer::IntoConsumerConfig::into_consumer_config(config),
                "action": "",
            }))
            .expect("encodes")
        };
        let subject_for_atlas = format!(
            "$JS.API.CONSUMER.CREATE.{STREAM}.{}.{}",
            consumer_name("atlas"),
            subject("atlas")
        );
        let answer = |body: Vec<u8>| {
            let client = client.clone();
            let subject_for_atlas = subject_for_atlas.clone();
            async move {
                let reply = client
                    .request(subject_for_atlas, body.into())
                    .await
                    .expect("an answer");
                serde_json::from_slice::<serde_json::Value>(&reply.payload).expect("json")
            }
        };
        let refused = answer(create(subject("argus"))).await;
        assert!(refused.get("error").is_some(), "{refused}");
        // Control: the matching filter is accepted on the same subject.
        let accepted = answer(create(subject("atlas"))).await;
        assert!(accepted.get("error").is_none(), "{accepted}");
    }
}