cortexkit-bus-inmemory 0.3.2

use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use cortexkit_bus_naming::AccountNames;
use cortexkit_bus_trait::{
    AsyncPermissionViolation, BusError, BusResult, DeliveryToken, DurableOwner, Message,
    RegisterUpdate, RegisterValue, Revision,
};
use tokio::sync::mpsc;

#[derive(Debug, Clone)]
pub struct WorkQueueConfig {
    pub subject: String,
    pub max_bytes: usize,
    pub max_deliveries: u32,
    pub max_nak_delay: Duration,
}

impl WorkQueueConfig {
    pub fn new(subject: impl Into<String>, max_bytes: usize) -> Self {
        Self {
            subject: subject.into(),
            max_bytes,
            max_deliveries: 5,
            max_nak_delay: Duration::from_secs(30),
        }
    }
}

/// The stream durables' ack wait unless a test sets another: the plane's
/// consumers use 30 seconds.
pub const DEFAULT_STREAM_ACK_WAIT: Duration = Duration::from_secs(30);

/// JetStream's default message-id deduplication window.
pub const DEFAULT_DEDUPLICATION_WINDOW: Duration = Duration::from_secs(120);

#[derive(Debug, Clone)]
pub struct InMemoryConfig {
    /// Durable names, `c_{agent_id}` or `m_{module_id}`.
    pub(crate) durables: BTreeSet<String>,
    pub(crate) denied_subjects: BTreeSet<String>,
    pub(crate) work_queue: Option<WorkQueueConfig>,
    pub(crate) stream_ack_wait: Duration,
    pub(crate) deduplication_window: Duration,
}

impl Default for InMemoryConfig {
    fn default() -> Self {
        Self {
            durables: BTreeSet::new(),
            denied_subjects: BTreeSet::new(),
            work_queue: None,
            stream_ack_wait: DEFAULT_STREAM_ACK_WAIT,
            deduplication_window: DEFAULT_DEDUPLICATION_WINDOW,
        }
    }
}

impl InMemoryConfig {
    /// Sets the message-id deduplication window for this bus's single stream.
    /// Use the production stream's window to reproduce its retry behavior.
    pub fn with_deduplication_window(mut self, window: Duration) -> Self {
        self.deduplication_window = window;
        self
    }

    /// How long a stream delivery may stay unsettled before another cursor on
    /// the same durable receives it again; applies to every durable.
    pub fn with_stream_ack_wait(mut self, ack_wait: Duration) -> Self {
        self.stream_ack_wait = ack_wait;
        self
    }

    /// Creates the owner's durable, as ck-bus would before any cursor binds.
    ///
    /// # Panics
    ///
    /// When the owner's id is not a valid agent or module id token: the
    /// fixture would otherwise hold a durable no cursor can bind.
    pub fn with_durable(mut self, owner: DurableOwner<'_>) -> Self {
        let name = durable_name(owner).expect("fixture durable owner must be a valid id token");
        self.durables.insert(name);
        self
    }

    pub fn deny_subject(mut self, subject: impl Into<String>) -> Self {
        self.denied_subjects.insert(subject.into());
        self
    }

    pub fn with_work_queue(mut self, config: WorkQueueConfig) -> Self {
        self.work_queue = Some(config);
        self
    }
}

