s2_sdk/session/

append.rs

use std::{
    collections::VecDeque,
    future::Future,
    num::NonZeroU32,
    pin::Pin,
    sync::{Arc, OnceLock},
    task::{Context, Poll},
    time::Duration,
};

use futures_util::StreamExt;
use tokio::{
    sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot},
    time::Instant,
};
use tokio_muxt::{CoalesceMode, MuxTimer};
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::task::AbortOnDropHandle;
use tracing::debug;

use crate::{
    api::{ApiError, BasinClient, Streaming, retry_builder},
    error::{AppendError, RequestError},
    frame_signal::FrameSignal,
    reconnect::{AdvisedReconnects, ReconnectAdvice},
    retry::RetryBackoffBuilder,
    types::{
        AccessTokenMode, AppendAck, AppendInput, AppendRetryPolicy, EncryptionKey, MeteredBytes,
        ONE_MIB, StreamName, StreamPosition, ValidationError,
    },
};

/// Errors returned by an append session.
#[derive(Debug, Clone, thiserror::Error)]
#[non_exhaustive]
pub enum AppendSessionError {
    /// An error with the append request underlying the session.
    #[error(transparent)]
    Append(#[from] AppendError),
    /// An append acknowledgement timed out.
    #[error("append acknowledgement timed out")]
    AckTimeout,
    /// The server disconnected during the session.
    #[error("server disconnected")]
    ServerDisconnected,
    /// The response stream closed while appends were in flight.
    #[error("response stream closed early while appends in flight")]
    StreamClosedEarly,
    /// The session was already closed.
    #[error("session already closed")]
    SessionClosed,
    /// The session is closing.
    #[error("session is closing")]
    SessionClosing,
    /// The session was dropped without being closed.
    #[error("session dropped without calling close")]
    SessionDropped,
    /// The server returned an invalid append acknowledgement.
    #[error("invalid append acknowledgement: {0}")]
    InvalidAck(String),
}

impl AppendSessionError {
    /// Whether retrying the operation is safe or sensible.
    pub fn is_retryable(&self) -> bool {
        match self {
            Self::Append(error) => error.is_retryable(),
            Self::AckTimeout | Self::ServerDisconnected => true,
            Self::StreamClosedEarly
            | Self::SessionClosed
            | Self::SessionClosing
            | Self::SessionDropped
            | Self::InvalidAck(_) => false,
        }
    }

    /// Whether retrying the operation cannot duplicate a mutation.
    pub fn has_no_side_effects(&self) -> bool {
        match self {
            Self::Append(error) => error.has_no_side_effects(),
            Self::SessionClosed | Self::SessionClosing => true,
            Self::AckTimeout
            | Self::ServerDisconnected
            | Self::StreamClosedEarly
            | Self::SessionDropped
            | Self::InvalidAck(_) => false,
        }
    }

    /// Return the underlying request error, if present.
    pub fn request_error(&self) -> Option<&RequestError> {
        match self {
            Self::Append(error) => error.request_error(),
            Self::AckTimeout
            | Self::ServerDisconnected
            | Self::StreamClosedEarly
            | Self::SessionClosed
            | Self::SessionClosing
            | Self::SessionDropped
            | Self::InvalidAck(_) => None,
        }
    }

    fn is_authentication_error(&self) -> bool {
        matches!(
            self,
            Self::Append(AppendError::Request(error)) if error.is_authentication_error()
        )
    }

