laser-sdk 0.8.0

use crate::agent::consumer::{AgentMessage, decode_agent_record};
use crate::error::LaserError;
use crate::types::MessageId;
use dashmap::DashMap;
use iggy::prelude::*;
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;
use tokio::sync::oneshot;

const REPLY_BATCH: u32 = 200;
const REPLY_POLL_INTERVAL: Duration = Duration::from_millis(20);
const PENDING_REPLY_LIMIT: usize = 64;

/// One shared reply consumer per `(stream, reply topic)`. It decodes each record
/// once and completes the one waiter whose correlation matches, so N concurrent
/// requests on the same reply topic read the topic once between them instead of
/// each running its own tail-seeded scan. Cached on the connection and driven by
/// a background task that stops when the last `Laser` clone drops.
#[derive(Clone)]
pub(crate) struct ReplyHub {
    inner: Arc<ReplyHubInner>,
}

struct ReplyHubInner {
    // Correlation -> the one-shot completing that request. A reply for an unknown
    // correlation (already completed, or never waited) is dropped.
    waiters: DashMap<String, ReplyWaiter>,
    // Woken by every subscribe, so an idle dispatcher parks instead of
    // polling the reply topic with nobody waiting.
    subscribed: tokio::sync::Notify,
    task: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
    #[cfg(feature = "sign")]
    verifier: Option<Arc<crate::sign::KeyRegistry>>,
}

struct ReplyWaiter {
    state: Mutex<ReplyWaiterState>,
    expected: ExpectedReply,
    #[cfg(feature = "sign")]
    expected_signer: Option<String>,
}

struct ReplyWaiterState {
    sender: Option<oneshot::Sender<AgentMessage>>,
    request: Option<Option<MessageId>>,
    pending: Vec<AgentMessage>,
}

// What a record must carry to answer one request: the request's session, a
// reply kind, and, when the request named its sender, that sender as the
// addressee.
pub(crate) struct ExpectedReply {
    pub(crate) session: crate::types::ConversationId,
    pub(crate) requester: Option<String>,
}

impl ExpectedReply {
    fn accepts(&self, message: &AgentMessage) -> bool {
        if message.provenance.conversation_id != self.session {
            return false;
        }
        let addressee = match &message.envelope {
            Some(envelope) => {
                if !matches!(
                    envelope.kind,
                    laser_wire::agent::AgentKind::Response | laser_wire::agent::AgentKind::Error
                ) {
                    return false;
                }
                envelope
                    .target
                    .as_ref()
                    .map(|target| target.as_str().to_owned())
            }
            None => message
                .provenance
                .target_agent_id
                .as_ref()
                .map(|target| target.as_str().to_owned()),
        };
        match (&self.requester, addressee) {
            (Some(requester), Some(addressee)) => *requester == addressee,
            (Some(_), None) => false,
            (None, _) => true,
        }
    }
}

impl Drop for ReplyHubInner {
    fn drop(&mut self) {
        if let Some(task) = self
            .task
            .lock()
            .expect("reply-hub task lock is not poisoned")
            .take()
        {
            task.abort();
        }
    }
}