/// The owner's durable name, `c_{agent_id}` or `m_{module_id}`, from
/// `AccountNames::consumer_name` and `AccountNames::module_consumer_name`.
/// Both refuse an id that is not one plain subject token, since a dot or a
/// wildcard in it would address some other consumer.
pub(crate) fn durable_name(owner: DurableOwner<'_>) -> BusResult<String> {
    let (id, name) = match owner {
        DurableOwner::Agent(agent_id) => (agent_id, AccountNames::consumer_name(agent_id)),
        DurableOwner::Module(module_id) => {
            (module_id, AccountNames::module_consumer_name(module_id))
        }
    };
    name.map_err(|error| BusError::denied(id, error.to_string()))
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BackendEvent {
    Published {
        subject: String,
        id: String,
    },
    Acknowledged {
        token: DeliveryToken,
    },
    NegativelyAcknowledged {
        token: DeliveryToken,
    },
    Terminated {
        token: DeliveryToken,
    },
    /// A stream delivery settled with `StreamCursor::term`.
    StreamTerminated {
        stream_seq: u64,
        reason: Option<String>,
    },
}

#[derive(Clone)]
pub struct InMemoryBus {
    pub(crate) inner: Arc<Mutex<State>>,
    pub(crate) deduplication_clock: Arc<dyn Fn() -> Instant + Send + Sync>,
    pub(crate) stream_clock: Arc<dyn Fn() -> Instant + Send + Sync>,
}

impl InMemoryBus {
    pub fn new(config: InMemoryConfig) -> Self {
        Self::new_with_deduplication_clock(config, Instant::now)
    }

    /// Injects a monotonic clock for deduplication without waiting in tests.
    /// Delivery leases and work-queue delays still use the real clock.
    /// Each bus has one stream; clones share its ids, while separate buses do not.
    pub fn new_with_deduplication_clock(
        config: InMemoryConfig,
        clock: impl Fn() -> Instant + Send + Sync + 'static,
    ) -> Self {
        let durable_cursors = config
            .durables
            .iter()
            .cloned()
            .map(|key| (key, DurableState::default()))
            .collect();
        Self {
            inner: Arc::new(Mutex::new(State {
                config,
                next_stream_seq: 1,
                stream_messages: Vec::new(),
                seen_ids: HashMap::new(),
                durable_cursors,
                register_revision: 0,
                register_values: BTreeMap::new(),
                watchers: Vec::new(),
                queue: VecDeque::new(),
                next_delivery_token: 1,
                events: Vec::new(),
                permission_violations: VecDeque::new(),
            })),
            deduplication_clock: Arc::new(clock),
            stream_clock: Arc::new(Instant::now),
        }
    }

    pub fn events(&self) -> Vec<BackendEvent> {
        self.inner
            .lock()
            .expect("in-memory bus poisoned")
            .events
            .clone()
    }

    pub fn queued_bytes(&self) -> usize {
        self.inner
            .lock()
            .expect("in-memory bus poisoned")
            .queue
            .iter()
            .map(|entry| entry.message.wire_size())
            .sum()
    }

    pub fn inject_permission_violation(&self, subject: impl Into<String>) {
        let subject = subject.into();
        self.inner
            .lock()
            .expect("in-memory bus poisoned")
            .permission_violations
            .push_back(AsyncPermissionViolation {
                subject,
                reason: "fixture-configured permission violation".into(),
            });
    }
}

pub(crate) struct State {
    pub config: InMemoryConfig,
    pub next_stream_seq: u64,
    pub stream_messages: Vec<StoredMessage>,
    pub seen_ids: HashMap<String, (u64, Instant)>,
    pub durable_cursors: HashMap<String, DurableState>,
    pub register_revision: u64,
    pub register_values: BTreeMap<String, RegisterValue>,
    pub watchers: Vec<Watcher>,
    pub queue: VecDeque<QueueEntry>,
    pub next_delivery_token: u64,
    pub events: Vec<BackendEvent>,
    pub permission_violations: VecDeque<AsyncPermissionViolation>,
}

#[derive(Clone)]
pub(crate) struct StoredMessage {
    pub stream_seq: u64,
    pub message: Message,
}

/// One durable's position. Delivery is strictly in order: the message at
/// `next_index` is the only one a cursor can receive until it is acked or
/// terminated.
#[derive(Default)]
pub(crate) struct DurableState {
    pub next_index: usize,
    /// Set by a nak: no cursor receives the head message before this instant.
    pub available_at: Option<Instant>,
    /// How many times the head message has been delivered.
    pub delivery_count: u64,
    /// The cursor holding the head message unsettled, if any.
    pub lease: Option<Lease>,
    pub next_cursor_id: u64,
}

/// An unsettled delivery of the head message. Until `deadline` no other
/// cursor receives the message; after it the next `next` on any cursor
/// redelivers it, as JetStream does once the ack wait expires.
pub(crate) struct Lease {
    pub cursor_id: u64,
    pub deadline: Instant,
}

pub(crate) struct Watcher {
    pub prefix: String,
    pub sender: mpsc::UnboundedSender<RegisterUpdate>,
}

pub(crate) struct QueueEntry {
    pub token: DeliveryToken,
    pub message: Message,
    pub delivery_count: u32,
    pub available_at: Instant,
    pub in_flight: bool,
}

impl State {
    pub fn next_revision(&mut self) -> Revision {
        self.register_revision += 1;
        Revision(self.register_revision)
    }

    pub fn notify_watchers(&mut self, update: RegisterUpdate) {
        let key = match &update {
            RegisterUpdate::Put(entry) => &entry.key,
            RegisterUpdate::Delete(tombstone) => &tombstone.key,
        };
        self.watchers.retain(|watcher| {
            !key.starts_with(&watcher.prefix) || watcher.sender.send(update.clone()).is_ok()
        });
    }
}