s2_sdk/session/
append.rs

1use std::{
2    collections::VecDeque,
3    future::Future,
4    num::NonZeroU32,
5    pin::Pin,
6    sync::{Arc, OnceLock},
7    task::{Context, Poll},
8    time::Duration,
9};
10
11use futures_util::StreamExt;
12use tokio::{
13    sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot},
14    time::Instant,
15};
16use tokio_muxt::{CoalesceMode, MuxTimer};
17use tokio_stream::wrappers::ReceiverStream;
18use tokio_util::task::AbortOnDropHandle;
19use tracing::debug;
20
21use crate::{
22    api::{ApiError, BasinClient, Streaming, retry_builder},
23    error::{AppendError, RequestError},
24    frame_signal::FrameSignal,
25    reconnect::{AdvisedReconnects, ReconnectAdvice},
26    retry::{AppendRetryError, RetryBackoffBuilder},
27    session::StreamHeaders,
28    types::{
29        AccessTokenMode, AppendAck, AppendInput, AppendRetryPolicy, MeteredBytes, ONE_MIB,
30        StreamConfig, StreamName, StreamPosition, ValidationError,
31    },
32};
33
34/// Errors returned by an append session.
35#[derive(Debug, Clone, thiserror::Error)]
36#[non_exhaustive]
37pub enum AppendSessionError {
38    /// An error with the append request underlying the session.
39    #[error(transparent)]
40    Append(#[from] AppendError),
41    /// An append acknowledgement timed out.
42    #[error("append acknowledgement timed out")]
43    AckTimeout,
44    /// The server disconnected during the session.
45    #[error("server disconnected")]
46    ServerDisconnected,
47    /// The response stream closed while appends were in flight.
48    #[error("response stream closed early while appends in flight")]
49    StreamClosedEarly,
50    /// The session was already closed.
51    #[error("session already closed")]
52    SessionClosed,
53    /// The session is closing.
54    #[error("session is closing")]
55    SessionClosing,
56    /// The session was dropped without being closed.
57    #[error("session dropped without calling close")]
58    SessionDropped,
59    /// The server returned an invalid append acknowledgement.
60    #[error("invalid append acknowledgement: {0}")]
61    InvalidAck(String),
62    /// The final attempt failed definitively, but an earlier attempt may have taken effect,
63    /// so the entire append operation is indeterminate.
64    #[error(
65        "append may have taken effect in an earlier attempt; final attempt failed: {final_attempt_error}"
66    )]
67    IndefiniteFailure {
68        /// The definite error returned by the final attempt.
69        #[source]
70        final_attempt_error: Box<Self>,
71    },
72}
73
74impl AppendSessionError {
75    /// Whether retrying the operation is safe or sensible.
76    pub fn is_retryable(&self) -> bool {
77        match self {
78            Self::Append(error) => error.is_retryable(),
79            Self::IndefiniteFailure {
80                final_attempt_error,
81            } => final_attempt_error.is_retryable(),
82            Self::AckTimeout | Self::ServerDisconnected => true,
83            Self::StreamClosedEarly
84            | Self::SessionClosed
85            | Self::SessionClosing
86            | Self::SessionDropped
87            | Self::InvalidAck(_) => false,
88        }
89    }
90
91    /// Whether retrying the operation cannot duplicate a mutation.
92    pub fn has_no_side_effects(&self) -> bool {
93        match self {
94            Self::Append(error) => error.has_no_side_effects(),
95            Self::IndefiniteFailure { .. } => false,
96            Self::SessionClosed | Self::SessionClosing => true,
97            Self::AckTimeout
98            | Self::ServerDisconnected
99            | Self::StreamClosedEarly
100            | Self::SessionDropped
101            | Self::InvalidAck(_) => false,
102        }
103    }
104
105    /// Return the underlying request error, if present.
106    pub fn request_error(&self) -> Option<&RequestError> {
107        match self {
108            Self::Append(error) => error.request_error(),
109            Self::IndefiniteFailure {
110                final_attempt_error,
111            } => final_attempt_error.request_error(),
112            Self::AckTimeout
113            | Self::ServerDisconnected
114            | Self::StreamClosedEarly
115            | Self::SessionClosed
116            | Self::SessionClosing
117            | Self::SessionDropped
118            | Self::InvalidAck(_) => None,
119        }
120    }
121
122    fn is_authentication_error(&self) -> bool {
123        matches!(
124            self,
125            Self::Append(AppendError::Request(error)) if error.is_authentication_error()
126        )
127    }
128
129    fn is_server_draining(&self) -> bool {
130        matches!(
131            self,
132            Self::Append(AppendError::Request(error)) if error.is_server_draining()
133        )
134    }
135}
136
137impl AppendRetryError for AppendSessionError {
138    fn has_no_side_effects(&self) -> bool {
139        Self::has_no_side_effects(self)
140    }
141
142    fn into_indefinite_failure(self) -> Self {
143        Self::IndefiniteFailure {
144            final_attempt_error: Box::new(self),
145        }
146    }
147}
148
149impl From<ApiError> for AppendSessionError {
150    fn from(error: ApiError) -> Self {
151        Self::Append(error.into())
152    }
153}
154
155/// A [`Future`] that resolves to an acknowledgement once the batch of records is appended.
156pub struct BatchSubmitTicket {
157    rx: oneshot::Receiver<Result<AppendAck, AppendSessionError>>,
158    terminal_err: Arc<OnceLock<AppendSessionError>>,
159}
160
161impl Future for BatchSubmitTicket {
162    type Output = Result<AppendAck, AppendSessionError>;
163
164    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
165        match Pin::new(&mut self.rx).poll(cx) {
166            Poll::Ready(Ok(res)) => Poll::Ready(res),
167            Poll::Ready(Err(_)) => Poll::Ready(Err(self
168                .terminal_err
169                .get()
170                .cloned()
171                .unwrap_or(AppendSessionError::SessionDropped))),
172            Poll::Pending => Poll::Pending,
173        }
174    }
175}
176
177#[derive(Debug, Clone)]
178/// Configuration for an [`AppendSession`].
179pub struct AppendSessionConfig {
180    max_unacked_bytes: u32,
181    max_unacked_batches: Option<u32>,
182    stream_config: Option<StreamConfig>,
183}
184
185impl Default for AppendSessionConfig {
186    fn default() -> Self {
187        Self {
188            max_unacked_bytes: 5 * ONE_MIB,
189            max_unacked_batches: None,
190            stream_config: None,
191        }
192    }
193}
194
195impl AppendSessionConfig {
196    /// Create a new [`AppendSessionConfig`] with default settings.
197    pub fn new() -> Self {
198        Self::default()
199    }
200
201    /// Set the limit on total metered bytes of unacknowledged [`AppendInput`]s held in memory.
202    ///
203    /// **Note:** It must be at least `1MiB`.
204    ///
205    /// Defaults to `5MiB`.
206    pub fn with_max_unacked_bytes(self, max_unacked_bytes: u32) -> Result<Self, ValidationError> {
207        if max_unacked_bytes < ONE_MIB {
208            return Err(format!("max_unacked_bytes must be at least {ONE_MIB}").into());
209        }
210        Ok(Self {
211            max_unacked_bytes,
212            ..self
213        })
214    }
215
216    /// Set the limit on number of unacknowledged [`AppendInput`]s held in memory.
217    ///
218    /// Defaults to no limit.
219    pub fn with_max_unacked_batches(self, max_unacked_batches: NonZeroU32) -> Self {
220        Self {
221            max_unacked_batches: Some(max_unacked_batches.get()),
222            ..self
223        }
224    }
225
226    /// Set the stream configuration to apply if the stream is created on append.
227    ///
228    /// Unset fields inherit the basin's default stream configuration. Ignored if the stream
229    /// already exists.
230    ///
231    /// Defaults to `None`.
232    pub fn with_stream_config(self, stream_config: StreamConfig) -> Self {
233        Self {
234            stream_config: Some(stream_config),
235            ..self
236        }
237    }
238
239    pub(crate) fn stream_config(&self) -> Option<&StreamConfig> {
240        self.stream_config.as_ref()
241    }
242}
243
244struct SessionState {
245    cmd_rx: mpsc::Receiver<Command>,
246    inflight_appends: VecDeque<InflightAppend>,
247    inflight_bytes: usize,
248    close_tx: Option<oneshot::Sender<Result<(), AppendSessionError>>>,
249    total_records: usize,
250    total_acked_records: usize,
251    prev_ack_end: Option<StreamPosition>,
252    stashed_submission: Option<StashedSubmission>,
253}
254
255impl SessionState {
256    fn is_close_complete(&self) -> bool {
257        self.close_tx.is_some()
258            && self.inflight_appends.is_empty()
259            && self.stashed_submission.is_none()
260    }
261}
262
263/// A session for high-throughput appending with backpressure control. It can be created from
264/// [`append_session`](crate::S2Stream::append_session).
265///
266/// Supports pipelining multiple [`AppendInput`]s while preserving submission order.
267pub struct AppendSession {
268    cmd_tx: mpsc::Sender<Command>,
269    permits: AppendPermits,
270    terminal_err: Arc<OnceLock<AppendSessionError>>,
271    _handle: AbortOnDropHandle<()>,
272}
273
274impl AppendSession {
275    pub(crate) fn new(
276        client: BasinClient,
277        stream: StreamName,
278        headers: StreamHeaders,
279        config: AppendSessionConfig,
280    ) -> Self {
281        let buffer_size = config
282            .max_unacked_batches
283            .map(|mib| mib as usize)
284            .unwrap_or(DEFAULT_CHANNEL_BUFFER_SIZE);
285        let (cmd_tx, cmd_rx) = mpsc::channel(buffer_size);
286        let permits = AppendPermits::new(config.max_unacked_batches, config.max_unacked_bytes);
287        let retry_builder = retry_builder(&client.config.retry);
288        let terminal_err = Arc::new(OnceLock::new());
289        let handle = AbortOnDropHandle::new(tokio::spawn(run_session_with_retry(
290            client,
291            stream,
292            headers,
293            cmd_rx,
294            retry_builder,
295            buffer_size,
296            terminal_err.clone(),
297        )));
298        Self {
299            cmd_tx,
300            permits,
301            terminal_err,
302            _handle: handle,
303        }
304    }
305
306    /// Submit a batch of records for appending.
307    ///
308    /// Internally, it waits on [`reserve`](Self::reserve), then submits using the permit.
309    /// This provides backpressure when inflight limits are reached.
310    /// For explicit control, use [`reserve`](Self::reserve) followed by
311    /// [`BatchSubmitPermit::submit`].
312    ///
313    /// **Note**: After all submits, you must call [`close`](Self::close) to ensure all batches are
314    /// appended.
315    pub async fn submit(
316        &self,
317        input: AppendInput,
318    ) -> Result<BatchSubmitTicket, AppendSessionError> {
319        let permit = self.reserve(input.records.metered_bytes() as u32).await?;
320        Ok(permit.submit(input))
321    }
322
323    /// Reserve capacity for a batch to be submitted. Useful in [`select!`](tokio::select) loops
324    /// where you want to interleave submission with other async work. See [`submit`](Self::submit)
325    /// for a simpler API.
326    ///
327    /// Waits when inflight limits are reached, providing explicit backpressure control.
328    /// The returned permit must be used to submit the batch.
329    ///
330    /// **Note**: After all submits, you must call [`close`](Self::close) to ensure all batches are
331    /// appended.
332    ///
333    /// # Cancel safety
334    ///
335    /// This method is cancel safe. Internally, it only awaits
336    /// [`Semaphore::acquire_many_owned`](tokio::sync::Semaphore::acquire_many_owned) and
337    /// [`Sender::reserve_owned`](tokio::sync::mpsc::Sender::reserve), both of which are cancel
338    /// safe.
339    pub async fn reserve(&self, bytes: u32) -> Result<BatchSubmitPermit, AppendSessionError> {
340        let append_permit = self.permits.acquire(bytes).await;
341        let cmd_tx_permit = self
342            .cmd_tx
343            .clone()
344            .reserve_owned()
345            .await
346            .map_err(|_| self.terminal_err())?;
347        Ok(BatchSubmitPermit {
348            append_permit,
349            cmd_tx_permit,
350            terminal_err: self.terminal_err.clone(),
351        })
352    }
353
354    /// Close the session and wait for all submitted batch of records to be appended.
355    pub async fn close(self) -> Result<(), AppendSessionError> {
356        let (done_tx, done_rx) = oneshot::channel();
357        self.cmd_tx
358            .send(Command::Close { done_tx })
359            .await
360            .map_err(|_| self.terminal_err())?;
361        done_rx.await.map_err(|_| self.terminal_err())??;
362        Ok(())
363    }
364
365    fn terminal_err(&self) -> AppendSessionError {
366        self.terminal_err
367            .get()
368            .cloned()
369            .unwrap_or(AppendSessionError::SessionClosed)
370    }
371}
372
373/// A permit to submit a batch after reserving capacity.
374pub struct BatchSubmitPermit {
375    append_permit: AppendPermit,
376    cmd_tx_permit: mpsc::OwnedPermit<Command>,
377    terminal_err: Arc<OnceLock<AppendSessionError>>,
378}
379
380impl BatchSubmitPermit {
381    /// Submit the batch using this permit.
382    pub fn submit(self, input: AppendInput) -> BatchSubmitTicket {
383        let (ack_tx, ack_rx) = oneshot::channel();
384        self.cmd_tx_permit.send(Command::Submit {
385            input,
386            ack_tx,
387            permit: Some(self.append_permit),
388        });
389        BatchSubmitTicket {
390            rx: ack_rx,
391            terminal_err: self.terminal_err,
392        }
393    }
394}
395
396pub(crate) struct AppendSessionInternal {
397    cmd_tx: mpsc::Sender<Command>,
398    terminal_err: Arc<OnceLock<AppendSessionError>>,
399    _handle: AbortOnDropHandle<()>,
400}
401
402impl AppendSessionInternal {
403    pub(crate) fn new(client: BasinClient, stream: StreamName, headers: StreamHeaders) -> Self {
404        let buffer_size = DEFAULT_CHANNEL_BUFFER_SIZE;
405        let (cmd_tx, cmd_rx) = mpsc::channel(buffer_size);
406        let retry_builder = retry_builder(&client.config.retry);
407        let terminal_err = Arc::new(OnceLock::new());
408        let handle = AbortOnDropHandle::new(tokio::spawn(run_session_with_retry(
409            client,
410            stream,
411            headers,
412            cmd_rx,
413            retry_builder,
414            buffer_size,
415            terminal_err.clone(),
416        )));
417        Self {
418            cmd_tx,
419            terminal_err,
420            _handle: handle,
421        }
422    }
423
424    pub(crate) fn submit(
425        &self,
426        input: AppendInput,
427    ) -> impl Future<Output = Result<BatchSubmitTicket, AppendSessionError>> + Send + 'static {
428        let cmd_tx = self.cmd_tx.clone();
429        let terminal_err = self.terminal_err.clone();
430        async move {
431            let (ack_tx, ack_rx) = oneshot::channel();
432            cmd_tx
433                .send(Command::Submit {
434                    input,
435                    ack_tx,
436                    permit: None,
437                })
438                .await
439                .map_err(|_| {
440                    terminal_err
441                        .get()
442                        .cloned()
443                        .unwrap_or(AppendSessionError::SessionClosed)
444                })?;
445            Ok(BatchSubmitTicket {
446                rx: ack_rx,
447                terminal_err,
448            })
449        }
450    }
451
452    pub(crate) async fn close(self) -> Result<(), AppendSessionError> {
453        let (done_tx, done_rx) = oneshot::channel();
454        self.cmd_tx
455            .send(Command::Close { done_tx })
456            .await
457            .map_err(|_| self.terminal_err())?;
458        done_rx.await.map_err(|_| self.terminal_err())??;
459        Ok(())
460    }
461
462    fn terminal_err(&self) -> AppendSessionError {
463        self.terminal_err
464            .get()
465            .cloned()
466            .unwrap_or(AppendSessionError::SessionClosed)
467    }
468}
469
470#[derive(Debug)]
471pub(crate) struct AppendPermit {
472    _count: Option<OwnedSemaphorePermit>,
473    _bytes: OwnedSemaphorePermit,
474}
475
476#[derive(Clone)]
477pub(crate) struct AppendPermits {
478    count: Option<Arc<Semaphore>>,
479    bytes: Arc<Semaphore>,
480}
481
482impl AppendPermits {
483    pub(crate) fn new(count_permits: Option<u32>, bytes_permits: u32) -> Self {
484        Self {
485            count: count_permits.map(|permits| Arc::new(Semaphore::new(permits as usize))),
486            bytes: Arc::new(Semaphore::new(bytes_permits as usize)),
487        }
488    }
489
490    pub(crate) async fn acquire(&self, bytes: u32) -> AppendPermit {
491        AppendPermit {
492            _count: if let Some(count) = self.count.as_ref() {
493                Some(
494                    count
495                        .clone()
496                        .acquire_many_owned(1)
497                        .await
498                        .expect("semaphore should not be closed"),
499                )
500            } else {
501                None
502            },
503            _bytes: self
504                .bytes
505                .clone()
506                .acquire_many_owned(bytes)
507                .await
508                .expect("semaphore should not be closed"),
509        }
510    }
511}
512
513async fn run_session_with_retry(
514    client: BasinClient,
515    stream: StreamName,
516    headers: StreamHeaders,
517    cmd_rx: mpsc::Receiver<Command>,
518    retry_builder: RetryBackoffBuilder,
519    buffer_size: usize,
520    terminal_err: Arc<OnceLock<AppendSessionError>>,
521) {
522    let access_token_mode = client.config.access_token.mode();
523    let frame_signal = match client.config.retry.append_retry_policy {
524        AppendRetryPolicy::NoSideEffects => Some(FrameSignal::new()),
525        AppendRetryPolicy::All => None,
526    };
527
528    let mut state = SessionState {
529        cmd_rx,
530        inflight_appends: VecDeque::new(),
531        inflight_bytes: 0,
532        close_tx: None,
533        total_records: 0,
534        total_acked_records: 0,
535        prev_ack_end: None,
536        stashed_submission: None,
537    };
538    let mut prev_total_acked_records = 0;
539    let mut retry_backoff = retry_builder.build();
540    let mut advised_reconnects = AdvisedReconnects::default();
541
542    loop {
543        let result = run_session(
544            &client,
545            &stream,
546            &headers,
547            &mut state,
548            buffer_size,
549            &frame_signal,
550            advised_reconnects,
551        )
552        .await;
553
554        match result {
555            Ok(SessionOutcome::Closed) => {
556                break;
557            }
558            Ok(SessionOutcome::ReconnectAdvised) => {
559                // The advised connection was already poisoned when the advice
560                // was first decoded, so reconnecting dials a fresh one.
561                advised_reconnects.record();
562                debug!(
563                    inflight_appends_len = state.inflight_appends.len(),
564                    advised_reconnects = advised_reconnects.count(),
565                    "reconnecting append session on server advice"
566                );
567            }
568            Err(err) if err.is_server_draining() && state.is_close_complete() => break,
569            Err(err) if err.is_server_draining() => {
570                advised_reconnects.record();
571                debug!(
572                    inflight_appends_len = state.inflight_appends.len(),
573                    advised_reconnects = advised_reconnects.count(),
574                    "reconnecting append session while server drains"
575                );
576            }
577            Err(err) => {
578                if prev_total_acked_records < state.total_acked_records {
579                    prev_total_acked_records = state.total_acked_records;
580                    retry_backoff.reset();
581                }
582
583                if is_safe_to_retry(
584                    &err,
585                    client.config.retry.append_retry_policy,
586                    !state.inflight_appends.is_empty(),
587                    frame_signal.as_ref(),
588                    access_token_mode,
589                ) && let Some(backoff) = retry_backoff.next()
590                {
591                    if err.attempt_may_have_side_effects(frame_signal.as_ref()) {
592                        for append in &mut state.inflight_appends {
593                            append.prior_uncertainty = true;
594                        }
595                    }
596                    debug!(
597                        %err,
598                        ?backoff,
599                        num_retries_remaining = retry_backoff.remaining(),
600                        "retrying append session"
601                    );
602                    tokio::time::sleep(backoff).await;
603                } else {
604                    debug!(
605                        %err,
606                        retries_exhausted = retry_backoff.is_exhausted(),
607                        "not retrying append session"
608                    );
609
610                    let session_err = err.clone().with_prior_uncertainty(
611                        state.inflight_appends.iter().any(|a| a.prior_uncertainty),
612                    );
613                    let _ = terminal_err.set(session_err.clone());
614
615                    for inflight_append in state.inflight_appends.drain(..) {
616                        let error = err
617                            .clone()
618                            .with_prior_uncertainty(inflight_append.prior_uncertainty);
619                        let _ = inflight_append.ack_tx.send(Err(error));
620                    }
621
622                    if let Some(stashed) = state.stashed_submission.take() {
623                        let _ = stashed.ack_tx.send(Err(err.clone()));
624                    }
625
626                    if let Some(done_tx) = state.close_tx.take() {
627                        let _ = done_tx.send(Err(session_err.clone()));
628                    }
629
630                    state.cmd_rx.close();
631                    while let Some(cmd) = state.cmd_rx.recv().await {
632                        let error = match &cmd {
633                            Command::Submit { .. } => &err,
634                            Command::Close { .. } => &session_err,
635                        };
636                        cmd.reject(error.clone());
637                    }
638                    break;
639                }
640            }
641        }
642    }
643
644    if let Some(done_tx) = state.close_tx.take() {
645        let _ = done_tx.send(Ok(()));
646    }
647}
648
649/// How a connection attempt ended without failing.
650enum SessionOutcome {
651    /// Everything submitted was acknowledged and the caller closed the session.
652    Closed,
653    /// The server advised reconnecting and this connection drained cleanly.
654    ReconnectAdvised,
655}
656
657async fn run_session(
658    client: &BasinClient,
659    stream: &StreamName,
660    headers: &StreamHeaders,
661    state: &mut SessionState,
662    buffer_size: usize,
663    frame_signal: &Option<FrameSignal>,
664    advised_reconnects: AdvisedReconnects,
665) -> Result<SessionOutcome, AppendSessionError> {
666    if let Some(s) = frame_signal {
667        s.reset();
668    }
669
670    let reconnect = ReconnectAdvice::default();
671    let (input_tx, mut acks) = connect(
672        client,
673        stream,
674        headers,
675        buffer_size,
676        frame_signal.clone(),
677        reconnect.clone(),
678    )
679    .await?;
680    let ack_timeout = client.config.request_timeout;
681
682    if !state.inflight_appends.is_empty() {
683        resend(state, &input_tx, &mut acks, ack_timeout).await?;
684
685        if let Some(s) = frame_signal {
686            s.reset();
687        }
688
689        assert!(state.inflight_appends.is_empty());
690        assert_eq!(state.inflight_bytes, 0);
691    }
692
693    if state.is_close_complete() {
694        return Ok(SessionOutcome::Closed);
695    }
696
697    let timer = MuxTimer::<N_TIMER_VARIANTS>::default();
698    tokio::pin!(timer);
699
700    let mut declined_advice = false;
701
702    loop {
703        if reconnect.is_advised() && state.close_tx.is_none() && !declined_advice {
704            if advised_reconnects.should_reconnect() {
705                drain_for_reconnect(input_tx, acks, state, timer.as_mut(), ack_timeout).await?;
706                return Ok(SessionOutcome::ReconnectAdvised);
707            }
708            declined_advice = true;
709        }
710
711        tokio::select! {
712            (event_ord, _deadline) = &mut timer, if timer.is_armed() => {
713                match TimerEvent::from(event_ord) {
714                    TimerEvent::AckDeadline => {
715                        return Err(AppendSessionError::AckTimeout);
716                    }
717                }
718            }
719
720            input_tx_permit = input_tx.reserve(), if state.stashed_submission.is_some() => {
721                let input_tx_permit = input_tx_permit
722                    .map_err(|_| AppendSessionError::ServerDisconnected)?;
723                let submission = state.stashed_submission
724                    .take()
725                    .expect("stashed_submission should not be None");
726
727                let ack_deadline = Instant::now() + ack_timeout;
728                input_tx_permit.send(submission.input.clone());
729
730                state.total_records += submission.input.records.len();
731                state.inflight_bytes += submission.input_metered_bytes;
732
733                timer.as_mut().fire_at(
734                    TimerEvent::AckDeadline,
735                    ack_deadline,
736                    CoalesceMode::Earliest,
737                );
738                state.inflight_appends.push_back(InflightAppend {
739                    input: submission.input,
740                    input_metered_bytes: submission.input_metered_bytes,
741                    ack_tx: submission.ack_tx,
742                    ack_deadline,
743                    _permit: submission.permit,
744                    prior_uncertainty: false,
745                });
746            }
747
748            cmd = state.cmd_rx.recv(), if state.stashed_submission.is_none() => {
749                match cmd {
750                    Some(Command::Submit { input, ack_tx, permit }) => {
751                        if state.close_tx.is_some() {
752                            let _ = ack_tx.send(
753                                Err(AppendSessionError::SessionClosing)
754                            );
755                        } else {
756                            let input_metered_bytes = input.records.metered_bytes();
757                            state.stashed_submission = Some(StashedSubmission {
758                                input,
759                                input_metered_bytes,
760                                ack_tx,
761                                permit,
762                            });
763                        }
764                    }
765                    Some(Command::Close { done_tx }) => {
766                        state.close_tx = Some(done_tx);
767                    }
768                    None => {
769                        return Err(AppendSessionError::SessionDropped);
770                    }
771                }
772            }
773
774            ack = acks.next() => {
775                match ack {
776                    Some(Ok(ack)) => {
777                        process_ack(
778                            ack,
779                            state,
780                            timer.as_mut(),
781                        )?;
782                    }
783                    Some(Err(err)) => {
784                        return Err(err.into());
785                    }
786                    None => {
787                        if !state.inflight_appends.is_empty() || state.stashed_submission.is_some() {
788                            return Err(AppendSessionError::StreamClosedEarly);
789                        }
790                        break;
791                    }
792                }
793            }
794        }
795
796        if state.is_close_complete() {
797            break;
798        }
799    }
800
801    assert!(state.inflight_appends.is_empty());
802    assert_eq!(state.inflight_bytes, 0);
803    assert!(state.stashed_submission.is_none());
804
805    Ok(SessionOutcome::Closed)
806}
807
808async fn resend(
809    state: &mut SessionState,
810    input_tx: &mpsc::Sender<AppendInput>,
811    acks: &mut Streaming<AppendAck>,
812    ack_timeout: Duration,
813) -> Result<(), AppendSessionError> {
814    debug!(
815        inflight_appends_len = state.inflight_appends.len(),
816        inflight_bytes = state.inflight_bytes,
817        "resending inflight appends"
818    );
819
820    let mut resend_index = 0;
821    let mut resend_finished = false;
822
823    let timer = MuxTimer::<N_TIMER_VARIANTS>::default();
824    tokio::pin!(timer);
825
826    while !state.inflight_appends.is_empty() {
827        tokio::select! {
828            (event_ord, _deadline) = &mut timer, if timer.is_armed() => {
829                match TimerEvent::from(event_ord) {
830                    TimerEvent::AckDeadline => {
831                        return Err(AppendSessionError::AckTimeout);
832                    }
833                }
834            }
835
836            input_tx_permit = input_tx.reserve(), if !resend_finished => {
837                let input_tx_permit = input_tx_permit
838                    .map_err(|_| AppendSessionError::ServerDisconnected)?;
839
840                if let Some(inflight_append) = state.inflight_appends.get_mut(resend_index) {
841                    inflight_append.ack_deadline = Instant::now() + ack_timeout;
842                    timer.as_mut().fire_at(
843                        TimerEvent::AckDeadline,
844                        inflight_append.ack_deadline,
845                        CoalesceMode::Latest,
846                    );
847                    input_tx_permit.send(inflight_append.input.clone());
848                    resend_index += 1;
849                } else {
850                    resend_finished = true;
851                }
852            }
853
854            ack = acks.next() => {
855                match ack {
856                    Some(Ok(ack)) => {
857                        process_ack(
858                            ack,
859                            state,
860                            timer.as_mut(),
861                        )?;
862                        resend_index = resend_index.checked_sub(1).ok_or_else(|| {
863                            AppendSessionError::InvalidAck(
864                                "received ack without a corresponding resent append in flight".to_string(),
865                            )
866                        })?;
867                    }
868                    Some(Err(err)) => {
869                        return Err(err.into());
870                    }
871                    None => {
872                        return Err(AppendSessionError::StreamClosedEarly);
873                    }
874                }
875            }
876        }
877    }
878
879    assert_eq!(
880        resend_index, 0,
881        "resend_index should be 0 after resend completes"
882    );
883    debug!("finished resending inflight appends");
884    Ok(())
885}
886
887/// Half-close so the server acknowledges everything it accepted and then ends
888/// the response cleanly. Every input reaches the server ahead of the request's
889/// end, so a clean end with appends still unacknowledged is a truncated
890/// response, and nothing is resent.
891async fn drain_for_reconnect(
892    input_tx: mpsc::Sender<AppendInput>,
893    mut acks: Streaming<AppendAck>,
894    state: &mut SessionState,
895    mut timer: Pin<&mut MuxTimer<N_TIMER_VARIANTS>>,
896    ack_timeout: Duration,
897) -> Result<(), AppendSessionError> {
898    drop(input_tx);
899    loop {
900        // Bound the wait for the server's end of stream, which is otherwise
901        // unbounded once nothing is in flight.
902        if !timer.is_armed() {
903            timer.as_mut().fire_at(
904                TimerEvent::AckDeadline,
905                Instant::now() + ack_timeout,
906                CoalesceMode::Earliest,
907            );
908        }
909
910        tokio::select! {
911            (event_ord, _deadline) = &mut timer, if timer.is_armed() => {
912                match TimerEvent::from(event_ord) {
913                    TimerEvent::AckDeadline => {
914                        return Err(AppendSessionError::AckTimeout);
915                    }
916                }
917            }
918
919            ack = acks.next() => {
920                match ack {
921                    Some(Ok(ack)) => {
922                        process_ack(ack, state, timer.as_mut())?;
923                    }
924                    Some(Err(err)) if err.is_server_draining() => {
925                        return Ok(());
926                    }
927                    Some(Err(err)) => {
928                        return Err(err.into());
929                    }
930                    None => {
931                        if !state.inflight_appends.is_empty() {
932                            return Err(AppendSessionError::StreamClosedEarly);
933                        }
934                        return Ok(());
935                    }
936                }
937            }
938        }
939    }
940}
941
942async fn connect(
943    client: &BasinClient,
944    stream: &StreamName,
945    headers: &StreamHeaders,
946    buffer_size: usize,
947    frame_signal: Option<FrameSignal>,
948    reconnect: ReconnectAdvice,
949) -> Result<(mpsc::Sender<AppendInput>, Streaming<AppendAck>), AppendSessionError> {
950    let (input_tx, input_rx) = mpsc::channel::<AppendInput>(buffer_size);
951    let ack_stream = Box::pin(
952        client
953            .append_session(
954                stream,
955                ReceiverStream::new(input_rx).map(|i| i.into()),
956                headers.encryption.as_ref(),
957                headers.stream_config.as_ref(),
958                frame_signal,
959                reconnect,
960            )
961            .await?
962            .map(|ack| match ack {
963                Ok(ack) => Ok(ack.into()),
964                Err(err) => Err(err),
965            }),
966    );
967    Ok((input_tx, ack_stream))
968}
969
970fn process_ack(
971    ack: AppendAck,
972    state: &mut SessionState,
973    timer: Pin<&mut MuxTimer<N_TIMER_VARIANTS>>,
974) -> Result<(), AppendSessionError> {
975    let corresponding_append = state.inflight_appends.pop_front().ok_or_else(|| {
976        AppendSessionError::InvalidAck(
977            "received ack without a corresponding append in flight".to_string(),
978        )
979    })?;
980
981    if ack.end.seq_num < ack.start.seq_num {
982        return Err(AppendSessionError::InvalidAck(
983            "ack end seq_num should be greater than or equal to start seq_num".to_string(),
984        ));
985    }
986
987    if state
988        .prev_ack_end
989        .is_some_and(|end| ack.end.seq_num <= end.seq_num)
990    {
991        return Err(AppendSessionError::InvalidAck(
992            "ack end seq_num should be greater than previous ack end".to_string(),
993        ));
994    }
995
996    let num_acked_records = (ack.end.seq_num - ack.start.seq_num) as usize;
997    let expected_records = corresponding_append.input.records.len();
998    if num_acked_records != expected_records {
999        return Err(AppendSessionError::InvalidAck(format!(
1000            "acked record count {num_acked_records} does not match submitted batch size {expected_records}"
1001        )));
1002    }
1003
1004    state.total_acked_records += num_acked_records;
1005    state.inflight_bytes -= corresponding_append.input_metered_bytes;
1006    state.prev_ack_end = Some(ack.end);
1007
1008    let _ = corresponding_append.ack_tx.send(Ok(ack));
1009
1010    if let Some(oldest_append) = state.inflight_appends.front() {
1011        timer.fire_at(
1012            TimerEvent::AckDeadline,
1013            oldest_append.ack_deadline,
1014            CoalesceMode::Latest,
1015        );
1016    } else {
1017        timer.cancel(TimerEvent::AckDeadline);
1018        assert_eq!(
1019            state.total_records, state.total_acked_records,
1020            "all records should be acked when inflight is empty"
1021        );
1022    }
1023
1024    Ok(())
1025}
1026
1027struct StashedSubmission {
1028    input: AppendInput,
1029    input_metered_bytes: usize,
1030    ack_tx: oneshot::Sender<Result<AppendAck, AppendSessionError>>,
1031    permit: Option<AppendPermit>,
1032}
1033
1034struct InflightAppend {
1035    input: AppendInput,
1036    input_metered_bytes: usize,
1037    ack_tx: oneshot::Sender<Result<AppendAck, AppendSessionError>>,
1038    ack_deadline: Instant,
1039    _permit: Option<AppendPermit>,
1040    prior_uncertainty: bool,
1041}
1042
1043enum Command {
1044    Submit {
1045        input: AppendInput,
1046        ack_tx: oneshot::Sender<Result<AppendAck, AppendSessionError>>,
1047        permit: Option<AppendPermit>,
1048    },
1049    Close {
1050        done_tx: oneshot::Sender<Result<(), AppendSessionError>>,
1051    },
1052}
1053
1054impl Command {
1055    fn reject(self, err: AppendSessionError) {
1056        match self {
1057            Command::Submit { ack_tx, .. } => {
1058                let _ = ack_tx.send(Err(err));
1059            }
1060            Command::Close { done_tx } => {
1061                let _ = done_tx.send(Err(err));
1062            }
1063        }
1064    }
1065}
1066
1067fn is_safe_to_retry(
1068    err: &AppendSessionError,
1069    policy: AppendRetryPolicy,
1070    has_inflight: bool,
1071    frame_signal: Option<&FrameSignal>,
1072    access_token_mode: AccessTokenMode,
1073) -> bool {
1074    let policy_compliant = match policy {
1075        AppendRetryPolicy::All => true,
1076        AppendRetryPolicy::NoSideEffects => {
1077            !has_inflight || !err.attempt_may_have_side_effects(frame_signal)
1078        }
1079    };
1080    policy_compliant
1081        && (err.is_retryable()
1082            || (access_token_mode.is_refreshable() && err.is_authentication_error()))
1083}
1084
1085const DEFAULT_CHANNEL_BUFFER_SIZE: usize = 100;
1086
1087#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1088enum TimerEvent {
1089    AckDeadline,
1090}
1091
1092const N_TIMER_VARIANTS: usize = 1;
1093
1094impl From<TimerEvent> for usize {
1095    fn from(event: TimerEvent) -> Self {
1096        match event {
1097            TimerEvent::AckDeadline => 0,
1098        }
1099    }
1100}
1101
1102impl From<usize> for TimerEvent {
1103    fn from(value: usize) -> Self {
1104        match value {
1105            0 => TimerEvent::AckDeadline,
1106            _ => panic!("invalid ordinal"),
1107        }
1108    }
1109}
1110
1111#[cfg(test)]
1112mod tests {
1113    use std::error::Error;
1114
1115    use http::StatusCode;
1116
1117    use super::{AppendSessionError, is_safe_to_retry};
1118    use crate::{
1119        api::{ApiError, ServerErrorBody},
1120        error::{AppendError, ProducerError, RequestError},
1121        frame_signal::FrameSignal,
1122        retry::AppendRetryError,
1123        types::{AccessTokenMode, AppendRetryPolicy},
1124    };
1125
1126    fn server_error(status: StatusCode, code: &str) -> AppendSessionError {
1127        AppendSessionError::Append(AppendError::Request(RequestError::from(ApiError::Server(
1128            status,
1129            ServerErrorBody {
1130                code: code.to_owned(),
1131                message: "test".to_owned(),
1132            },
1133        ))))
1134    }
1135
1136    #[rstest::rstest]
1137    #[case(StatusCode::FORBIDDEN, "permission_denied", false, false)]
1138    #[case(StatusCode::FORBIDDEN, "permission_denied", true, true)]
1139    #[case(StatusCode::SERVICE_UNAVAILABLE, "unavailable", false, false)]
1140    #[case(StatusCode::SERVICE_UNAVAILABLE, "unavailable", true, false)]
1141    #[case(StatusCode::TOO_MANY_REQUESTS, "rate_limited", true, true)]
1142    #[test]
1143    fn session_failure_preserves_uncertainty_and_latest_error(
1144        #[case] status: StatusCode,
1145        #[case] code: &str,
1146        #[case] prior_uncertainty: bool,
1147        #[case] wrapped: bool,
1148    ) {
1149        let latest = server_error(status, code);
1150        let error = latest.clone().with_prior_uncertainty(prior_uncertainty);
1151        assert_eq!(
1152            matches!(error, AppendSessionError::IndefiniteFailure { .. }),
1153            wrapped
1154        );
1155        assert_eq!(error.is_retryable(), latest.is_retryable());
1156        assert_eq!(
1157            error.has_no_side_effects(),
1158            !wrapped && latest.has_no_side_effects()
1159        );
1160        assert_eq!(
1161            error.request_error().unwrap().server_error().unwrap().code,
1162            code
1163        );
1164        if wrapped {
1165            let source = error
1166                .source()
1167                .unwrap()
1168                .downcast_ref::<Box<AppendSessionError>>()
1169                .unwrap();
1170            assert_eq!(source.to_string(), latest.to_string());
1171            assert!(source.has_no_side_effects());
1172        }
1173
1174        let producer_error = ProducerError::from(error.clone());
1175        assert_eq!(producer_error.is_retryable(), error.is_retryable());
1176        assert_eq!(
1177            producer_error.has_no_side_effects(),
1178            error.has_no_side_effects()
1179        );
1180        assert_eq!(
1181            producer_error
1182                .request_error()
1183                .unwrap()
1184                .server_error()
1185                .unwrap()
1186                .code,
1187            code
1188        );
1189    }
1190
1191    #[test]
1192    fn safe_to_retry_session_all_policy() {
1193        let retryable = server_error(StatusCode::INTERNAL_SERVER_ERROR, "internal");
1194        let non_retryable = server_error(StatusCode::BAD_REQUEST, "bad_request");
1195        let policy = AppendRetryPolicy::All;
1196        let static_mode = AccessTokenMode::Static;
1197
1198        // All policy — always policy-compliant, just needs retryable.
1199        assert!(is_safe_to_retry(
1200            &retryable,
1201            policy,
1202            true,
1203            None,
1204            static_mode
1205        ));
1206        assert!(!is_safe_to_retry(
1207            &non_retryable,
1208            policy,
1209            true,
1210            None,
1211            static_mode,
1212        ));
1213
1214        let unauthorized = server_error(StatusCode::UNAUTHORIZED, "authn");
1215        #[cfg(feature = "_hidden")]
1216        assert!(is_safe_to_retry(
1217            &unauthorized,
1218            policy,
1219            true,
1220            None,
1221            AccessTokenMode::Refreshable,
1222        ));
1223        assert!(!is_safe_to_retry(
1224            &unauthorized,
1225            policy,
1226            true,
1227            None,
1228            static_mode,
1229        ));
1230
1231        #[cfg(feature = "_hidden")]
1232        let unrelated_unauthorized = server_error(StatusCode::UNAUTHORIZED, "other");
1233        #[cfg(feature = "_hidden")]
1234        assert!(!is_safe_to_retry(
1235            &unrelated_unauthorized,
1236            policy,
1237            true,
1238            None,
1239            AccessTokenMode::Refreshable,
1240        ));
1241    }
1242
1243    #[test]
1244    fn safe_to_retry_session_no_side_effects_policy() {
1245        let retryable = server_error(StatusCode::INTERNAL_SERVER_ERROR, "internal");
1246        let no_side_effect = server_error(StatusCode::TOO_MANY_REQUESTS, "rate_limited");
1247        let policy = AppendRetryPolicy::NoSideEffects;
1248        let signal = FrameSignal::new();
1249        let mode = AccessTokenMode::Static;
1250
1251        // No inflight — always safe.
1252        signal.signal();
1253        assert!(is_safe_to_retry(
1254            &retryable,
1255            policy,
1256            false,
1257            Some(&signal),
1258            mode,
1259        ));
1260
1261        // Inflight + signal not set — safe (no data sent this attempt).
1262        signal.reset();
1263        assert!(is_safe_to_retry(
1264            &retryable,
1265            policy,
1266            true,
1267            Some(&signal),
1268            mode,
1269        ));
1270
1271        // Inflight + signal set + error with possible side effects — not safe.
1272        signal.signal();
1273        assert!(!is_safe_to_retry(
1274            &retryable,
1275            policy,
1276            true,
1277            Some(&signal),
1278            mode,
1279        ));
1280
1281        // Inflight + signal set + no-side-effect error — safe.
1282        assert!(is_safe_to_retry(
1283            &no_side_effect,
1284            policy,
1285            true,
1286            Some(&signal),
1287            mode,
1288        ));
1289
1290        // AckTimeout — retryable but has possible side effects.
1291        assert!(!is_safe_to_retry(
1292            &AppendSessionError::AckTimeout,
1293            policy,
1294            true,
1295            Some(&signal),
1296            mode,
1297        ));
1298    }
1299}
s2_sdk/session/append.rs

s2_sdk/session/
append.rs
