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

s2_sdk/session/
append.rs