    fn is_server_draining(&self) -> bool {
        matches!(
            self,
            Self::Append(AppendError::Request(error)) if error.is_server_draining()
        )
    }
}

impl From<ApiError> for AppendSessionError {
    fn from(error: ApiError) -> Self {
        match error {
            ApiError::AppendConditionFailed(condition) => {
                Self::Append(AppendError::ConditionFailed(condition.into()))
            }
            other => Self::Append(AppendError::Request(other.into())),
        }
    }
}

/// A [`Future`] that resolves to an acknowledgement once the batch of records is appended.
pub struct BatchSubmitTicket {
    rx: oneshot::Receiver<Result<AppendAck, AppendSessionError>>,
    terminal_err: Arc<OnceLock<AppendSessionError>>,
}

impl Future for BatchSubmitTicket {
    type Output = Result<AppendAck, AppendSessionError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match Pin::new(&mut self.rx).poll(cx) {
            Poll::Ready(Ok(res)) => Poll::Ready(res),
            Poll::Ready(Err(_)) => Poll::Ready(Err(self
                .terminal_err
                .get()
                .cloned()
                .unwrap_or(AppendSessionError::SessionDropped))),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[derive(Debug, Clone)]
/// Configuration for an [`AppendSession`].
pub struct AppendSessionConfig {
    max_unacked_bytes: u32,
    max_unacked_batches: Option<u32>,
}

impl Default for AppendSessionConfig {
    fn default() -> Self {
        Self {
            max_unacked_bytes: 5 * ONE_MIB,
            max_unacked_batches: None,
        }
    }
}

impl AppendSessionConfig {
    /// Create a new [`AppendSessionConfig`] with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the limit on total metered bytes of unacknowledged [`AppendInput`]s held in memory.
    ///
    /// **Note:** It must be at least `1MiB`.
    ///
    /// Defaults to `5MiB`.
    pub fn with_max_unacked_bytes(self, max_unacked_bytes: u32) -> Result<Self, ValidationError> {
        if max_unacked_bytes < ONE_MIB {
            return Err(format!("max_unacked_bytes must be at least {ONE_MIB}").into());
        }
        Ok(Self {
            max_unacked_bytes,
            ..self
        })
    }

    /// Set the limit on number of unacknowledged [`AppendInput`]s held in memory.
    ///
    /// Defaults to no limit.
    pub fn with_max_unacked_batches(self, max_unacked_batches: NonZeroU32) -> Self {
        Self {
            max_unacked_batches: Some(max_unacked_batches.get()),
            ..self
        }
    }
}

struct SessionState {
    cmd_rx: mpsc::Receiver<Command>,
    inflight_appends: VecDeque<InflightAppend>,
    inflight_bytes: usize,
    close_tx: Option<oneshot::Sender<Result<(), AppendSessionError>>>,
    total_records: usize,
    total_acked_records: usize,
    prev_ack_end: Option<StreamPosition>,
    stashed_submission: Option<StashedSubmission>,
}

impl SessionState {
    fn is_close_complete(&self) -> bool {
        self.close_tx.is_some()
            && self.inflight_appends.is_empty()
            && self.stashed_submission.is_none()
    }
}

/// A session for high-throughput appending with backpressure control. It can be created from
/// [`append_session`](crate::S2Stream::append_session).
///
/// Supports pipelining multiple [`AppendInput`]s while preserving submission order.
pub struct AppendSession {
    cmd_tx: mpsc::Sender<Command>,
    permits: AppendPermits,
    terminal_err: Arc<OnceLock<AppendSessionError>>,
    _handle: AbortOnDropHandle<()>,
}

impl AppendSession {
    pub(crate) fn new(
        client: BasinClient,
        stream: StreamName,
        encryption: Option<EncryptionKey>,
        config: AppendSessionConfig,
    ) -> Self {
        let buffer_size = config
            .max_unacked_batches
            .map(|mib| mib as usize)
            .unwrap_or(DEFAULT_CHANNEL_BUFFER_SIZE);
        let (cmd_tx, cmd_rx) = mpsc::channel(buffer_size);
        let permits = AppendPermits::new(config.max_unacked_batches, config.max_unacked_bytes);
        let retry_builder = retry_builder(&client.config.retry);
        let terminal_err = Arc::new(OnceLock::new());
        let handle = AbortOnDropHandle::new(tokio::spawn(run_session_with_retry(
            client,
            stream,
            encryption,
            cmd_rx,
            retry_builder,
            buffer_size,
            terminal_err.clone(),
        )));
        Self {
            cmd_tx,
            permits,
            terminal_err,
            _handle: handle,
        }
    }

    /// Submit a batch of records for appending.
    ///
    /// Internally, it waits on [`reserve`](Self::reserve), then submits using the permit.
    /// This provides backpressure when inflight limits are reached.
    /// For explicit control, use [`reserve`](Self::reserve) followed by
    /// [`BatchSubmitPermit::submit`].
    ///
    /// **Note**: After all submits, you must call [`close`](Self::close) to ensure all batches are
    /// appended.
    pub async fn submit(
        &self,
        input: AppendInput,
    ) -> Result<BatchSubmitTicket, AppendSessionError> {
        let permit = self.reserve(input.records.metered_bytes() as u32).await?;
        Ok(permit.submit(input))
    }

    /// Reserve capacity for a batch to be submitted. Useful in [`select!`](tokio::select) loops
    /// where you want to interleave submission with other async work. See [`submit`](Self::submit)
    /// for a simpler API.
    ///
    /// Waits when inflight limits are reached, providing explicit backpressure control.
    /// The returned permit must be used to submit the batch.
    ///
    /// **Note**: After all submits, you must call [`close`](Self::close) to ensure all batches are
    /// appended.
    ///
    /// # Cancel safety
    ///
    /// This method is cancel safe. Internally, it only awaits
    /// [`Semaphore::acquire_many_owned`](tokio::sync::Semaphore::acquire_many_owned) and
    /// [`Sender::reserve_owned`](tokio::sync::mpsc::Sender::reserve), both of which are cancel
    /// safe.
    pub async fn reserve(&self, bytes: u32) -> Result<BatchSubmitPermit, AppendSessionError> {
        let append_permit = self.permits.acquire(bytes).await;
        let cmd_tx_permit = self
            .cmd_tx
            .clone()
            .reserve_owned()
            .await
            .map_err(|_| self.terminal_err())?;
        Ok(BatchSubmitPermit {
            append_permit,
            cmd_tx_permit,
            terminal_err: self.terminal_err.clone(),
        })
    }

    /// Close the session and wait for all submitted batch of records to be appended.
    pub async fn close(self) -> Result<(), AppendSessionError> {
        let (done_tx, done_rx) = oneshot::channel();
        self.cmd_tx
            .send(Command::Close { done_tx })
            .await
            .map_err(|_| self.terminal_err())?;
        done_rx.await.map_err(|_| self.terminal_err())??;
        Ok(())
    }

    fn terminal_err(&self) -> AppendSessionError {
        self.terminal_err
            .get()
            .cloned()
            .unwrap_or(AppendSessionError::SessionClosed)
    }
}

/// A permit to submit a batch after reserving capacity.
pub struct BatchSubmitPermit {
    append_permit: AppendPermit,
    cmd_tx_permit: mpsc::OwnedPermit<Command>,
    terminal_err: Arc<OnceLock<AppendSessionError>>,
}

impl BatchSubmitPermit {
    /// Submit the batch using this permit.
    pub fn submit(self, input: AppendInput) -> BatchSubmitTicket {
        let (ack_tx, ack_rx) = oneshot::channel();
        self.cmd_tx_permit.send(Command::Submit {
            input,
            ack_tx,
            permit: Some(self.append_permit),
        });
        BatchSubmitTicket {
            rx: ack_rx,
            terminal_err: self.terminal_err,
        }
    }
}

pub(crate) struct AppendSessionInternal {
    cmd_tx: mpsc::Sender<Command>,
    terminal_err: Arc<OnceLock<AppendSessionError>>,
    _handle: AbortOnDropHandle<()>,
}

impl AppendSessionInternal {
    pub(crate) fn new(
        client: BasinClient,
        stream: StreamName,
        encryption: Option<EncryptionKey>,
    ) -> Self {
        let buffer_size = DEFAULT_CHANNEL_BUFFER_SIZE;
        let (cmd_tx, cmd_rx) = mpsc::channel(buffer_size);
        let retry_builder = retry_builder(&client.config.retry);
        let terminal_err = Arc::new(OnceLock::new());
        let handle = AbortOnDropHandle::new(tokio::spawn(run_session_with_retry(
            client,
            stream,
            encryption,
            cmd_rx,
            retry_builder,
            buffer_size,
            terminal_err.clone(),
        )));
        Self {
            cmd_tx,
            terminal_err,
            _handle: handle,
        }
    }

    pub(crate) fn submit(
        &self,
        input: AppendInput,
    ) -> impl Future<Output = Result<BatchSubmitTicket, AppendSessionError>> + Send + 'static {
        let cmd_tx = self.cmd_tx.clone();
        let terminal_err = self.terminal_err.clone();
        async move {
            let (ack_tx, ack_rx) = oneshot::channel();
            cmd_tx
                .send(Command::Submit {
                    input,
                    ack_tx,
                    permit: None,
                })
                .await
                .map_err(|_| {
                    terminal_err
                        .get()
                        .cloned()
                        .unwrap_or(AppendSessionError::SessionClosed)
                })?;
            Ok(BatchSubmitTicket {
                rx: ack_rx,
                terminal_err,
            })
        }
    }

    pub(crate) async fn close(self) -> Result<(), AppendSessionError> {
        let (done_tx, done_rx) = oneshot::channel();
        self.cmd_tx
            .send(Command::Close { done_tx })
            .await
            .map_err(|_| self.terminal_err())?;
        done_rx.await.map_err(|_| self.terminal_err())??;
        Ok(())
    }

    fn terminal_err(&self) -> AppendSessionError {
        self.terminal_err
            .get()
            .cloned()
            .unwrap_or(AppendSessionError::SessionClosed)
    }
}

#[derive(Debug)]
pub(crate) struct AppendPermit {
    _count: Option<OwnedSemaphorePermit>,
    _bytes: OwnedSemaphorePermit,
}

#[derive(Clone)]
pub(crate) struct AppendPermits {
    count: Option<Arc<Semaphore>>,
    bytes: Arc<Semaphore>,
}

impl AppendPermits {
    pub(crate) fn new(count_permits: Option<u32>, bytes_permits: u32) -> Self {
        Self {
            count: count_permits.map(|permits| Arc::new(Semaphore::new(permits as usize))),
            bytes: Arc::new(Semaphore::new(bytes_permits as usize)),
        }
    }