impl ReplyHub {
    /// Create the hub and spawn its dispatcher, seeded at the reply topic's tail
    /// so it reads only records appended after creation. A waiter must
    /// [`subscribe`](Self::subscribe) before sending its request (registration is
    /// synchronous), so a reply, which cannot exist until the request is sent, is
    /// dispatched rather than missed.
    pub(crate) async fn create(
        client: Arc<IggyClient>,
        stream: String,
        topic: Identifier,
        #[cfg(feature = "sign")] verifier: Option<Arc<crate::sign::KeyRegistry>>,
    ) -> Result<Self, LaserError> {
        let stream_id = Identifier::named(&stream)?;
        let consumer = Consumer::new(Identifier::named("laser-reply-hub")?);
        let mut offsets: Vec<u64> = Vec::new();
        if let Some(details) = client.get_topic(&stream_id, &topic).await? {
            let partitions = crate::poll::bounded_partitions(details.partitions_count);
            offsets = vec![0u64; partitions as usize];
            for partition in 0..partitions {
                let polled = client
                    .poll_messages(
                        &stream_id,
                        &topic,
                        Some(partition),
                        &consumer,
                        &PollingStrategy::last(),
                        1,
                        false,
                    )
                    .await?;
                if let Some(last) = polled.messages.last() {
                    offsets[partition as usize] = last.header.offset.saturating_add(1);
                }
            }
        }
        let inner = Arc::new(ReplyHubInner {
            waiters: DashMap::new(),
            subscribed: tokio::sync::Notify::new(),
            task: std::sync::Mutex::new(None),
            #[cfg(feature = "sign")]
            verifier,
        });
        let handle = tokio::spawn(dispatch_loop(
            client,
            stream_id,
            topic,
            consumer,
            offsets,
            Arc::downgrade(&inner),
        ));
        *inner
            .task
            .lock()
            .expect("reply-hub task lock is not poisoned") = Some(handle);
        Ok(Self { inner })
    }

    /// Register a waiter for `correlation`, returning a ticket. Synchronous, so a
    /// caller registers before sending the request the reply answers.
    pub(crate) fn subscribe(
        &self,
        correlation: String,
        expected: ExpectedReply,
        expected_signer: Option<String>,
    ) -> ReplyTicket {
        let (tx, rx) = oneshot::channel();
        #[cfg(not(feature = "sign"))]
        let _ = expected_signer;
        self.inner.waiters.insert(
            correlation.clone(),
            ReplyWaiter {
                state: Mutex::new(ReplyWaiterState {
                    sender: Some(tx),
                    request: None,
                    pending: Vec::new(),
                }),
                expected,
                #[cfg(feature = "sign")]
                expected_signer,
            },
        );
        self.inner.subscribed.notify_one();
        ReplyTicket {
            hub: Arc::downgrade(&self.inner),
            correlation,
            rx,
        }
    }
}

/// A registered wait for one correlated reply. Consumed by [`wait`](Self::wait).
/// Dropping it (here or on timeout) deregisters the waiter.
pub(crate) struct ReplyTicket {
    hub: Weak<ReplyHubInner>,
    correlation: String,
    rx: oneshot::Receiver<AgentMessage>,
}

impl ReplyTicket {
    /// Bind the committed request address before a correlated record can complete the wait.
    pub(crate) fn arm(&self, request: Option<MessageId>) {
        let Some(hub) = self.hub.upgrade() else {
            return;
        };
        let Some(waiter) = hub.waiters.get(&self.correlation) else {
            return;
        };
        let mut state = waiter
            .state
            .lock()
            .expect("reply waiter lock is not poisoned");
        state.request = Some(request);
        let reply = state
            .pending
            .iter()
            .position(|message| request.is_none_or(|address| message.id != address))
            .map(|index| state.pending.remove(index));
        let sender = reply.as_ref().and_then(|_| state.sender.take());
        drop(state);
        drop(waiter);
        if let (Some(sender), Some(reply)) = (sender, reply) {
            hub.waiters.remove(&self.correlation);
            let _ = sender.send(reply);
        }
    }

    /// Await the correlated reply up to `timeout`, then deregister.
    pub(crate) async fn wait(mut self, timeout: Duration) -> Result<AgentMessage, LaserError> {
        match tokio::time::timeout(timeout, &mut self.rx).await {
            Ok(Ok(message)) => Ok(message),
            // The sender dropped without a reply (the hub went away), or the wait
            // timed out. Either way the caller sees a reply timeout.
            Ok(Err(_)) | Err(_) => Err(LaserError::Timeout("reply")),
        }
    }
}

impl Drop for ReplyTicket {
    fn drop(&mut self) {
        if let Some(hub) = self.hub.upgrade() {
            hub.waiters.remove(&self.correlation);
        }
    }
}

// Poll the reply topic forward and dispatch each record to its waiter. Holds only
// a `Weak` to the hub, upgrading per pass, so the task exits once the last `Laser`
// clone (and any outstanding ticket) is gone.
/// How often a parked dispatcher checks that its hub still exists.
const IDLE_HUB_CHECK: std::time::Duration = std::time::Duration::from_secs(5);

