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

//! An agent's own terminal stream: its creation and the consumer a reader
//! replays it through. The names, subjects, grants and replay start are
//! `swarm_wire::agent_term`.
//!
//! The stream captures exactly one literal subject, which
//! `swarm_wire::subagent_term`'s stream never overlaps. The `JetStream` server
//! refuses to create a stream whose subjects overlap another's.

use swarm_wire::agent_term::{IDLE_HEARTBEAT, MAX_AGE, stream_name, stream_subject};

/// How large one agent's stream may grow before the discard policy reclaims
/// space. The same cap as the subagent stream, for the same reason: it
/// bounds the disk one agent can claim.
const MAX_BYTES: i64 = 64 * 1024 * 1024;

/// The config [`open_or_create`] creates `agent`'s stream with.
///
/// `max_message_size` tracks the queue's live `max_payload`, as the subagent
/// stream's does, so the stream never rejects a row the queue accepted.
fn stream_config(agent: &str, max_message_size: i32) -> async_nats::jetstream::stream::Config {
    async_nats::jetstream::stream::Config {
        name: stream_name(agent),
        description: Some(format!("Terminal rows of {agent}")),
        subjects: vec![stream_subject(agent)],
        max_age: MAX_AGE,
        max_bytes: MAX_BYTES,
        discard: async_nats::jetstream::stream::DiscardPolicy::Old,
        max_message_size,
        ..Default::default()
    }
}

/// Open `agent`'s terminal stream, creating it if it does not exist yet. An
/// existing stream is opened as it is.
pub async fn open_or_create(
    client: &async_nats::Client,
    agent: &str,
) -> Result<async_nats::jetstream::stream::Stream, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    let name = stream_name(agent);
    match js.get_stream(&name).await {
        Ok(stream) => Ok(stream),
        Err(e) => {
            tracing::info!(stream = %name, reason = %e, "agent terminal stream not available, creating it");
            let max_message_size = i32::try_from(crate::max_payload(client)).unwrap_or(i32::MAX);
            js.create_stream(stream_config(agent, max_message_size))
                .await
                .map_err(|source| crate::Error::CreateAgentStream {
                    stream: name,
                    source,
                })
        }
    }
}

/// The consumer that replays the rows on `filter_subject` from
/// `start_sequence` and then follows new ones, delivered to `deliver_subject`
/// in stream order. The filter is an agent's [`stream_subject`] or one
/// subagent's `swarm_wire::subagent_term::subject`.
///
/// Ephemeral and unacknowledged: the server drops it once nothing listens
/// on `deliver_subject`, so a reader that goes away needs no delete grant.
/// `name` must be set, because the server only takes the filter from the
/// subject `swarm_wire::agent_term::consumer_create_grant` (or
/// `swarm_wire::subagent_term::consumer_create_grant`) covers when the create
/// names a consumer.
#[must_use]
pub fn replay_consumer_config(
    filter_subject: String,
    name: String,
    deliver_subject: String,
    start_sequence: u64,
) -> async_nats::jetstream::consumer::push::Config {
    async_nats::jetstream::consumer::push::Config {
        name: Some(name),
        deliver_subject,
        filter_subject,
        deliver_policy: async_nats::jetstream::consumer::DeliverPolicy::ByStartSequence {
            start_sequence,
        },
        ack_policy: async_nats::jetstream::consumer::AckPolicy::None,
        idle_heartbeat: IDLE_HEARTBEAT,
        ..Default::default()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The consumer starts at a stored sequence, not at the next new row, so
    /// the replay reaches the reader before the live tail, in stream order,
    /// through the one subscription.
    #[test]
    fn the_consumer_replays_before_it_follows() {
        let cfg = replay_consumer_config(
            stream_subject("iris"),
            "r1".to_owned(),
            "_INBOX.x".to_owned(),
            901,
        );
        assert_eq!(
            cfg.deliver_policy,
            async_nats::jetstream::consumer::DeliverPolicy::ByStartSequence {
                start_sequence: 901
            }
        );
        assert_eq!(cfg.filter_subject, "$SWARM.term.iris");
        assert_eq!(cfg.name.as_deref(), Some("r1"));
        assert_eq!(
            cfg.ack_policy,
            async_nats::jetstream::consumer::AckPolicy::None
        );
        assert!(cfg.durable_name.is_none());
        // Without it a dropped consumer is never noticed.
        assert_eq!(cfg.idle_heartbeat, IDLE_HEARTBEAT);
        assert!(!cfg.idle_heartbeat.is_zero());
    }

    /// Every limit is read from the subagent stream's own config, so the two
    /// fail this test the moment they drift apart.
    #[cfg(feature = "subagent-term")]
    #[test]
    fn stream_config_matches_the_subagent_streams_limits() {
        let main = stream_config("iris", 8_388_608);
        let sub = crate::subagent_term::stream_config("iris", 8_388_608);
        assert_eq!(main.subjects, ["$SWARM.term.iris"]);
        assert_eq!(main.max_age, sub.max_age);
        assert_eq!(main.max_bytes, sub.max_bytes);
        assert_eq!(main.discard, sub.discard);
        assert_eq!(main.max_message_size, sub.max_message_size);
        assert_eq!(main.retention, sub.retention);
        assert_eq!(main.storage, sub.storage);
        assert_eq!(main.num_replicas, sub.num_replicas);
    }
}