    pub(crate) async fn acquire(&self, bytes: u32) -> AppendPermit {
        AppendPermit {
            _count: if let Some(count) = self.count.as_ref() {
                Some(
                    count
                        .clone()
                        .acquire_many_owned(1)
                        .await
                        .expect("semaphore should not be closed"),
                )
            } else {
                None
            },
            _bytes: self
                .bytes
                .clone()
                .acquire_many_owned(bytes)
                .await
                .expect("semaphore should not be closed"),
        }
    }
}

async fn run_session_with_retry(
    client: BasinClient,
    stream: StreamName,
    encryption: Option<EncryptionKey>,
    cmd_rx: mpsc::Receiver<Command>,
    retry_builder: RetryBackoffBuilder,
    buffer_size: usize,
    terminal_err: Arc<OnceLock<AppendSessionError>>,
) {
    let access_token_mode = client.config.access_token.mode();
    let frame_signal = match client.config.retry.append_retry_policy {
        AppendRetryPolicy::NoSideEffects => Some(FrameSignal::new()),
        AppendRetryPolicy::All => None,
    };

    let mut state = SessionState {
        cmd_rx,
        inflight_appends: VecDeque::new(),
        inflight_bytes: 0,
        close_tx: None,
        total_records: 0,
        total_acked_records: 0,
        prev_ack_end: None,
        stashed_submission: None,
    };
    let mut prev_total_acked_records = 0;
    let mut retry_backoff = retry_builder.build();
    let mut advised_reconnects = AdvisedReconnects::default();

    loop {
        let result = run_session(
            &client,
            &stream,
            encryption.as_ref(),
            &mut state,
            buffer_size,
            &frame_signal,
            advised_reconnects,
        )
        .await;

        match result {
            Ok(SessionOutcome::Closed) => {
                break;
            }
            Ok(SessionOutcome::ReconnectAdvised) => {
                // The advised connection was already poisoned when the advice
                // was first decoded, so reconnecting dials a fresh one.
                advised_reconnects.record();
                debug!(
                    inflight_appends_len = state.inflight_appends.len(),
                    advised_reconnects = advised_reconnects.count(),
                    "reconnecting append session on server advice"
                );
            }
            Err(err) if err.is_server_draining() && state.is_close_complete() => break,
            Err(err) if err.is_server_draining() => {
                advised_reconnects.record();
                debug!(
                    inflight_appends_len = state.inflight_appends.len(),
                    advised_reconnects = advised_reconnects.count(),
                    "reconnecting append session while server drains"
                );
            }
            Err(err) => {
                if prev_total_acked_records < state.total_acked_records {
                    prev_total_acked_records = state.total_acked_records;
                    retry_backoff.reset();
                }

                if is_safe_to_retry(
                    &err,
                    client.config.retry.append_retry_policy,
                    !state.inflight_appends.is_empty(),
                    frame_signal.as_ref(),
                    access_token_mode,
                ) && let Some(backoff) = retry_backoff.next()
                {
                    debug!(
                        %err,
                        ?backoff,
                        num_retries_remaining = retry_backoff.remaining(),
                        "retrying append session"
                    );
                    tokio::time::sleep(backoff).await;
                } else {
                    debug!(
                        %err,
                        retries_exhausted = retry_backoff.is_exhausted(),
                        "not retrying append session"
                    );

                    let err: AppendSessionError = err;

                    let _ = terminal_err.set(err.clone());

                    for inflight_append in state.inflight_appends.drain(..) {
                        let _ = inflight_append.ack_tx.send(Err(err.clone()));
                    }

                    if let Some(stashed) = state.stashed_submission.take() {
                        let _ = stashed.ack_tx.send(Err(err.clone()));
                    }

                    if let Some(done_tx) = state.close_tx.take() {
                        let _ = done_tx.send(Err(err.clone()));
                    }

                    state.cmd_rx.close();
                    while let Some(cmd) = state.cmd_rx.recv().await {
                        cmd.reject(err.clone());
                    }
                    break;
                }
            }
        }
    }

    if let Some(done_tx) = state.close_tx.take() {
        let _ = done_tx.send(Ok(()));
    }
}

/// How a connection attempt ended without failing.
enum SessionOutcome {
    /// Everything submitted was acknowledged and the caller closed the session.
    Closed,
    /// The server advised reconnecting and this connection drained cleanly.
    ReconnectAdvised,
}

async fn run_session(
    client: &BasinClient,
    stream: &StreamName,
    encryption: Option<&EncryptionKey>,
    state: &mut SessionState,
    buffer_size: usize,
    frame_signal: &Option<FrameSignal>,
    advised_reconnects: AdvisedReconnects,
) -> Result<SessionOutcome, AppendSessionError> {
    if let Some(s) = frame_signal {
        s.reset();
    }

    let reconnect = ReconnectAdvice::default();
    let (input_tx, mut acks) = connect(
        client,
        stream,
        encryption,
        buffer_size,
        frame_signal.clone(),
        reconnect.clone(),
    )
    .await?;
    let ack_timeout = client.config.request_timeout;

    if !state.inflight_appends.is_empty() {
        resend(state, &input_tx, &mut acks, ack_timeout).await?;

        if let Some(s) = frame_signal {
            s.reset();
        }

        assert!(state.inflight_appends.is_empty());
        assert_eq!(state.inflight_bytes, 0);
    }

    if state.is_close_complete() {
        return Ok(SessionOutcome::Closed);
    }

    let timer = MuxTimer::<N_TIMER_VARIANTS>::default();
    tokio::pin!(timer);

    let mut declined_advice = false;

    loop {
        if reconnect.is_advised() && state.close_tx.is_none() && !declined_advice {
            if advised_reconnects.should_reconnect() {
                drain_for_reconnect(input_tx, acks, state, timer.as_mut(), ack_timeout).await?;
                return Ok(SessionOutcome::ReconnectAdvised);
            }
            declined_advice = true;
        }

        tokio::select! {
            (event_ord, _deadline) = &mut timer, if timer.is_armed() => {
                match TimerEvent::from(event_ord) {
                    TimerEvent::AckDeadline => {
                        return Err(AppendSessionError::AckTimeout);
                    }
                }
            }

            input_tx_permit = input_tx.reserve(), if state.stashed_submission.is_some() => {
                let input_tx_permit = input_tx_permit
                    .map_err(|_| AppendSessionError::ServerDisconnected)?;
                let submission = state.stashed_submission
                    .take()
                    .expect("stashed_submission should not be None");

                let ack_deadline = Instant::now() + ack_timeout;
                input_tx_permit.send(submission.input.clone());

                state.total_records += submission.input.records.len();
                state.inflight_bytes += submission.input_metered_bytes;

                timer.as_mut().fire_at(
                    TimerEvent::AckDeadline,
                    ack_deadline,
                    CoalesceMode::Earliest,
                );
                state.inflight_appends.push_back(InflightAppend {
                    input: submission.input,
                    input_metered_bytes: submission.input_metered_bytes,
                    ack_tx: submission.ack_tx,
                    ack_deadline,
                    _permit: submission.permit,
                });
            }

            cmd = state.cmd_rx.recv(), if state.stashed_submission.is_none() => {
                match cmd {
                    Some(Command::Submit { input, ack_tx, permit }) => {
                        if state.close_tx.is_some() {
                            let _ = ack_tx.send(
                                Err(AppendSessionError::SessionClosing)
                            );
                        } else {
                            let input_metered_bytes = input.records.metered_bytes();
                            state.stashed_submission = Some(StashedSubmission {
                                input,
                                input_metered_bytes,
                                ack_tx,
                                permit,
                            });
                        }
                    }
                    Some(Command::Close { done_tx }) => {
                        state.close_tx = Some(done_tx);
                    }
                    None => {
                        return Err(AppendSessionError::SessionDropped);
                    }
                }
            }

            ack = acks.next() => {
                match ack {
                    Some(Ok(ack)) => {
                        process_ack(
                            ack,
                            state,
                            timer.as_mut(),
                        )?;
                    }
                    Some(Err(err)) => {
                        return Err(err.into());
                    }
                    None => {
                        if !state.inflight_appends.is_empty() || state.stashed_submission.is_some() {
                            return Err(AppendSessionError::StreamClosedEarly);
                        }
                        break;
                    }
                }
            }
        }

        if state.is_close_complete() {
            break;
        }
    }

    assert!(state.inflight_appends.is_empty());
    assert_eq!(state.inflight_bytes, 0);
    assert!(state.stashed_submission.is_none());

    Ok(SessionOutcome::Closed)
}

async fn resend(
    state: &mut SessionState,
    input_tx: &mpsc::Sender<AppendInput>,
    acks: &mut Streaming<AppendAck>,
    ack_timeout: Duration,
) -> Result<(), AppendSessionError> {
    debug!(
        inflight_appends_len = state.inflight_appends.len(),
        inflight_bytes = state.inflight_bytes,
        "resending inflight appends"
    );

    let mut resend_index = 0;
    let mut resend_finished = false;

    let timer = MuxTimer::<N_TIMER_VARIANTS>::default();
    tokio::pin!(timer);

    while !state.inflight_appends.is_empty() {
        tokio::select! {
            (event_ord, _deadline) = &mut timer, if timer.is_armed() => {
                match TimerEvent::from(event_ord) {
                    TimerEvent::AckDeadline => {
                        return Err(AppendSessionError::AckTimeout);
                    }
                }
            }

            input_tx_permit = input_tx.reserve(), if !resend_finished => {
                let input_tx_permit = input_tx_permit
                    .map_err(|_| AppendSessionError::ServerDisconnected)?;

                if let Some(inflight_append) = state.inflight_appends.get_mut(resend_index) {
                    inflight_append.ack_deadline = Instant::now() + ack_timeout;
                    timer.as_mut().fire_at(
                        TimerEvent::AckDeadline,
                        inflight_append.ack_deadline,
                        CoalesceMode::Latest,
                    );
                    input_tx_permit.send(inflight_append.input.clone());
                    resend_index += 1;
                } else {
                    resend_finished = true;
                }
            }

            ack = acks.next() => {
                match ack {
                    Some(Ok(ack)) => {
                        process_ack(
                            ack,
                            state,
                            timer.as_mut(),
                        )?;
                        resend_index = resend_index.checked_sub(1).ok_or_else(|| {
                            AppendSessionError::InvalidAck(
                                "received ack without a corresponding resent append in flight".to_string(),
                            )
                        })?;
                    }
                    Some(Err(err)) => {
                        return Err(err.into());
                    }
                    None => {
                        return Err(AppendSessionError::StreamClosedEarly);
                    }
                }
            }
        }
    }

    assert_eq!(
        resend_index, 0,
        "resend_index should be 0 after resend completes"
    );
    debug!("finished resending inflight appends");
    Ok(())
}

/// Half-close so the server acknowledges everything it accepted and then ends
/// the response cleanly. Every input reaches the server ahead of the request's
/// end, so a clean end with appends still unacknowledged is a truncated
/// response, and nothing is resent.
async fn drain_for_reconnect(
    input_tx: mpsc::Sender<AppendInput>,
    mut acks: Streaming<AppendAck>,
    state: &mut SessionState,
    mut timer: Pin<&mut MuxTimer<N_TIMER_VARIANTS>>,
    ack_timeout: Duration,
) -> Result<(), AppendSessionError> {
    drop(input_tx);
    loop {
        // Bound the wait for the server's end of stream, which is otherwise
        // unbounded once nothing is in flight.
        if !timer.is_armed() {
            timer.as_mut().fire_at(
                TimerEvent::AckDeadline,
                Instant::now() + ack_timeout,
                CoalesceMode::Earliest,
            );
        }

        tokio::select! {
            (event_ord, _deadline) = &mut timer, if timer.is_armed() => {
                match TimerEvent::from(event_ord) {
                    TimerEvent::AckDeadline => {
                        return Err(AppendSessionError::AckTimeout);
                    }
                }
            }

            ack = acks.next() => {
                match ack {
                    Some(Ok(ack)) => {
                        process_ack(ack, state, timer.as_mut())?;
                    }
                    Some(Err(err)) if err.is_server_draining() => {
                        return Ok(());
                    }
                    Some(Err(err)) => {
                        return Err(err.into());
                    }
                    None => {
                        if !state.inflight_appends.is_empty() {
                            return Err(AppendSessionError::StreamClosedEarly);
                        }
                        return Ok(());
                    }
                }
            }
        }
    }
}

async fn connect(
    client: &BasinClient,
    stream: &StreamName,
    encryption: Option<&EncryptionKey>,
    buffer_size: usize,
    frame_signal: Option<FrameSignal>,
    reconnect: ReconnectAdvice,
) -> Result<(mpsc::Sender<AppendInput>, Streaming<AppendAck>), AppendSessionError> {
    let (input_tx, input_rx) = mpsc::channel::<AppendInput>(buffer_size);
    let ack_stream = Box::pin(
        client
            .append_session(
                stream,
                ReceiverStream::new(input_rx).map(|i| i.into()),
                encryption,
                frame_signal,
                reconnect,
            )
            .await?
            .map(|ack| match ack {
                Ok(ack) => Ok(ack.into()),
                Err(err) => Err(err),
            }),
    );
    Ok((input_tx, ack_stream))
}

fn process_ack(
    ack: AppendAck,
    state: &mut SessionState,
    timer: Pin<&mut MuxTimer<N_TIMER_VARIANTS>>,
) -> Result<(), AppendSessionError> {
    let corresponding_append = state.inflight_appends.pop_front().ok_or_else(|| {
        AppendSessionError::InvalidAck(
            "received ack without a corresponding append in flight".to_string(),
        )
    })?;

    if ack.end.seq_num < ack.start.seq_num {
        return Err(AppendSessionError::InvalidAck(
            "ack end seq_num should be greater than or equal to start seq_num".to_string(),
        ));
    }

    if state
        .prev_ack_end
        .is_some_and(|end| ack.end.seq_num <= end.seq_num)
    {
        return Err(AppendSessionError::InvalidAck(
            "ack end seq_num should be greater than previous ack end".to_string(),
        ));
    }

    let num_acked_records = (ack.end.seq_num - ack.start.seq_num) as usize;
    let expected_records = corresponding_append.input.records.len();
    if num_acked_records != expected_records {
        return Err(AppendSessionError::InvalidAck(format!(
            "acked record count {num_acked_records} does not match submitted batch size {expected_records}"
        )));
    }

    state.total_acked_records += num_acked_records;
    state.inflight_bytes -= corresponding_append.input_metered_bytes;
    state.prev_ack_end = Some(ack.end);

    let _ = corresponding_append.ack_tx.send(Ok(ack));

    if let Some(oldest_append) = state.inflight_appends.front() {
        timer.fire_at(
            TimerEvent::AckDeadline,
            oldest_append.ack_deadline,
            CoalesceMode::Latest,
        );
    } else {
        timer.cancel(TimerEvent::AckDeadline);
        assert_eq!(
            state.total_records, state.total_acked_records,
            "all records should be acked when inflight is empty"
        );
    }

    Ok(())
}

struct StashedSubmission {
    input: AppendInput,
    input_metered_bytes: usize,
    ack_tx: oneshot::Sender<Result<AppendAck, AppendSessionError>>,
    permit: Option<AppendPermit>,
}

struct InflightAppend {
    input: AppendInput,
    input_metered_bytes: usize,
    ack_tx: oneshot::Sender<Result<AppendAck, AppendSessionError>>,
    ack_deadline: Instant,
    _permit: Option<AppendPermit>,
}

enum Command {
    Submit {
        input: AppendInput,
        ack_tx: oneshot::Sender<Result<AppendAck, AppendSessionError>>,
        permit: Option<AppendPermit>,
    },
    Close {
        done_tx: oneshot::Sender<Result<(), AppendSessionError>>,
    },
}

impl Command {
    fn reject(self, err: AppendSessionError) {
        match self {
            Command::Submit { ack_tx, .. } => {
                let _ = ack_tx.send(Err(err));
            }
            Command::Close { done_tx } => {
                let _ = done_tx.send(Err(err));
            }
        }
    }
}

fn is_safe_to_retry(
    err: &AppendSessionError,
    policy: AppendRetryPolicy,
    has_inflight: bool,
    frame_signal: Option<&FrameSignal>,
    access_token_mode: AccessTokenMode,
) -> bool {
    let policy_compliant = match policy {
        AppendRetryPolicy::All => true,
        AppendRetryPolicy::NoSideEffects => {
            !has_inflight
                || !frame_signal.is_none_or(|s| s.is_signalled())
                || err.has_no_side_effects()
        }
    };
    policy_compliant
        && (err.is_retryable()
            || (access_token_mode.is_refreshable() && err.is_authentication_error()))
}

const DEFAULT_CHANNEL_BUFFER_SIZE: usize = 100;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TimerEvent {
    AckDeadline,
}

const N_TIMER_VARIANTS: usize = 1;

impl From<TimerEvent> for usize {
    fn from(event: TimerEvent) -> Self {
        match event {
            TimerEvent::AckDeadline => 0,
        }
    }
}

impl From<usize> for TimerEvent {
    fn from(value: usize) -> Self {
        match value {
            0 => TimerEvent::AckDeadline,
            _ => panic!("invalid ordinal"),
        }
    }
}

#[cfg(test)]
mod tests {
    use http::StatusCode;