async fn dispatch_loop(
    client: Arc<IggyClient>,
    stream: Identifier,
    topic: Identifier,
    consumer: Consumer,
    mut offsets: Vec<u64>,
    hub: Weak<ReplyHubInner>,
) {
    loop {
        let Some(inner) = hub.upgrade() else {
            return;
        };
        // Nobody waits: park until a subscribe. The permit `notify_one`
        // stores covers a subscribe that lands before this wait starts.
        if inner.waiters.is_empty() {
            let parked = async {
                inner.subscribed.notified().await;
            };
            // Wake now and then to notice a hub that was dropped while idle.
            let _ = tokio::time::timeout(IDLE_HUB_CHECK, parked).await;
            drop(inner);
            continue;
        }
        // The reply topic may not exist at creation (a fan-out reply topic is
        // created by the first reply), so resolve partitions each pass until it is.
        let partitions = match client.get_topic(&stream, &topic).await {
            Ok(Some(details)) => crate::poll::bounded_partitions(details.partitions_count),
            _ => {
                drop(inner);
                tokio::time::sleep(REPLY_POLL_INTERVAL).await;
                continue;
            }
        };
        if (offsets.len() as u32) < partitions {
            offsets.resize(partitions as usize, 0);
        }
        let mut dispatched = false;
        for partition in 0..partitions {
            let from = offsets[partition as usize];
            let batch = match crate::poll::drain_partition(
                &client,
                &stream,
                &topic,
                &consumer,
                crate::poll::DrainRange::open(partition, from),
                REPLY_BATCH,
            )
            .await
            {
                Ok(batch) => batch,
                Err(_) => continue,
            };
            offsets[partition as usize] = batch.next_offset;
            for message in batch.messages {
                let Ok(decoded) = decode_agent_record(&message, laser_wire::agent::features::NONE)
                else {
                    continue;
                };
                let provenance = decoded.provenance;
                let Some(correlation) = provenance.correlation_id.clone() else {
                    continue;
                };
                let Some(waiter) = inner.waiters.get(&correlation) else {
                    continue;
                };
                #[cfg(feature = "sign")]
                let verified_principal = match &inner.verifier {
                    Some(registry) => {
                        let Some(envelope) = decoded.envelope.as_ref() else {
                            continue;
                        };
                        let Some(context) = decoded.signature_context.as_ref() else {
                            continue;
                        };
                        let Ok(verified) = registry.verify_observed_at(
                            envelope,
                            context,
                            decoded.observed_at_micros,
                        ) else {
                            continue;
                        };
                        if waiter
                            .expected_signer
                            .as_deref()
                            .is_some_and(|expected| expected != verified.principal.as_str())
                        {
                            continue;
                        }
                        Some(verified.principal)
                    }
                    None => None,
                };
                #[cfg(not(feature = "sign"))]
                let verified_principal = None;
                let reply = AgentMessage {
                    provenance,
                    id: MessageId::new(partition, message.header.offset),
                    payload: message.payload.to_vec(),
                    envelope: decoded.envelope,
                    content_type: decoded.content_type,
                    verified_principal,
                };
                if !waiter.expected.accepts(&reply) {
                    continue;
                }
                let mut state = waiter
                    .state
                    .lock()
                    .expect("reply waiter lock is not poisoned");
                let Some(request) = state.request else {
                    if state.pending.len() == PENDING_REPLY_LIMIT {
                        state.pending.remove(0);
                    }
                    state.pending.push(reply);
                    dispatched = true;
                    continue;
                };
                if request.is_some_and(|address| reply.id == address) {
                    continue;
                }
                let sender = state.sender.take();
                drop(state);
                drop(waiter);
                if let Some(sender) = sender {
                    inner.waiters.remove(&correlation);
                    let _ = sender.send(reply);
                    dispatched = true;
                }
            }
        }
        drop(inner);
        if !dispatched {
            tokio::time::sleep(REPLY_POLL_INTERVAL).await;
        }
    }
}