    use super::{AppendSessionError, is_safe_to_retry};
    use crate::{
        api::{ApiError, ServerErrorBody},
        error::{AppendError, RequestError},
        frame_signal::FrameSignal,
        types::{AccessTokenMode, AppendRetryPolicy},
    };

    fn server_error(status: StatusCode, code: &str) -> AppendSessionError {
        AppendSessionError::Append(AppendError::Request(RequestError::from(ApiError::Server(
            status,
            ServerErrorBody {
                code: code.to_owned(),
                message: "test".to_owned(),
            },
        ))))
    }

    #[test]
    fn safe_to_retry_session_all_policy() {
        let retryable = server_error(StatusCode::INTERNAL_SERVER_ERROR, "internal");
        let non_retryable = server_error(StatusCode::BAD_REQUEST, "bad_request");
        let policy = AppendRetryPolicy::All;
        let static_mode = AccessTokenMode::Static;

        // All policy — always policy-compliant, just needs retryable.
        assert!(is_safe_to_retry(
            &retryable,
            policy,
            true,
            None,
            static_mode
        ));
        assert!(!is_safe_to_retry(
            &non_retryable,
            policy,
            true,
            None,
            static_mode,
        ));

        let unauthorized = server_error(StatusCode::UNAUTHORIZED, "authn");
        #[cfg(feature = "_hidden")]
        assert!(is_safe_to_retry(
            &unauthorized,
            policy,
            true,
            None,
            AccessTokenMode::Refreshable,
        ));
        assert!(!is_safe_to_retry(
            &unauthorized,
            policy,
            true,
            None,
            static_mode,
        ));

        #[cfg(feature = "_hidden")]
        let unrelated_unauthorized = server_error(StatusCode::UNAUTHORIZED, "other");
        #[cfg(feature = "_hidden")]
        assert!(!is_safe_to_retry(
            &unrelated_unauthorized,
            policy,
            true,
            None,
            AccessTokenMode::Refreshable,
        ));
    }

    #[test]
    fn safe_to_retry_session_no_side_effects_policy() {
        let retryable = server_error(StatusCode::INTERNAL_SERVER_ERROR, "internal");
        let no_side_effect = server_error(StatusCode::TOO_MANY_REQUESTS, "rate_limited");
        let policy = AppendRetryPolicy::NoSideEffects;
        let signal = FrameSignal::new();
        let mode = AccessTokenMode::Static;

        // No inflight — always safe.
        signal.signal();
        assert!(is_safe_to_retry(
            &retryable,
            policy,
            false,
            Some(&signal),
            mode,
        ));

        // Inflight + signal not set — safe (no data sent this attempt).
        signal.reset();
        assert!(is_safe_to_retry(
            &retryable,
            policy,
            true,
            Some(&signal),
            mode,
        ));

        // Inflight + signal set + error with possible side effects — not safe.
        signal.signal();
        assert!(!is_safe_to_retry(
            &retryable,
            policy,
            true,
            Some(&signal),
            mode,
        ));

        // Inflight + signal set + no-side-effect error — safe.
        assert!(is_safe_to_retry(
            &no_side_effect,
            policy,
            true,
            Some(&signal),
            mode,
        ));

        // AckTimeout — retryable but has possible side effects.
        assert!(!is_safe_to_retry(
            &AppendSessionError::AckTimeout,
            policy,
            true,
            Some(&signal),
            mode,
        ));
    }
}