Skip to main content

scopedb_client/
append_stream.rs

1// Copyright 2024 ScopeDB, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::error::Error as StdError;
16use std::fmt;
17use std::panic::AssertUnwindSafe;
18use std::sync::Arc;
19use std::sync::Mutex;
20use std::sync::MutexGuard;
21use std::time::Duration;
22use std::time::Instant;
23use std::time::SystemTime;
24
25use serde::Serialize;
26use tokio::sync::Notify;
27use tokio::sync::OwnedSemaphorePermit;
28use tokio::sync::Semaphore;
29use tokio::sync::mpsc;
30use tokio::sync::oneshot;
31use tokio::task::JoinHandle;
32use tokio::task::JoinSet;
33
34use crate::AppendErrorDetails;
35use crate::AppendRowsResult;
36use crate::AppendState;
37use crate::Client;
38use crate::Error;
39use crate::ErrorKind;
40use crate::client::MAX_APPEND_BODY_BYTES;
41use crate::client::MAX_APPEND_ROWS;
42
43const DEFAULT_BATCH_BYTES: usize = 4 * 1024 * 1024;
44const DEFAULT_MAX_BATCH_ROWS: usize = MAX_APPEND_ROWS;
45const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_secs(1);
46const DEFAULT_CHANNEL_CAPACITY: usize = 1024;
47const DEFAULT_MAX_IN_FLIGHT_REQUESTS: usize = 4;
48const DEFAULT_MAX_PENDING_BYTES: usize = 64 * 1024 * 1024;
49const DEFAULT_MAX_RETRIES: usize = 8;
50const DEFAULT_INITIAL_BACKOFF: Duration = Duration::from_millis(100);
51const DEFAULT_MAX_BACKOFF: Duration = Duration::from_secs(5);
52const DEFAULT_CIRCUIT_FAILURE_THRESHOLD: usize = 5;
53const DEFAULT_CIRCUIT_COOLDOWN: Duration = Duration::from_secs(30);
54
55/// Determines what an append stream does after one HTTP batch fails.
56#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
57#[non_exhaustive]
58pub enum AppendFailurePolicy {
59    /// Stop accepting rows and fail the stream's delivery barrier.
60    #[default]
61    Stop,
62    /// Record the loss in the delivery report and continue with later batches.
63    Continue,
64}
65
66/// The lifecycle state of an append stream.
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68#[non_exhaustive]
69pub enum AppendStreamState {
70    Open,
71    Closing,
72    Closed,
73    Failed,
74}
75
76/// The state of the continue-mode availability circuit breaker.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78#[non_exhaustive]
79pub enum AppendCircuitState {
80    Closed,
81    Open,
82    HalfOpen,
83}
84
85/// The aggregate result of a delivery barrier.
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87#[non_exhaustive]
88pub enum AppendDeliveryOutcome {
89    /// Every covered row was committed and no local rows were dropped.
90    Ok,
91    /// Some covered rows committed and some were lost or have an unknown outcome.
92    Partial,
93    /// No covered row committed and all losses are known not to have committed.
94    Failed,
95    /// No covered row is known to have committed and at least one may have committed.
96    Unknown,
97}
98
99/// Rows accepted by a bulk local-admission operation.
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub struct AppendAdmissionResult {
102    pub accepted_rows: u64,
103}
104
105/// Aggregate settlement information since the previous delivery barrier.
106/// Counts do not identify individual input records. Use application-owned
107/// business keys to reconcile rows before replaying them.
108#[derive(Debug, Clone, PartialEq, Eq)]
109pub struct AppendDeliveryReport {
110    pub outcome: AppendDeliveryOutcome,
111    pub accepted_rows: u64,
112    /// Logical input rows acknowledged as committed at least once. Retrying an
113    /// unknown outcome may insert additional copies that this count cannot measure.
114    pub committed_rows: u64,
115    pub failed_rows: u64,
116    pub unknown_rows: u64,
117    pub dropped_rows: u64,
118    pub committed_batches: u64,
119    pub failed_batches: u64,
120    pub unknown_batches: u64,
121    pub retries: u64,
122    /// Wall time spent waiting for this barrier.
123    /// A shared terminal failure report starts at the earliest outstanding barrier.
124    pub duration: Duration,
125}
126
127/// Local drop counters grouped by admission failure reason.
128#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
129pub struct AppendDroppedRows {
130    pub buffer_full: u64,
131    pub invalid_record: u64,
132    pub record_too_large: u64,
133    pub circuit_open: u64,
134    pub closed: u64,
135}
136
137/// The most recent remote append failure observed by the stream.
138#[derive(Debug, Clone)]
139pub struct AppendLastFailure {
140    pub at: SystemTime,
141    pub error: Error,
142}
143
144/// A point-in-time snapshot of append stream counters and state.
145#[derive(Debug, Clone)]
146pub struct AppendStreamStats {
147    pub state: AppendStreamState,
148    pub circuit_state: AppendCircuitState,
149    pub accepted_rows: u64,
150    pub committed_rows: u64,
151    pub failed_rows: u64,
152    pub unknown_rows: u64,
153    pub dropped_rows: u64,
154    pub dropped_by_reason: AppendDroppedRows,
155    pub retries: u64,
156    pub pending_rows: u64,
157    pub pending_bytes: usize,
158    pub in_flight_batches: usize,
159    pub last_failure: Option<AppendLastFailure>,
160    pub last_report: Option<AppendDeliveryReport>,
161}
162
163/// Action taken after a batch failure listener is invoked.
164#[derive(Debug, Clone, Copy, PartialEq, Eq)]
165#[non_exhaustive]
166pub enum AppendBatchFailureAction {
167    Continuing,
168    CircuitOpened,
169    Stopped,
170}
171
172/// Diagnostic event emitted after a rejected or ambiguous append batch.
173pub struct AppendBatchFailureEvent {
174    pub error: Error,
175    pub batch_rows: usize,
176    pub batch_bytes: usize,
177    pub outcome: AppendState,
178    pub action: AppendBatchFailureAction,
179}
180
181impl fmt::Debug for AppendBatchFailureEvent {
182    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183        f.debug_struct("AppendBatchFailureEvent")
184            .field("error", &self.error)
185            .field("batch_rows", &self.batch_rows)
186            .field("batch_bytes", &self.batch_bytes)
187            .field("outcome", &self.outcome)
188            .field("action", &self.action)
189            .finish()
190    }
191}
192
193/// Continue-mode circuit breaker settings.
194#[derive(Debug, Clone, Copy, PartialEq, Eq)]
195#[non_exhaustive]
196pub struct AppendCircuitBreakerOptions {
197    pub failure_threshold: usize,
198    pub cooldown: Duration,
199}
200
201impl AppendCircuitBreakerOptions {
202    pub fn new(failure_threshold: usize, cooldown: Duration) -> Self {
203        Self {
204            failure_threshold,
205            cooldown,
206        }
207    }
208}
209
210impl Default for AppendCircuitBreakerOptions {
211    fn default() -> Self {
212        Self::new(DEFAULT_CIRCUIT_FAILURE_THRESHOLD, DEFAULT_CIRCUIT_COOLDOWN)
213    }
214}
215
216/// A synchronous local-admission failure from [`AppendStream::try_send`].
217#[derive(Debug)]
218#[non_exhaustive]
219pub enum AppendTrySendError {
220    Serialization(serde_json::Error),
221    InvalidRecord,
222    RecordTooLarge { bytes: usize, limit: usize },
223    BufferFull,
224    CircuitOpen,
225    Closed,
226}
227
228impl fmt::Display for AppendTrySendError {
229    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
230        match self {
231            Self::Serialization(err) => write!(f, "failed to serialize append row: {err}"),
232            Self::InvalidRecord => write!(f, "append row must serialize to a JSON object"),
233            Self::RecordTooLarge { bytes, limit } => {
234                write!(
235                    f,
236                    "append row is {bytes} bytes, exceeding the {limit}-byte limit"
237                )
238            }
239            Self::BufferFull => write!(f, "append stream local buffer is full"),
240            Self::CircuitOpen => write!(f, "append stream circuit breaker is open"),
241            Self::Closed => write!(f, "append stream is closed"),
242        }
243    }
244}
245
246impl StdError for AppendTrySendError {
247    fn source(&self) -> Option<&(dyn StdError + 'static)> {
248        match self {
249            Self::Serialization(err) => Some(err),
250            _ => None,
251        }
252    }
253}
254
255type BatchFailureListener = Arc<dyn Fn(&AppendBatchFailureEvent) + Send + Sync>;
256
257#[derive(Debug, Clone, Copy)]
258struct RetryConfig {
259    max_retries: usize,
260    retry_unknown: bool,
261    initial_backoff: Duration,
262    max_backoff: Duration,
263}
264
265/// Configures a bounded, concurrent table append stream.
266///
267/// Retries temporary rejections by default. Replaying unknown outcomes requires
268/// [`Self::retry_unknown`] because an unacknowledged attempt may have committed.
269///
270/// # Telemetry
271///
272/// For a log or telemetry sink that should keep running after batch failures,
273/// use [`AppendFailurePolicy::Continue`] and keep the circuit breaker enabled.
274/// Start with one concurrent batch per stream when producers run on multiple
275/// application replicas; tune the total concurrency across all producers.
276///
277/// ```no_run
278/// use scopedb_client::AppendFailurePolicy;
279///
280/// # async fn example(client: scopedb_client::Client) -> Result<(), scopedb_client::Error> {
281/// let stream = client.table("events").append_stream()
282///     .failure_policy(AppendFailurePolicy::Continue)
283///     .max_concurrent_batches(1)
284///     .build()?;
285///
286/// if let Err(error) = stream.try_send(&serde_json::json!({"name": "request.completed"})) {
287///     eprintln!("telemetry row not admitted: {error}");
288/// }
289/// let report = stream.shutdown().await?;
290/// println!("Write results: {report:?}");
291/// # Ok(())
292/// # }
293/// ```
294///
295/// Continue mode is best effort. Monitor reports and [`AppendStream::stats`]
296/// for failed, unknown, and locally dropped rows. Send failure diagnostics to
297/// a separate sink. [`AppendStream::try_send`] never waits for capacity;
298/// [`AppendStream::send`] provides backpressure instead.
299pub struct AppendStreamBuilder {
300    client: Client,
301    database: String,
302    schema: String,
303    table: String,
304    batch_bytes: usize,
305    max_batch_rows: usize,
306    flush_interval: Duration,
307    channel_capacity: usize,
308    max_pending_bytes: usize,
309    max_in_flight_requests: usize,
310    retry: RetryConfig,
311    failure_policy: AppendFailurePolicy,
312    attempt_timeout: Option<Duration>,
313    circuit_breaker: Option<AppendCircuitBreakerOptions>,
314    batch_failure_listeners: Vec<BatchFailureListener>,
315}
316
317impl AppendStreamBuilder {
318    pub(crate) fn new(client: Client, database: String, schema: String, table: String) -> Self {
319        Self {
320            client,
321            database,
322            schema,
323            table,
324            batch_bytes: DEFAULT_BATCH_BYTES,
325            max_batch_rows: DEFAULT_MAX_BATCH_ROWS,
326            flush_interval: DEFAULT_FLUSH_INTERVAL,
327            channel_capacity: DEFAULT_CHANNEL_CAPACITY,
328            max_pending_bytes: DEFAULT_MAX_PENDING_BYTES,
329            max_in_flight_requests: DEFAULT_MAX_IN_FLIGHT_REQUESTS,
330            retry: RetryConfig {
331                max_retries: DEFAULT_MAX_RETRIES,
332                retry_unknown: false,
333                initial_backoff: DEFAULT_INITIAL_BACKOFF,
334                max_backoff: DEFAULT_MAX_BACKOFF,
335            },
336            failure_policy: AppendFailurePolicy::Stop,
337            attempt_timeout: None,
338            circuit_breaker: Some(AppendCircuitBreakerOptions::default()),
339            batch_failure_listeners: Vec::new(),
340        }
341    }
342
343    /// Chooses whether to stop or continue after a batch fails. Defaults to Stop.
344    pub fn failure_policy(mut self, failure_policy: AppendFailurePolicy) -> Self {
345        self.failure_policy = failure_policy;
346        self
347    }
348
349    /// Sets the target NDJSON payload size (default 4 MiB, maximum 8 MiB).
350    /// One row may exceed this target up to 8 MiB.
351    pub fn target_batch_bytes(mut self, target_batch_bytes: usize) -> Self {
352        self.batch_bytes = target_batch_bytes;
353        self
354    }
355
356    /// Sets the maximum rows in one HTTP batch, up to the 200,000-row protocol limit.
357    pub fn max_batch_rows(mut self, max_batch_rows: usize) -> Self {
358        self.max_batch_rows = max_batch_rows;
359        self
360    }
361
362    /// Sets the maximum delay between buffering the first row and dispatching its batch.
363    pub fn flush_interval(mut self, flush_interval: Duration) -> Self {
364        self.flush_interval = flush_interval;
365        self
366    }
367
368    pub fn channel_capacity(mut self, channel_capacity: usize) -> Self {
369        self.channel_capacity = channel_capacity;
370        self
371    }
372
373    /// Sets the maximum bytes admitted locally but not yet settled remotely.
374    pub fn max_buffered_bytes(mut self, max_buffered_bytes: usize) -> Self {
375        self.max_pending_bytes = max_buffered_bytes;
376        self
377    }
378
379    /// Sets the maximum number of append batches sent concurrently per stream (default 4).
380    pub fn max_concurrent_batches(mut self, max_concurrent_batches: usize) -> Self {
381        self.max_in_flight_requests = max_concurrent_batches;
382        self
383    }
384
385    pub fn max_retries(mut self, max_retries: usize) -> Self {
386        self.retry.max_retries = max_retries;
387        self
388    }
389
390    /// Opts into replaying transient unknown outcomes. This can duplicate rows.
391    /// Client timeouts are never replayed automatically, even when enabled.
392    pub fn retry_unknown(mut self, enabled: bool) -> Self {
393        self.retry.retry_unknown = enabled;
394        self
395    }
396
397    pub fn initial_backoff(mut self, initial_backoff: Duration) -> Self {
398        self.retry.initial_backoff = initial_backoff;
399        self
400    }
401
402    pub fn max_backoff(mut self, max_backoff: Duration) -> Self {
403        self.retry.max_backoff = max_backoff;
404        self
405    }
406
407    /// Sets a timeout for each HTTP attempt, overriding the client's write timeout.
408    /// A timeout has an unknown commit outcome. By default the client uses 30 seconds.
409    /// The duration must be finite and nonzero; this deadline cannot be disabled.
410    pub fn attempt_timeout(mut self, attempt_timeout: Duration) -> Self {
411        self.attempt_timeout = Some(attempt_timeout);
412        self
413    }
414
415    /// Configures the continue-mode availability circuit, or disables it with `None`.
416    /// Enabled by default: five consecutive availability failures pause dispatch
417    /// for 30 seconds before probing recovery.
418    pub fn circuit_breaker(mut self, circuit_breaker: Option<AppendCircuitBreakerOptions>) -> Self {
419        self.circuit_breaker = circuit_breaker;
420        self
421    }
422
423    /// Observes remote batch failures. Listener panics are isolated from the worker.
424    pub fn on_batch_failure<F>(mut self, listener: F) -> Self
425    where
426        F: Fn(&AppendBatchFailureEvent) + Send + Sync + 'static,
427    {
428        self.batch_failure_listeners.push(Arc::new(listener));
429        self
430    }
431
432    pub fn build(self) -> Result<AppendStream, Error> {
433        validate_builder(&self)?;
434        tokio::runtime::Handle::try_current().map_err(|err| {
435            Error::new(
436                ErrorKind::ConfigInvalid,
437                "append stream must be built inside a Tokio runtime",
438            )
439            .set_source(err)
440        })?;
441
442        Ok(AppendStream::new(AppendStreamConfig {
443            client: self.client,
444            database: self.database,
445            schema: self.schema,
446            table: self.table,
447            batch_bytes: self.batch_bytes,
448            max_batch_rows: self.max_batch_rows,
449            flush_interval: self.flush_interval,
450            channel_capacity: self.channel_capacity,
451            max_pending_bytes: self.max_pending_bytes,
452            max_in_flight_requests: self.max_in_flight_requests,
453            retry: self.retry,
454            failure_policy: self.failure_policy,
455            attempt_timeout: self.attempt_timeout,
456            circuit_breaker: self.circuit_breaker,
457            batch_failure_listeners: self.batch_failure_listeners,
458        }))
459    }
460}
461
462fn validate_builder(builder: &AppendStreamBuilder) -> Result<(), Error> {
463    builder.client.make_resource_url(&[
464        "catalog",
465        "databases",
466        &builder.database,
467        "schemas",
468        &builder.schema,
469        "tables",
470        &builder.table,
471        "ingest",
472    ])?;
473    validate_deadline("initial_backoff", builder.retry.initial_backoff)?;
474    validate_deadline("max_backoff", builder.retry.max_backoff)?;
475    if builder.channel_capacity > Semaphore::MAX_PERMITS {
476        return Err(config_error("channel_capacity is too large"));
477    }
478    validate_positive("target_batch_bytes", builder.batch_bytes)?;
479    if builder.batch_bytes > MAX_APPEND_BODY_BYTES {
480        return Err(config_error(format!(
481            "target_batch_bytes must not exceed {MAX_APPEND_BODY_BYTES}"
482        )));
483    }
484    validate_positive("max_batch_rows", builder.max_batch_rows)?;
485    if builder.max_batch_rows > MAX_APPEND_ROWS {
486        return Err(config_error(format!(
487            "max_batch_rows must not exceed {MAX_APPEND_ROWS}"
488        )));
489    }
490    if builder.flush_interval.is_zero() {
491        return Err(config_error("flush_interval must be greater than zero"));
492    }
493    validate_deadline("flush_interval", builder.flush_interval)?;
494    validate_positive("channel_capacity", builder.channel_capacity)?;
495    validate_positive("max_buffered_bytes", builder.max_pending_bytes)?;
496    if builder.max_pending_bytes > (u32::MAX as usize).min(Semaphore::MAX_PERMITS) {
497        return Err(config_error(format!(
498            "max_buffered_bytes must not exceed {}",
499            u32::MAX
500        )));
501    }
502    validate_positive("max_concurrent_batches", builder.max_in_flight_requests)?;
503    if builder
504        .attempt_timeout
505        .is_some_and(|attempt_timeout| attempt_timeout.is_zero())
506    {
507        return Err(config_error("attempt_timeout must be greater than zero"));
508    }
509    if let Some(attempt_timeout) = builder.attempt_timeout {
510        validate_deadline("attempt_timeout", attempt_timeout)?;
511    }
512    if let Some(circuit) = builder.circuit_breaker {
513        validate_positive(
514            "circuit_breaker.failure_threshold",
515            circuit.failure_threshold,
516        )?;
517        if circuit.cooldown.is_zero() {
518            return Err(config_error(
519                "circuit_breaker.cooldown must be greater than zero",
520            ));
521        }
522        validate_deadline("circuit_breaker.cooldown", circuit.cooldown)?;
523    }
524    Ok(())
525}
526
527fn validate_deadline(name: &str, duration: Duration) -> Result<(), Error> {
528    if Instant::now().checked_add(duration).is_none() {
529        Err(config_error(format!("{name} is too large")))
530    } else {
531        Ok(())
532    }
533}
534
535fn validate_positive(name: &str, value: usize) -> Result<(), Error> {
536    if value == 0 {
537        Err(config_error(format!("{name} must be greater than zero")))
538    } else {
539        Ok(())
540    }
541}
542
543fn config_error(message: impl Into<String>) -> Error {
544    Error::new(ErrorKind::ConfigInvalid, message)
545}
546
547struct AppendStreamConfig {
548    client: Client,
549    database: String,
550    schema: String,
551    table: String,
552    batch_bytes: usize,
553    max_batch_rows: usize,
554    flush_interval: Duration,
555    channel_capacity: usize,
556    max_pending_bytes: usize,
557    max_in_flight_requests: usize,
558    retry: RetryConfig,
559    failure_policy: AppendFailurePolicy,
560    attempt_timeout: Option<Duration>,
561    circuit_breaker: Option<AppendCircuitBreakerOptions>,
562    batch_failure_listeners: Vec<BatchFailureListener>,
563}
564
565enum AppendCommand {
566    Record(BufferedRecord),
567    Flush(BarrierCommand),
568    Shutdown(BarrierCommand),
569}
570
571struct BarrierCommand {
572    dropped_rows_at_barrier: u64,
573    started_at: Instant,
574    ack: oneshot::Sender<Result<AppendDeliveryReport, Error>>,
575}
576
577struct BufferedRecord {
578    payload: String,
579    reserved_bytes: usize,
580    _reservation: OwnedSemaphorePermit,
581}
582
583struct PendingBytesBudget {
584    capacity: usize,
585    semaphore: Arc<Semaphore>,
586}
587
588impl PendingBytesBudget {
589    fn new(capacity: usize) -> Self {
590        Self {
591            capacity,
592            semaphore: Arc::new(Semaphore::new(capacity)),
593        }
594    }
595
596    async fn acquire(&self, bytes: usize) -> Result<OwnedSemaphorePermit, PendingAcquireError> {
597        if bytes > self.capacity {
598            return Err(PendingAcquireError::ExceedsCapacity);
599        }
600        self.semaphore
601            .clone()
602            .acquire_many_owned(bytes as u32)
603            .await
604            .map_err(|_| PendingAcquireError::Closed)
605    }
606
607    fn try_acquire(&self, bytes: usize) -> Result<OwnedSemaphorePermit, PendingAcquireError> {
608        if bytes > self.capacity {
609            return Err(PendingAcquireError::ExceedsCapacity);
610        }
611        self.semaphore
612            .clone()
613            .try_acquire_many_owned(bytes as u32)
614            .map_err(|err| match err {
615                tokio::sync::TryAcquireError::Closed => PendingAcquireError::Closed,
616                tokio::sync::TryAcquireError::NoPermits => PendingAcquireError::Full,
617            })
618    }
619
620    fn close(&self) {
621        self.semaphore.close();
622    }
623}
624
625#[derive(Debug, Clone, Copy, PartialEq, Eq)]
626enum PendingAcquireError {
627    Closed,
628    Full,
629    ExceedsCapacity,
630}
631
632#[derive(Debug, Clone, Copy, Default)]
633struct DeliveryCounters {
634    accepted_rows: u64,
635    committed_rows: u64,
636    failed_rows: u64,
637    unknown_rows: u64,
638    committed_batches: u64,
639    failed_batches: u64,
640    unknown_batches: u64,
641    retries: u64,
642}
643
644struct CircuitRuntime {
645    state: AppendCircuitState,
646    opened_until: Option<Instant>,
647    consecutive_failures: usize,
648    probe_admission_claimed: bool,
649}
650
651impl Default for CircuitRuntime {
652    fn default() -> Self {
653        Self {
654            state: AppendCircuitState::Closed,
655            opened_until: None,
656            consecutive_failures: 0,
657            probe_admission_claimed: false,
658        }
659    }
660}
661
662struct SharedState {
663    lifecycle: AppendStreamState,
664    fatal: Option<Error>,
665    final_result: Option<Result<AppendDeliveryReport, Error>>,
666    accepted_rows: u64,
667    committed_rows: u64,
668    failed_rows: u64,
669    unknown_rows: u64,
670    dropped_rows: u64,
671    dropped_by_reason: AppendDroppedRows,
672    retries: u64,
673    pending_rows: u64,
674    pending_bytes: usize,
675    in_flight_batches: usize,
676    last_failure: Option<AppendLastFailure>,
677    last_report: Option<AppendDeliveryReport>,
678    circuit: CircuitRuntime,
679}
680
681impl Default for SharedState {
682    fn default() -> Self {
683        Self {
684            lifecycle: AppendStreamState::Open,
685            fatal: None,
686            final_result: None,
687            accepted_rows: 0,
688            committed_rows: 0,
689            failed_rows: 0,
690            unknown_rows: 0,
691            dropped_rows: 0,
692            dropped_by_reason: AppendDroppedRows::default(),
693            retries: 0,
694            pending_rows: 0,
695            pending_bytes: 0,
696            in_flight_batches: 0,
697            last_failure: None,
698            last_report: None,
699            circuit: CircuitRuntime::default(),
700        }
701    }
702}
703
704struct AppendStreamInner {
705    tx: mpsc::Sender<AppendCommand>,
706    shared: Arc<Mutex<SharedState>>,
707    pending_bytes: Arc<PendingBytesBudget>,
708    shutdown_notify: Arc<Notify>,
709}
710
711/// A cloneable, bounded producer handle for asynchronous table appends.
712/// Dropping every handle drains admitted rows in the background while the Tokio
713/// runtime remains alive. Use shutdown() to observe settlement before exiting.
714#[derive(Clone)]
715pub struct AppendStream {
716    inner: Arc<AppendStreamInner>,
717}
718
719impl AppendStream {
720    fn new(config: AppendStreamConfig) -> Self {
721        let (tx, rx) = mpsc::channel(config.channel_capacity);
722        let shared = Arc::new(Mutex::new(SharedState::default()));
723        let pending_bytes = Arc::new(PendingBytesBudget::new(config.max_pending_bytes));
724        let shutdown_notify = Arc::new(Notify::new());
725
726        let worker = tokio::spawn(run_append_worker(
727            config,
728            rx,
729            shared.clone(),
730            pending_bytes.clone(),
731            shutdown_notify.clone(),
732        ));
733        tokio::spawn(monitor_append_worker(
734            worker,
735            shared.clone(),
736            pending_bytes.clone(),
737            shutdown_notify.clone(),
738        ));
739
740        Self {
741            inner: Arc::new(AppendStreamInner {
742                tx,
743                shared,
744                pending_bytes,
745                shutdown_notify,
746            }),
747        }
748    }
749
750    /// Serializes and admits one row, waiting for bounded local capacity.
751    ///
752    /// Success means local admission only. Use [`Self::flush`] or
753    /// [`Self::shutdown`] to wait for remote settlement.
754    pub async fn send<T>(&self, row: &T) -> Result<(), Error>
755    where
756        T: Serialize + ?Sized,
757    {
758        self.check_open()?;
759        let payload = prepare_record(row)?;
760        let reserved_bytes = payload.len().saturating_add(1);
761        let reservation = self
762            .inner
763            .pending_bytes
764            .acquire(reserved_bytes)
765            .await
766            .map_err(|err| self.map_pending_error(err, reserved_bytes))?;
767        let permit = match self.inner.tx.reserve().await {
768            Ok(permit) => permit,
769            Err(_) => {
770                drop(reservation);
771                return Err(self.closed_or_fatal_error());
772            }
773        };
774
775        let mut shared = lock(&self.inner.shared);
776        if shared.lifecycle != AppendStreamState::Open {
777            drop(shared);
778            drop(permit);
779            drop(reservation);
780            return Err(self.closed_or_fatal_error());
781        }
782        shared.accepted_rows = shared.accepted_rows.saturating_add(1);
783        shared.pending_rows = shared.pending_rows.saturating_add(1);
784        shared.pending_bytes = shared.pending_bytes.saturating_add(reserved_bytes);
785        permit.send(AppendCommand::Record(BufferedRecord {
786            payload,
787            reserved_bytes,
788            _reservation: reservation,
789        }));
790        Ok(())
791    }
792
793    /// Attempts synchronous local admission without waiting for capacity.
794    /// Returns an error if the circuit is open or local capacity is full.
795    pub fn try_send<T>(&self, row: &T) -> Result<(), AppendTrySendError>
796    where
797        T: Serialize + ?Sized,
798    {
799        {
800            let shared = lock(&self.inner.shared);
801            if shared.lifecycle != AppendStreamState::Open {
802                drop(shared);
803                note_drop(&mut lock(&self.inner.shared), DropReason::Closed);
804                return Err(AppendTrySendError::Closed);
805            }
806            if !circuit_admission_available(&shared) {
807                drop(shared);
808                note_drop(&mut lock(&self.inner.shared), DropReason::CircuitOpen);
809                return Err(AppendTrySendError::CircuitOpen);
810            }
811        }
812
813        let payload = match prepare_record_for_try_send(row) {
814            Ok(payload) => payload,
815            Err(err) => {
816                let reason = match err {
817                    AppendTrySendError::RecordTooLarge { .. } => DropReason::RecordTooLarge,
818                    _ => DropReason::InvalidRecord,
819                };
820                note_drop(&mut lock(&self.inner.shared), reason);
821                return Err(err);
822            }
823        };
824        let reserved_bytes = payload.len().saturating_add(1);
825        let reservation = match self.inner.pending_bytes.try_acquire(reserved_bytes) {
826            Ok(reservation) => reservation,
827            Err(PendingAcquireError::Closed) => {
828                note_drop(&mut lock(&self.inner.shared), DropReason::Closed);
829                return Err(AppendTrySendError::Closed);
830            }
831            Err(PendingAcquireError::Full) => {
832                note_drop(&mut lock(&self.inner.shared), DropReason::BufferFull);
833                return Err(AppendTrySendError::BufferFull);
834            }
835            Err(PendingAcquireError::ExceedsCapacity) => {
836                note_drop(&mut lock(&self.inner.shared), DropReason::RecordTooLarge);
837                return Err(AppendTrySendError::RecordTooLarge {
838                    bytes: reserved_bytes,
839                    limit: self.inner.pending_bytes.capacity,
840                });
841            }
842        };
843        let permit = match self.inner.tx.try_reserve() {
844            Ok(permit) => permit,
845            Err(mpsc::error::TrySendError::Full(_)) => {
846                drop(reservation);
847                note_drop(&mut lock(&self.inner.shared), DropReason::BufferFull);
848                return Err(AppendTrySendError::BufferFull);
849            }
850            Err(mpsc::error::TrySendError::Closed(_)) => {
851                drop(reservation);
852                note_drop(&mut lock(&self.inner.shared), DropReason::Closed);
853                return Err(AppendTrySendError::Closed);
854            }
855        };
856
857        let mut shared = lock(&self.inner.shared);
858        if shared.lifecycle != AppendStreamState::Open {
859            drop(shared);
860            drop(permit);
861            drop(reservation);
862            note_drop(&mut lock(&self.inner.shared), DropReason::Closed);
863            return Err(AppendTrySendError::Closed);
864        }
865        if !try_claim_circuit_admission(&mut shared) {
866            drop(shared);
867            drop(permit);
868            drop(reservation);
869            note_drop(&mut lock(&self.inner.shared), DropReason::CircuitOpen);
870            return Err(AppendTrySendError::CircuitOpen);
871        }
872        shared.accepted_rows = shared.accepted_rows.saturating_add(1);
873        shared.pending_rows = shared.pending_rows.saturating_add(1);
874        shared.pending_bytes = shared.pending_bytes.saturating_add(reserved_bytes);
875        permit.send(AppendCommand::Record(BufferedRecord {
876            payload,
877            reserved_bytes,
878            _reservation: reservation,
879        }));
880        Ok(())
881    }
882
883    /// Admits an iterator of rows sequentially with bounded memory.
884    ///
885    /// If admission fails partway, the error includes the accepted count.
886    /// Call [`Self::shutdown`] to settle those rows; [`Self::stats`] retains
887    /// cumulative counts.
888    pub async fn send_all<I>(&self, rows: I) -> Result<AppendAdmissionResult, Error>
889    where
890        I: IntoIterator,
891        I::Item: Serialize,
892    {
893        let mut accepted_rows = 0u64;
894        for row in rows {
895            self.send(&row)
896                .await
897                .map_err(|error| error.with_context("accepted_rows", accepted_rows))?;
898            accepted_rows = accepted_rows.saturating_add(1);
899        }
900        Ok(AppendAdmissionResult { accepted_rows })
901    }
902
903    /// Dispatches rows admitted before this barrier and waits for their outcomes.
904    ///
905    /// The report covers activity since the previous barrier. On a Stop-mode
906    /// failure, partial counts remain available in [`Self::stats`].
907    ///
908    /// Committed rows have at least one successful acknowledgement; earlier
909    /// unknown attempts may also have committed. A failed barrier may include
910    /// successful batches, so replaying its entire input can duplicate rows.
911    ///
912    /// Once enqueued, dropping this future does not cancel remote settlement.
913    /// Keep the future alive to receive its interval report; the latest completed
914    /// report also remains observable through [`Self::stats`].
915    pub async fn flush(&self) -> Result<AppendDeliveryReport, Error> {
916        loop {
917            let notified = self.inner.shutdown_notify.notified();
918            match self.current_terminal_result() {
919                TerminalResult::Open => {}
920                TerminalResult::Waiting => {
921                    notified.await;
922                    continue;
923                }
924                TerminalResult::Complete(result) => return result,
925            }
926
927            let permit = match self.inner.tx.reserve().await {
928                Ok(permit) => permit,
929                Err(_) => return self.wait_for_terminal().await,
930            };
931            let (ack, response) = oneshot::channel();
932            {
933                let shared = lock(&self.inner.shared);
934                if shared.lifecycle != AppendStreamState::Open {
935                    drop(shared);
936                    drop(permit);
937                    continue;
938                }
939                permit.send(AppendCommand::Flush(BarrierCommand {
940                    dropped_rows_at_barrier: shared.dropped_rows,
941                    started_at: Instant::now(),
942                    ack,
943                }));
944            }
945            return match response.await {
946                Ok(Ok(report)) => Ok(report),
947                Ok(Err(_)) | Err(_) => self.wait_for_terminal().await,
948            };
949        }
950    }
951
952    /// Permanently closes the stream after settling every admitted row.
953    ///
954    /// This method is idempotent across cloned producer handles. Once polled,
955    /// dropping its future stops waiting but settlement continues in the worker.
956    /// Call shutdown again to observe completion. Failure reports remain in stats().
957    pub async fn shutdown(&self) -> Result<AppendDeliveryReport, Error> {
958        let should_start = {
959            let mut shared = lock(&self.inner.shared);
960            match shared.lifecycle {
961                AppendStreamState::Open => {
962                    shared.lifecycle = AppendStreamState::Closing;
963                    true
964                }
965                AppendStreamState::Closing | AppendStreamState::Failed => false,
966                AppendStreamState::Closed if shared.final_result.is_some() => {
967                    return terminal_result(&shared);
968                }
969                AppendStreamState::Closed => false,
970            }
971        };
972
973        if should_start {
974            let inner = self.inner.clone();
975            tokio::spawn(async move {
976                let _ = enqueue_shutdown(inner).await;
977            });
978        }
979
980        self.wait_for_terminal().await
981    }
982
983    /// Returns a synchronous snapshot suitable for metrics and shutdown diagnostics.
984    pub fn stats(&self) -> AppendStreamStats {
985        let shared = lock(&self.inner.shared);
986        AppendStreamStats {
987            state: shared.lifecycle,
988            circuit_state: shared.circuit.state,
989            accepted_rows: shared.accepted_rows,
990            committed_rows: shared.committed_rows,
991            failed_rows: shared.failed_rows,
992            unknown_rows: shared.unknown_rows,
993            dropped_rows: shared.dropped_rows,
994            dropped_by_reason: shared.dropped_by_reason,
995            retries: shared.retries,
996            pending_rows: shared.pending_rows,
997            pending_bytes: shared.pending_bytes,
998            in_flight_batches: shared.in_flight_batches,
999            last_failure: shared.last_failure.clone(),
1000            last_report: shared.last_report.clone(),
1001        }
1002    }
1003
1004    fn check_open(&self) -> Result<(), Error> {
1005        let shared = lock(&self.inner.shared);
1006        if shared.lifecycle == AppendStreamState::Open {
1007            Ok(())
1008        } else {
1009            Err(error_from_shared(&shared))
1010        }
1011    }
1012
1013    fn closed_or_fatal_error(&self) -> Error {
1014        error_from_shared(&lock(&self.inner.shared))
1015    }
1016
1017    fn map_pending_error(&self, error: PendingAcquireError, requested: usize) -> Error {
1018        match error {
1019            PendingAcquireError::Closed => self.closed_or_fatal_error(),
1020            PendingAcquireError::Full => Error::new(
1021                ErrorKind::Unexpected,
1022                "append stream local byte budget is full",
1023            )
1024            .set_temporary(),
1025            PendingAcquireError::ExceedsCapacity => Error::new(
1026                ErrorKind::Unexpected,
1027                format!(
1028                    "append row requires {requested} buffered bytes, exceeding max_pending_bytes={}",
1029                    self.inner.pending_bytes.capacity
1030                ),
1031            )
1032            .set_permanent(),
1033        }
1034    }
1035
1036    fn current_terminal_result(&self) -> TerminalResult {
1037        let shared = lock(&self.inner.shared);
1038        match shared.lifecycle {
1039            AppendStreamState::Open => TerminalResult::Open,
1040            AppendStreamState::Closing => TerminalResult::Waiting,
1041            AppendStreamState::Closed | AppendStreamState::Failed
1042                if shared.final_result.is_none() =>
1043            {
1044                TerminalResult::Waiting
1045            }
1046            AppendStreamState::Closed | AppendStreamState::Failed => {
1047                TerminalResult::Complete(terminal_result(&shared))
1048            }
1049        }
1050    }
1051
1052    async fn wait_for_terminal(&self) -> Result<AppendDeliveryReport, Error> {
1053        loop {
1054            let notified = self.inner.shutdown_notify.notified();
1055            {
1056                let shared = lock(&self.inner.shared);
1057                if matches!(
1058                    shared.lifecycle,
1059                    AppendStreamState::Closed | AppendStreamState::Failed
1060                ) && shared.final_result.is_some()
1061                {
1062                    return terminal_result(&shared);
1063                }
1064            }
1065            notified.await;
1066        }
1067    }
1068}
1069
1070enum TerminalResult {
1071    Open,
1072    Waiting,
1073    Complete(Result<AppendDeliveryReport, Error>),
1074}
1075
1076async fn enqueue_shutdown(inner: Arc<AppendStreamInner>) -> Result<AppendDeliveryReport, Error> {
1077    let (ack, response) = oneshot::channel();
1078    let dropped_rows_at_barrier = lock(&inner.shared).dropped_rows;
1079    inner
1080        .tx
1081        .send(AppendCommand::Shutdown(BarrierCommand {
1082            dropped_rows_at_barrier,
1083            started_at: Instant::now(),
1084            ack,
1085        }))
1086        .await
1087        .map_err(|_| error_from_shared(&lock(&inner.shared)))?;
1088    response
1089        .await
1090        .unwrap_or_else(|_| Err(error_from_shared(&lock(&inner.shared))))
1091}
1092
1093fn terminal_result(shared: &SharedState) -> Result<AppendDeliveryReport, Error> {
1094    match &shared.final_result {
1095        Some(Ok(report)) => Ok(report.clone()),
1096        Some(Err(error)) => Err(error.clone()),
1097        None => Err(error_from_shared(shared)),
1098    }
1099}
1100
1101fn error_from_shared(shared: &SharedState) -> Error {
1102    shared.fatal.as_ref().cloned().unwrap_or_else(|| {
1103        Error::new(ErrorKind::Unexpected, "append stream is closed").set_persistent()
1104    })
1105}
1106
1107fn prepare_record<T>(row: &T) -> Result<String, Error>
1108where
1109    T: Serialize + ?Sized,
1110{
1111    let payload = serde_json::to_string(row).map_err(|err| {
1112        Error::new(ErrorKind::Unexpected, "failed to serialize append row").set_source(err)
1113    })?;
1114    validate_record_payload(&payload).map_err(|err| match err {
1115        AppendTrySendError::RecordTooLarge { bytes, limit } => Error::new(
1116            ErrorKind::Unexpected,
1117            format!("append row is {bytes} bytes, exceeding the {limit}-byte limit"),
1118        ),
1119        _ => Error::new(
1120            ErrorKind::Unexpected,
1121            "append row must serialize to a JSON object",
1122        ),
1123    })?;
1124    Ok(payload)
1125}
1126
1127fn prepare_record_for_try_send<T>(row: &T) -> Result<String, AppendTrySendError>
1128where
1129    T: Serialize + ?Sized,
1130{
1131    let payload = serde_json::to_string(row).map_err(AppendTrySendError::Serialization)?;
1132    validate_record_payload(&payload)?;
1133    Ok(payload)
1134}
1135
1136fn validate_record_payload(payload: &str) -> Result<(), AppendTrySendError> {
1137    let bytes = payload.len();
1138    if bytes > MAX_APPEND_BODY_BYTES {
1139        return Err(AppendTrySendError::RecordTooLarge {
1140            bytes,
1141            limit: MAX_APPEND_BODY_BYTES,
1142        });
1143    }
1144    if !payload.starts_with('{') {
1145        return Err(AppendTrySendError::InvalidRecord);
1146    }
1147    Ok(())
1148}
1149
1150struct BatchOutcome {
1151    records: Vec<BufferedRecord>,
1152    batch_bytes: usize,
1153    result: Result<AppendRowsResult, Error>,
1154    retries: usize,
1155    circuit_probe: bool,
1156}
1157
1158struct AppendBatchRequest {
1159    client: Client,
1160    database: String,
1161    schema: String,
1162    table: String,
1163    retry: RetryConfig,
1164    attempt_timeout: Option<Duration>,
1165}
1166
1167struct AppendWorker {
1168    config: AppendStreamConfig,
1169    shared: Arc<Mutex<SharedState>>,
1170    pending_bytes: Arc<PendingBytesBudget>,
1171    shutdown_notify: Arc<Notify>,
1172    rows: Vec<BufferedRecord>,
1173    current_bytes: usize,
1174    batch_deadline: Option<Instant>,
1175    in_flight: JoinSet<BatchOutcome>,
1176    interval: DeliveryCounters,
1177    reported_dropped_rows: u64,
1178    fatal: Option<Error>,
1179}
1180
1181async fn run_append_worker(
1182    config: AppendStreamConfig,
1183    mut rx: mpsc::Receiver<AppendCommand>,
1184    shared: Arc<Mutex<SharedState>>,
1185    pending_bytes: Arc<PendingBytesBudget>,
1186    shutdown_notify: Arc<Notify>,
1187) {
1188    let mut worker = AppendWorker {
1189        config,
1190        shared,
1191        pending_bytes,
1192        shutdown_notify,
1193        rows: Vec::new(),
1194        current_bytes: 0,
1195        batch_deadline: None,
1196        in_flight: JoinSet::new(),
1197        interval: DeliveryCounters::default(),
1198        reported_dropped_rows: 0,
1199        fatal: None,
1200    };
1201
1202    loop {
1203        if worker.fatal.is_some() {
1204            worker.finish_fatal(&mut rx, None).await;
1205            return;
1206        }
1207
1208        let deadline = worker
1209            .batch_deadline
1210            .unwrap_or_else(|| Instant::now() + Duration::from_secs(86400 * 365));
1211        tokio::select! {
1212            biased;
1213            outcome = worker.in_flight.join_next(), if !worker.in_flight.is_empty() => {
1214                worker.handle_joined(outcome).await;
1215            }
1216            _ = tokio::time::sleep_until(deadline.into()), if worker.batch_deadline.is_some() => {
1217                worker.dispatch_buffered().await;
1218            }
1219            command = rx.recv() => {
1220                match command {
1221                    Some(AppendCommand::Record(record)) => worker.buffer_record(record).await,
1222                    Some(AppendCommand::Flush(barrier)) => {
1223                        if let Some(barrier) = worker.complete_barrier(barrier).await {
1224                            worker.finish_fatal(&mut rx, Some(barrier)).await;
1225                            return;
1226                        }
1227                    }
1228                    Some(AppendCommand::Shutdown(barrier)) => {
1229                        if let Some(barrier) = worker.complete_barrier(barrier).await {
1230                            worker.finish_fatal(&mut rx, Some(barrier)).await;
1231                        } else {
1232                            worker.finish_closed();
1233                        }
1234                        return;
1235                    }
1236                    None => {
1237                        worker.dispatch_buffered().await;
1238                        worker.drain_in_flight().await;
1239                        if worker.fatal.is_some() {
1240                            worker.finish_fatal(&mut rx, None).await;
1241                        } else {
1242                            let dropped = lock(&worker.shared).dropped_rows;
1243                            let report = worker.take_report(dropped, Instant::now());
1244                            worker.finish_closed_with_report(report);
1245                        }
1246                        return;
1247                    }
1248                }
1249            }
1250        }
1251    }
1252}
1253
1254async fn monitor_append_worker(
1255    worker: JoinHandle<()>,
1256    shared: Arc<Mutex<SharedState>>,
1257    pending_bytes: Arc<PendingBytesBudget>,
1258    shutdown_notify: Arc<Notify>,
1259) {
1260    let worker_result = worker.await;
1261    let mut shared = lock(&shared);
1262    if shared.final_result.is_some() {
1263        return;
1264    }
1265
1266    let error = match worker_result {
1267        Ok(()) => Error::new(
1268            ErrorKind::Unexpected,
1269            "append stream worker stopped before publishing a final result",
1270        )
1271        .set_persistent(),
1272        Err(error) => {
1273            unknown_append_error("append stream worker terminated unexpectedly", Some(error))
1274        }
1275    };
1276    let snapshot = error.clone();
1277    let resolved = shared
1278        .committed_rows
1279        .saturating_add(shared.failed_rows)
1280        .saturating_add(shared.unknown_rows);
1281    let unresolved = shared.accepted_rows.saturating_sub(resolved);
1282    shared.unknown_rows = shared.unknown_rows.saturating_add(unresolved);
1283    shared.pending_rows = 0;
1284    shared.pending_bytes = 0;
1285    shared.in_flight_batches = 0;
1286    shared.lifecycle = AppendStreamState::Failed;
1287    shared.fatal = Some(snapshot.clone());
1288    shared.final_result = Some(Err(snapshot));
1289    drop(shared);
1290    pending_bytes.close();
1291    shutdown_notify.notify_waiters();
1292}
1293
1294impl AppendWorker {
1295    async fn buffer_record(&mut self, record: BufferedRecord) {
1296        self.interval.accepted_rows = self.interval.accepted_rows.saturating_add(1);
1297        let separator = usize::from(!self.rows.is_empty());
1298        if !self.rows.is_empty()
1299            && self
1300                .current_bytes
1301                .saturating_add(separator)
1302                .saturating_add(record.payload.len())
1303                > self.config.batch_bytes
1304        {
1305            self.dispatch_buffered().await;
1306        }
1307        if self.fatal.is_some() {
1308            self.mark_records_failed(vec![record]);
1309            return;
1310        }
1311
1312        if self.rows.is_empty() {
1313            self.batch_deadline = Some(deadline_after(self.config.flush_interval));
1314        } else {
1315            self.current_bytes = self.current_bytes.saturating_add(1);
1316        }
1317        self.current_bytes = self.current_bytes.saturating_add(record.payload.len());
1318        self.rows.push(record);
1319        if self.current_bytes >= self.config.batch_bytes
1320            || self.rows.len() >= self.config.max_batch_rows
1321        {
1322            self.dispatch_buffered().await;
1323        }
1324    }
1325
1326    async fn dispatch_buffered(&mut self) {
1327        if self.rows.is_empty() || self.fatal.is_some() {
1328            return;
1329        }
1330        while self.in_flight.len() >= self.config.max_in_flight_requests {
1331            let outcome = self.in_flight.join_next().await;
1332            self.handle_joined(outcome).await;
1333            if self.fatal.is_some() {
1334                return;
1335            }
1336        }
1337
1338        let circuit_probe = self.wait_for_circuit().await;
1339        if self.fatal.is_some() {
1340            return;
1341        }
1342        let records = std::mem::take(&mut self.rows);
1343        let batch_bytes = std::mem::take(&mut self.current_bytes);
1344        self.batch_deadline = None;
1345        let payload = records
1346            .iter()
1347            .map(|record| record.payload.as_str())
1348            .collect::<Vec<_>>()
1349            .join("\n");
1350        let request = AppendBatchRequest {
1351            client: self.config.client.clone(),
1352            database: self.config.database.clone(),
1353            schema: self.config.schema.clone(),
1354            table: self.config.table.clone(),
1355            retry: self.config.retry,
1356            attempt_timeout: self.config.attempt_timeout,
1357        };
1358        lock(&self.shared).in_flight_batches += 1;
1359        self.in_flight.spawn(async move {
1360            let (result, retries) = append_batch(request, payload).await;
1361            BatchOutcome {
1362                records,
1363                batch_bytes,
1364                result,
1365                retries,
1366                circuit_probe,
1367            }
1368        });
1369    }
1370
1371    async fn wait_for_circuit(&mut self) -> bool {
1372        if self.config.failure_policy != AppendFailurePolicy::Continue
1373            || self.config.circuit_breaker.is_none()
1374        {
1375            return false;
1376        }
1377
1378        loop {
1379            let (state, remaining) = {
1380                let shared = lock(&self.shared);
1381                let remaining = shared
1382                    .circuit
1383                    .opened_until
1384                    .and_then(|until| until.checked_duration_since(Instant::now()));
1385                (shared.circuit.state, remaining)
1386            };
1387            match state {
1388                AppendCircuitState::Closed => return false,
1389                AppendCircuitState::Open => {
1390                    if let Some(remaining) = remaining {
1391                        if self.in_flight.is_empty() {
1392                            tokio::time::sleep(remaining).await;
1393                        } else {
1394                            tokio::select! {
1395                                _ = tokio::time::sleep(remaining) => {}
1396                                outcome = self.in_flight.join_next() => {
1397                                    self.handle_joined(outcome).await;
1398                                    if self.fatal.is_some() {
1399                                        return false;
1400                                    }
1401                                }
1402                            }
1403                        }
1404                        continue;
1405                    }
1406                    if !self.in_flight.is_empty() {
1407                        self.drain_in_flight().await;
1408                        if self.fatal.is_some() {
1409                            return false;
1410                        }
1411                        continue;
1412                    }
1413                    let mut shared = lock(&self.shared);
1414                    if shared.circuit.state == AppendCircuitState::Open {
1415                        shared.circuit.state = AppendCircuitState::HalfOpen;
1416                        return true;
1417                    }
1418                }
1419                AppendCircuitState::HalfOpen => {
1420                    if self.in_flight.is_empty() {
1421                        return true;
1422                    }
1423                    let outcome = self.in_flight.join_next().await;
1424                    self.handle_joined(outcome).await;
1425                    if self.fatal.is_some() {
1426                        return false;
1427                    }
1428                }
1429            }
1430        }
1431    }
1432
1433    async fn complete_barrier(&mut self, barrier: BarrierCommand) -> Option<BarrierCommand> {
1434        self.dispatch_buffered().await;
1435        self.drain_in_flight().await;
1436        if self.fatal.is_some() {
1437            return Some(barrier);
1438        }
1439        let report = self.take_report(barrier.dropped_rows_at_barrier, barrier.started_at);
1440        let _ = barrier.ack.send(Ok(report));
1441        None
1442    }
1443
1444    async fn drain_in_flight(&mut self) {
1445        while let Some(outcome) = self.in_flight.join_next().await {
1446            self.handle_joined(Some(outcome)).await;
1447        }
1448    }
1449
1450    async fn handle_joined(
1451        &mut self,
1452        outcome: Option<Result<BatchOutcome, tokio::task::JoinError>>,
1453    ) {
1454        let Some(outcome) = outcome else {
1455            return;
1456        };
1457        let outcome = match outcome {
1458            Ok(outcome) => outcome,
1459            Err(error) => {
1460                let fatal = unknown_append_error(
1461                    "append stream background request task failed",
1462                    Some(error),
1463                );
1464                self.set_fatal(fatal.clone());
1465                return;
1466            }
1467        };
1468
1469        let rows = outcome.records.len() as u64;
1470        let reserved_bytes = outcome
1471            .records
1472            .iter()
1473            .map(|record| record.reserved_bytes)
1474            .sum::<usize>();
1475        let mut failure_event = None;
1476        {
1477            let mut shared = lock(&self.shared);
1478            shared.in_flight_batches = shared.in_flight_batches.saturating_sub(1);
1479            shared.pending_rows = shared.pending_rows.saturating_sub(rows);
1480            shared.pending_bytes = shared.pending_bytes.saturating_sub(reserved_bytes);
1481            shared.retries = shared.retries.saturating_add(outcome.retries as u64);
1482            self.interval.retries = self.interval.retries.saturating_add(outcome.retries as u64);
1483
1484            match &outcome.result {
1485                Ok(result) => {
1486                    let committed = result.num_rows_inserted.max(0) as u64;
1487                    shared.committed_rows = shared.committed_rows.saturating_add(committed);
1488                    self.interval.committed_rows =
1489                        self.interval.committed_rows.saturating_add(committed);
1490                    self.interval.committed_batches =
1491                        self.interval.committed_batches.saturating_add(1);
1492                    record_circuit_success(&mut shared, outcome.circuit_probe);
1493                }
1494                Err(error) => {
1495                    let snapshot = error.clone();
1496                    let status = snapshot
1497                        .append_details()
1498                        .map(|d| d.status)
1499                        .unwrap_or(AppendState::Unknown);
1500                    if status == AppendState::Unknown {
1501                        shared.unknown_rows = shared.unknown_rows.saturating_add(rows);
1502                        self.interval.unknown_rows =
1503                            self.interval.unknown_rows.saturating_add(rows);
1504                        self.interval.unknown_batches =
1505                            self.interval.unknown_batches.saturating_add(1);
1506                    } else {
1507                        shared.failed_rows = shared.failed_rows.saturating_add(rows);
1508                        self.interval.failed_rows = self.interval.failed_rows.saturating_add(rows);
1509                        self.interval.failed_batches =
1510                            self.interval.failed_batches.saturating_add(1);
1511                    }
1512                    shared.last_failure = Some(AppendLastFailure {
1513                        at: SystemTime::now(),
1514                        error: error.clone(),
1515                    });
1516                    let circuit_opened = self.config.failure_policy
1517                        == AppendFailurePolicy::Continue
1518                        && record_circuit_failure(
1519                            &mut shared,
1520                            self.config.circuit_breaker,
1521                            error,
1522                            outcome.circuit_probe,
1523                        );
1524                    let action = if self.config.failure_policy == AppendFailurePolicy::Stop {
1525                        AppendBatchFailureAction::Stopped
1526                    } else if circuit_opened {
1527                        AppendBatchFailureAction::CircuitOpened
1528                    } else {
1529                        AppendBatchFailureAction::Continuing
1530                    };
1531                    failure_event = Some(AppendBatchFailureEvent {
1532                        error: snapshot.clone(),
1533                        batch_rows: rows as usize,
1534                        batch_bytes: outcome.batch_bytes,
1535                        outcome: status,
1536                        action,
1537                    });
1538                    if self.config.failure_policy == AppendFailurePolicy::Stop {
1539                        drop(shared);
1540                        self.set_fatal(snapshot);
1541                    }
1542                }
1543            }
1544        }
1545        drop(outcome.records);
1546
1547        if let Some(event) = &failure_event {
1548            for listener in &self.config.batch_failure_listeners {
1549                let _ = std::panic::catch_unwind(AssertUnwindSafe(|| listener(event)));
1550            }
1551        }
1552    }
1553
1554    fn set_fatal(&mut self, snapshot: Error) {
1555        let replace = match &self.fatal {
1556            None => true,
1557            Some(current) => {
1558                current.append_details().map(|d| d.status) != Some(AppendState::Unknown)
1559                    && snapshot.append_details().map(|d| d.status) == Some(AppendState::Unknown)
1560            }
1561        };
1562        if replace {
1563            self.fatal = Some(snapshot.clone());
1564        }
1565        let fatal = self.fatal.clone().expect("fatal state must be set");
1566        let mut shared = lock(&self.shared);
1567        shared.lifecycle = AppendStreamState::Failed;
1568        shared.fatal = Some(fatal);
1569        drop(shared);
1570        self.pending_bytes.close();
1571    }
1572
1573    fn mark_records_failed(&mut self, records: Vec<BufferedRecord>) {
1574        if records.is_empty() {
1575            return;
1576        }
1577        let rows = records.len() as u64;
1578        let bytes = records
1579            .iter()
1580            .map(|record| record.reserved_bytes)
1581            .sum::<usize>();
1582        let mut shared = lock(&self.shared);
1583        shared.failed_rows = shared.failed_rows.saturating_add(rows);
1584        shared.pending_rows = shared.pending_rows.saturating_sub(rows);
1585        shared.pending_bytes = shared.pending_bytes.saturating_sub(bytes);
1586        self.interval.failed_rows = self.interval.failed_rows.saturating_add(rows);
1587        drop(shared);
1588        drop(records);
1589    }
1590
1591    async fn finish_fatal(
1592        &mut self,
1593        rx: &mut mpsc::Receiver<AppendCommand>,
1594        barrier: Option<BarrierCommand>,
1595    ) {
1596        rx.close();
1597        let buffered = std::mem::take(&mut self.rows);
1598        self.mark_records_failed(buffered);
1599        self.current_bytes = 0;
1600        self.batch_deadline = None;
1601
1602        let mut barriers: Vec<_> = barrier.into_iter().collect();
1603        while let Some(command) = rx.recv().await {
1604            match command {
1605                AppendCommand::Record(record) => {
1606                    self.interval.accepted_rows = self.interval.accepted_rows.saturating_add(1);
1607                    self.mark_records_failed(vec![record]);
1608                }
1609                AppendCommand::Flush(barrier) | AppendCommand::Shutdown(barrier) => {
1610                    barriers.push(barrier);
1611                }
1612            }
1613        }
1614        // All outstanding barriers observe this settlement; retain the earliest wait.
1615        let started_at = barriers
1616            .iter()
1617            .map(|barrier| barrier.started_at)
1618            .min()
1619            .unwrap_or_else(Instant::now);
1620        self.drain_in_flight().await;
1621
1622        let fatal = self
1623            .fatal
1624            .clone()
1625            .unwrap_or_else(|| Error::new(ErrorKind::Unexpected, "append stream failed"));
1626        {
1627            let mut shared = lock(&self.shared);
1628            let resolved = shared
1629                .committed_rows
1630                .saturating_add(shared.failed_rows)
1631                .saturating_add(shared.unknown_rows);
1632            let unresolved = shared.accepted_rows.saturating_sub(resolved);
1633            if unresolved > 0 {
1634                shared.unknown_rows = shared.unknown_rows.saturating_add(unresolved);
1635                self.interval.unknown_rows = self.interval.unknown_rows.saturating_add(unresolved);
1636                self.interval.unknown_batches = self.interval.unknown_batches.saturating_add(1);
1637            }
1638            shared.pending_rows = 0;
1639            shared.pending_bytes = 0;
1640            shared.in_flight_batches = 0;
1641        }
1642        let dropped = lock(&self.shared).dropped_rows;
1643        self.take_report(dropped, started_at);
1644        {
1645            let mut shared = lock(&self.shared);
1646            shared.lifecycle = AppendStreamState::Failed;
1647            shared.fatal = Some(fatal.clone());
1648            shared.final_result = Some(Err(fatal.clone()));
1649        }
1650        for barrier in barriers {
1651            let _ = barrier.ack.send(Err(fatal.clone()));
1652        }
1653        self.pending_bytes.close();
1654        self.shutdown_notify.notify_waiters();
1655    }
1656
1657    fn take_report(
1658        &mut self,
1659        dropped_rows_at_barrier: u64,
1660        started_at: Instant,
1661    ) -> AppendDeliveryReport {
1662        let dropped_rows = dropped_rows_at_barrier.saturating_sub(self.reported_dropped_rows);
1663        self.reported_dropped_rows = self.reported_dropped_rows.max(dropped_rows_at_barrier);
1664        let lost_rows = self
1665            .interval
1666            .failed_rows
1667            .saturating_add(self.interval.unknown_rows)
1668            .saturating_add(dropped_rows);
1669        let outcome = if lost_rows == 0 {
1670            AppendDeliveryOutcome::Ok
1671        } else if self.interval.committed_rows > 0 {
1672            AppendDeliveryOutcome::Partial
1673        } else if self.interval.unknown_rows > 0 {
1674            AppendDeliveryOutcome::Unknown
1675        } else {
1676            AppendDeliveryOutcome::Failed
1677        };
1678        let report = AppendDeliveryReport {
1679            outcome,
1680            accepted_rows: self.interval.accepted_rows,
1681            committed_rows: self.interval.committed_rows,
1682            failed_rows: self.interval.failed_rows,
1683            unknown_rows: self.interval.unknown_rows,
1684            dropped_rows,
1685            committed_batches: self.interval.committed_batches,
1686            failed_batches: self.interval.failed_batches,
1687            unknown_batches: self.interval.unknown_batches,
1688            retries: self.interval.retries,
1689            duration: started_at.elapsed(),
1690        };
1691        self.interval = DeliveryCounters::default();
1692        let mut shared = lock(&self.shared);
1693        shared.last_report = Some(report.clone());
1694        report
1695    }
1696
1697    fn finish_closed(&mut self) {
1698        let report = lock(&self.shared)
1699            .last_report
1700            .clone()
1701            .unwrap_or_else(empty_report);
1702        self.finish_closed_with_report(report);
1703    }
1704
1705    fn finish_closed_with_report(&mut self, report: AppendDeliveryReport) {
1706        let mut shared = lock(&self.shared);
1707        shared.lifecycle = AppendStreamState::Closed;
1708        shared.last_report = Some(report.clone());
1709        shared.final_result = Some(Ok(report));
1710        drop(shared);
1711        self.pending_bytes.close();
1712        self.shutdown_notify.notify_waiters();
1713    }
1714}
1715
1716async fn append_batch(
1717    request: AppendBatchRequest,
1718    payload: String,
1719) -> (Result<AppendRowsResult, Error>, usize) {
1720    let mut retries = 0usize;
1721    let mut backoff = request.retry.initial_backoff.min(request.retry.max_backoff);
1722    let mut saw_unknown = false;
1723    loop {
1724        let result = request
1725            .client
1726            .append_rows_with_timeout(
1727                &request.database,
1728                &request.schema,
1729                &request.table,
1730                payload.clone(),
1731                request.attempt_timeout,
1732            )
1733            .await;
1734
1735        let error = match result {
1736            Ok(result) => return (Ok(result), retries),
1737            Err(error) => error,
1738        };
1739        saw_unknown |= error
1740            .append_details()
1741            .is_some_and(|details| details.status == AppendState::Unknown);
1742        let retryable = append_retryable(&error, request.retry.retry_unknown);
1743        if !retryable || retries >= request.retry.max_retries {
1744            let error = if retryable {
1745                error.with_context("retries", retries).set_persistent()
1746            } else {
1747                error
1748            };
1749            return (Err(preserve_unknown_outcome(error, saw_unknown)), retries);
1750        }
1751        let retry_delay = retry_delay(backoff, &error, request.retry.max_backoff);
1752        if !retry_delay.is_zero() {
1753            tokio::time::sleep(retry_delay).await;
1754        }
1755        retries += 1;
1756        backoff = next_backoff(backoff, request.retry.max_backoff);
1757    }
1758}
1759
1760fn append_retryable(error: &Error, retry_unknown: bool) -> bool {
1761    if error.write_details().is_some() || error.server_retryable() == Some(false) {
1762        return false;
1763    }
1764    match error.append_details().map(|details| details.status) {
1765        Some(AppendState::Rejected) => error.is_temporary(),
1766        Some(AppendState::Unknown) if retry_unknown => error.http_status().is_none_or(|status| {
1767            status.is_success()
1768                || status == reqwest::StatusCode::REQUEST_TIMEOUT
1769                || status == reqwest::StatusCode::TOO_MANY_REQUESTS
1770                || status.is_server_error()
1771        }),
1772        _ => false,
1773    }
1774}
1775
1776fn preserve_unknown_outcome(error: Error, saw_unknown: bool) -> Error {
1777    if !saw_unknown
1778        || error
1779            .append_details()
1780            .is_some_and(|details| details.status == AppendState::Unknown)
1781    {
1782        return error;
1783    }
1784    // A later rejection cannot establish whether an earlier attempt committed.
1785    let mut details = error
1786        .append_details()
1787        .cloned()
1788        .unwrap_or(AppendErrorDetails {
1789            status: AppendState::Unknown,
1790            row_errors: Vec::new(),
1791            row_errors_truncated: false,
1792        });
1793    details.status = AppendState::Unknown;
1794    error
1795        .with_context("delivery", "an earlier append attempt may have committed")
1796        .set_append_details(details)
1797        .set_persistent()
1798}
1799
1800fn unknown_append_error<E>(message: impl Into<String>, source: Option<E>) -> Error
1801where
1802    E: Into<anyhow::Error>,
1803{
1804    let mut error =
1805        Error::new(ErrorKind::AppendRowsFailed, message).set_append_details(AppendErrorDetails {
1806            status: AppendState::Unknown,
1807            row_errors: Vec::new(),
1808            row_errors_truncated: false,
1809        });
1810    if let Some(source) = source {
1811        error = error.set_source(source);
1812    }
1813    error.set_persistent()
1814}
1815
1816fn next_backoff(current: Duration, max_backoff: Duration) -> Duration {
1817    if current.is_zero() {
1818        return Duration::ZERO;
1819    }
1820    current
1821        .checked_mul(2)
1822        .unwrap_or(max_backoff)
1823        .min(max_backoff)
1824}
1825
1826fn retry_delay(backoff: Duration, error: &Error, max_backoff: Duration) -> Duration {
1827    backoff
1828        .max(error.retry_after().unwrap_or(Duration::ZERO))
1829        .min(max_backoff)
1830}
1831
1832fn try_claim_circuit_admission(shared: &mut SharedState) -> bool {
1833    match shared.circuit.state {
1834        AppendCircuitState::Closed => true,
1835        AppendCircuitState::HalfOpen => false,
1836        AppendCircuitState::Open => {
1837            let cooled_down = shared
1838                .circuit
1839                .opened_until
1840                .is_none_or(|until| Instant::now() >= until);
1841            if cooled_down && !shared.circuit.probe_admission_claimed {
1842                shared.circuit.probe_admission_claimed = true;
1843                true
1844            } else {
1845                false
1846            }
1847        }
1848    }
1849}
1850
1851fn circuit_admission_available(shared: &SharedState) -> bool {
1852    match shared.circuit.state {
1853        AppendCircuitState::Closed => true,
1854        AppendCircuitState::HalfOpen => false,
1855        AppendCircuitState::Open => {
1856            shared
1857                .circuit
1858                .opened_until
1859                .is_none_or(|until| Instant::now() >= until)
1860                && !shared.circuit.probe_admission_claimed
1861        }
1862    }
1863}
1864
1865fn record_circuit_success(shared: &mut SharedState, probe: bool) {
1866    shared.circuit.consecutive_failures = 0;
1867    if probe {
1868        shared.circuit.state = AppendCircuitState::Closed;
1869        shared.circuit.opened_until = None;
1870        shared.circuit.probe_admission_claimed = false;
1871    }
1872}
1873
1874fn record_circuit_failure(
1875    shared: &mut SharedState,
1876    options: Option<AppendCircuitBreakerOptions>,
1877    error: &Error,
1878    probe: bool,
1879) -> bool {
1880    let Some(options) = options else {
1881        return false;
1882    };
1883    let availability_failure = error.is_temporary()
1884        || error.is_persistent()
1885        || error
1886            .append_details()
1887            .is_some_and(|details| details.status == AppendState::Unknown);
1888    if !availability_failure {
1889        if probe {
1890            record_circuit_success(shared, true);
1891        }
1892        return false;
1893    }
1894
1895    shared.circuit.consecutive_failures = shared.circuit.consecutive_failures.saturating_add(1);
1896    if probe || shared.circuit.consecutive_failures >= options.failure_threshold {
1897        shared.circuit.state = AppendCircuitState::Open;
1898        shared.circuit.opened_until = Some(deadline_after(options.cooldown));
1899        shared.circuit.probe_admission_claimed = false;
1900        true
1901    } else {
1902        false
1903    }
1904}
1905
1906#[derive(Debug, Clone, Copy)]
1907enum DropReason {
1908    BufferFull,
1909    InvalidRecord,
1910    RecordTooLarge,
1911    CircuitOpen,
1912    Closed,
1913}
1914
1915fn note_drop(shared: &mut SharedState, reason: DropReason) {
1916    shared.dropped_rows = shared.dropped_rows.saturating_add(1);
1917    let counter = match reason {
1918        DropReason::BufferFull => &mut shared.dropped_by_reason.buffer_full,
1919        DropReason::InvalidRecord => &mut shared.dropped_by_reason.invalid_record,
1920        DropReason::RecordTooLarge => &mut shared.dropped_by_reason.record_too_large,
1921        DropReason::CircuitOpen => &mut shared.dropped_by_reason.circuit_open,
1922        DropReason::Closed => &mut shared.dropped_by_reason.closed,
1923    };
1924    *counter = counter.saturating_add(1);
1925}
1926
1927fn empty_report() -> AppendDeliveryReport {
1928    AppendDeliveryReport {
1929        outcome: AppendDeliveryOutcome::Ok,
1930        accepted_rows: 0,
1931        committed_rows: 0,
1932        failed_rows: 0,
1933        unknown_rows: 0,
1934        dropped_rows: 0,
1935        committed_batches: 0,
1936        failed_batches: 0,
1937        unknown_batches: 0,
1938        retries: 0,
1939        duration: Duration::ZERO,
1940    }
1941}
1942
1943fn deadline_after(duration: Duration) -> Instant {
1944    let now = Instant::now();
1945    now.checked_add(duration).unwrap_or(now)
1946}
1947
1948fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
1949    mutex
1950        .lock()
1951        .unwrap_or_else(|poisoned| poisoned.into_inner())
1952}
1953
1954#[cfg(test)]
1955mod tests {
1956    use std::collections::HashMap;
1957    use std::io::Read;
1958    use std::io::Write;
1959    use std::net::TcpListener;
1960    use std::net::TcpStream;
1961    use std::sync::atomic::AtomicBool;
1962    use std::sync::atomic::AtomicUsize;
1963    use std::sync::atomic::Ordering;
1964    use std::thread;
1965
1966    use super::*;
1967
1968    #[derive(Debug, Clone)]
1969    struct RecordedRequest {
1970        target: String,
1971        headers: HashMap<String, String>,
1972        body: String,
1973    }
1974
1975    #[derive(Debug, Clone)]
1976    struct MockResponse {
1977        status: u16,
1978        body: Vec<u8>,
1979        delay: Duration,
1980        headers: Vec<(&'static str, String)>,
1981        body_delay: Duration,
1982        disconnect: Option<MockDisconnect>,
1983    }
1984
1985    #[derive(Debug, Clone, Copy)]
1986    enum MockDisconnect {
1987        BeforeHeaders,
1988        DuringBody,
1989    }
1990
1991    impl MockResponse {
1992        fn json(status: u16, body: impl Into<String>) -> Self {
1993            Self {
1994                status,
1995                body: body.into().into_bytes(),
1996                delay: Duration::ZERO,
1997                headers: Vec::new(),
1998                body_delay: Duration::ZERO,
1999                disconnect: None,
2000            }
2001        }
2002
2003        fn committed(request: &RecordedRequest) -> Self {
2004            let rows = request
2005                .body
2006                .lines()
2007                .filter(|line| !line.trim().is_empty())
2008                .count();
2009            Self::json(
2010                200,
2011                format!(r#"{{"status":"committed","num_rows_inserted":{rows}}}"#),
2012            )
2013        }
2014
2015        fn with_delay(mut self, delay: Duration) -> Self {
2016            self.delay = delay;
2017            self
2018        }
2019
2020        fn with_header(mut self, name: &'static str, value: &str) -> Self {
2021            self.headers.push((name, value.to_string()));
2022            self
2023        }
2024
2025        fn without_ack(mut self) -> Self {
2026            self.disconnect = Some(MockDisconnect::BeforeHeaders);
2027            self
2028        }
2029
2030        fn with_body_delay(mut self, delay: Duration) -> Self {
2031            self.body_delay = delay;
2032            self
2033        }
2034
2035        fn with_disconnect(mut self, disconnect: MockDisconnect) -> Self {
2036            self.disconnect = Some(disconnect);
2037            self
2038        }
2039    }
2040
2041    type Responder = dyn Fn(usize, &RecordedRequest) -> MockResponse + Send + Sync;
2042
2043    struct MockServer {
2044        endpoint: String,
2045        requests: Arc<Mutex<Vec<RecordedRequest>>>,
2046        max_active: Arc<AtomicUsize>,
2047        stop: Arc<AtomicBool>,
2048        accept_task: Option<thread::JoinHandle<()>>,
2049    }
2050
2051    impl MockServer {
2052        fn start<F>(responder: F) -> Self
2053        where
2054            F: Fn(usize, &RecordedRequest) -> MockResponse + Send + Sync + 'static,
2055        {
2056            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
2057            listener.set_nonblocking(true).unwrap();
2058            let address = listener.local_addr().unwrap();
2059            let endpoint = format!("http://{address}");
2060            let requests = Arc::new(Mutex::new(Vec::new()));
2061            let active = Arc::new(AtomicUsize::new(0));
2062            let max_active = Arc::new(AtomicUsize::new(0));
2063            let stop = Arc::new(AtomicBool::new(false));
2064            let responder: Arc<Responder> = Arc::new(responder);
2065
2066            let accept_task = thread::spawn({
2067                let requests = requests.clone();
2068                let active = active.clone();
2069                let max_active = max_active.clone();
2070                let stop = stop.clone();
2071                move || {
2072                    while !stop.load(Ordering::Acquire) {
2073                        match listener.accept() {
2074                            Ok((stream, _)) => {
2075                                if stop.load(Ordering::Acquire) {
2076                                    break;
2077                                }
2078                                let requests = requests.clone();
2079                                let active = active.clone();
2080                                let max_active = max_active.clone();
2081                                let responder = responder.clone();
2082                                thread::spawn(move || {
2083                                    handle_connection(
2084                                        stream, requests, active, max_active, responder,
2085                                    );
2086                                });
2087                            }
2088                            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
2089                                thread::sleep(Duration::from_millis(1));
2090                            }
2091                            Err(_) => break,
2092                        }
2093                    }
2094                }
2095            });
2096
2097            Self {
2098                endpoint,
2099                requests,
2100                max_active,
2101                stop,
2102                accept_task: Some(accept_task),
2103            }
2104        }
2105
2106        fn client(&self) -> Client {
2107            let http_client = reqwest::Client::builder().no_proxy().build().unwrap();
2108            Client::new(&self.endpoint, http_client).unwrap()
2109        }
2110
2111        fn client_with_write_timeout(&self, timeout: Duration) -> Client {
2112            Client::builder(&self.endpoint)
2113                .http_client(reqwest::Client::builder().no_proxy().build().unwrap())
2114                .write_timeout(timeout)
2115                .build()
2116                .unwrap()
2117        }
2118
2119        fn requests(&self) -> Vec<RecordedRequest> {
2120            lock(&self.requests).clone()
2121        }
2122
2123        fn max_active(&self) -> usize {
2124            self.max_active.load(Ordering::Acquire)
2125        }
2126    }
2127
2128    impl Drop for MockServer {
2129        fn drop(&mut self) {
2130            self.stop.store(true, Ordering::Release);
2131            let _ = TcpStream::connect(self.endpoint.trim_start_matches("http://"));
2132            if let Some(task) = self.accept_task.take() {
2133                let _ = task.join();
2134            }
2135        }
2136    }
2137
2138    fn handle_connection(
2139        mut stream: TcpStream,
2140        requests: Arc<Mutex<Vec<RecordedRequest>>>,
2141        active: Arc<AtomicUsize>,
2142        max_active: Arc<AtomicUsize>,
2143        responder: Arc<Responder>,
2144    ) {
2145        stream.set_nonblocking(false).unwrap();
2146        stream
2147            .set_read_timeout(Some(Duration::from_secs(2)))
2148            .unwrap();
2149        let Some(request) = read_request(&mut stream) else {
2150            return;
2151        };
2152        let index = {
2153            let mut recorded = lock(&requests);
2154            let index = recorded.len();
2155            recorded.push(request.clone());
2156            index
2157        };
2158        let now_active = active.fetch_add(1, Ordering::AcqRel) + 1;
2159        max_active.fetch_max(now_active, Ordering::AcqRel);
2160        let response = responder(index, &request);
2161        if !response.delay.is_zero() {
2162            thread::sleep(response.delay);
2163        }
2164        if matches!(response.disconnect, Some(MockDisconnect::BeforeHeaders)) {
2165            active.fetch_sub(1, Ordering::AcqRel);
2166            return;
2167        }
2168        let reason = match response.status {
2169            200 => "OK",
2170            400 => "Bad Request",
2171            422 => "Unprocessable Entity",
2172            503 => "Service Unavailable",
2173            _ => "Test Response",
2174        };
2175        let headers = response
2176            .headers
2177            .iter()
2178            .map(|(name, value)| format!("{name}: {value}\r\n"))
2179            .collect::<String>();
2180        let wire = format!(
2181            "HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n",
2182            response.status,
2183            reason,
2184            response.body.len(),
2185            headers,
2186        );
2187        let _ = stream.write_all(wire.as_bytes());
2188        let _ = stream.flush();
2189        if !response.body_delay.is_zero() {
2190            thread::sleep(response.body_delay);
2191        }
2192        let body = response.body.as_slice();
2193        let body = if matches!(response.disconnect, Some(MockDisconnect::DuringBody)) {
2194            &body[..body.len().min(1)]
2195        } else {
2196            body
2197        };
2198        let _ = stream.write_all(body);
2199        let _ = stream.flush();
2200        active.fetch_sub(1, Ordering::AcqRel);
2201    }
2202
2203    fn read_request(stream: &mut TcpStream) -> Option<RecordedRequest> {
2204        let mut bytes = Vec::new();
2205        let mut buffer = [0u8; 4096];
2206        let header_end = loop {
2207            let read = stream.read(&mut buffer).ok()?;
2208            if read == 0 {
2209                return None;
2210            }
2211            bytes.extend_from_slice(&buffer[..read]);
2212            if let Some(position) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
2213                break position + 4;
2214            }
2215        };
2216        let headers_text = String::from_utf8(bytes[..header_end].to_vec()).ok()?;
2217        let mut lines = headers_text.split("\r\n");
2218        let target = lines.next()?.split_whitespace().nth(1)?.to_string();
2219        let mut headers = HashMap::new();
2220        for line in lines.filter(|line| !line.is_empty()) {
2221            let (name, value) = line.split_once(':')?;
2222            headers.insert(name.trim().to_ascii_lowercase(), value.trim().to_string());
2223        }
2224        let content_length = headers
2225            .get("content-length")
2226            .and_then(|value| value.parse::<usize>().ok())
2227            .unwrap_or(0);
2228        while bytes.len().saturating_sub(header_end) < content_length {
2229            let read = stream.read(&mut buffer).ok()?;
2230            if read == 0 {
2231                return None;
2232            }
2233            bytes.extend_from_slice(&buffer[..read]);
2234        }
2235        let wire_body = &bytes[header_end..header_end.saturating_add(content_length)];
2236        let decoded_body = match headers.get("content-encoding").map(String::as_str) {
2237            Some("zstd") => zstd::stream::decode_all(wire_body).ok()?,
2238            None | Some("identity") => wire_body.to_vec(),
2239            Some(_) => return None,
2240        };
2241        let body = String::from_utf8(decoded_body).ok()?;
2242        Some(RecordedRequest {
2243            target,
2244            headers,
2245            body,
2246        })
2247    }
2248
2249    #[test]
2250    fn mock_server_waits_for_delayed_request_bytes() {
2251        let server = MockServer::start(|_, request| MockResponse::committed(request));
2252        let mut stream = TcpStream::connect(server.endpoint.trim_start_matches("http://")).unwrap();
2253        stream
2254            .set_read_timeout(Some(Duration::from_secs(1)))
2255            .unwrap();
2256
2257        thread::sleep(Duration::from_millis(20));
2258        let body = "{\"id\":1}\n";
2259        let request = format!(
2260            "POST /rows HTTP/1.1\r\nHost: localhost\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
2261            body.len()
2262        );
2263        stream.write_all(request.as_bytes()).unwrap();
2264
2265        let mut response = String::new();
2266        stream.read_to_string(&mut response).unwrap();
2267        assert!(response.starts_with("HTTP/1.1 200 OK"), "{response}");
2268        assert_eq!(server.requests().len(), 1);
2269    }
2270
2271    #[test]
2272    fn rejects_non_object_rows() {
2273        assert!(matches!(
2274            prepare_record_for_try_send(&vec![1, 2, 3]),
2275            Err(AppendTrySendError::InvalidRecord)
2276        ));
2277    }
2278
2279    #[test]
2280    fn report_outcomes_include_drops_and_unknown_rows() {
2281        let config = AppendStreamConfig {
2282            client: Client::new("http://127.0.0.1:1", reqwest::Client::new()).unwrap(),
2283            database: "scopedb".to_string(),
2284            schema: "public".to_string(),
2285            table: "events".to_string(),
2286            batch_bytes: 1024,
2287            max_batch_rows: 100,
2288            flush_interval: Duration::from_secs(1),
2289            channel_capacity: 1,
2290            max_pending_bytes: 1024,
2291            max_in_flight_requests: 1,
2292            retry: RetryConfig {
2293                max_retries: 0,
2294                retry_unknown: false,
2295                initial_backoff: Duration::ZERO,
2296                max_backoff: Duration::ZERO,
2297            },
2298            failure_policy: AppendFailurePolicy::Continue,
2299            attempt_timeout: Some(Duration::from_secs(1)),
2300            circuit_breaker: None,
2301            batch_failure_listeners: Vec::new(),
2302        };
2303        let shared = Arc::new(Mutex::new(SharedState::default()));
2304        let budget = Arc::new(PendingBytesBudget::new(1024));
2305        let mut worker = AppendWorker {
2306            config,
2307            shared,
2308            pending_bytes: budget,
2309            shutdown_notify: Arc::new(Notify::new()),
2310            rows: Vec::new(),
2311            current_bytes: 0,
2312            batch_deadline: None,
2313            in_flight: JoinSet::new(),
2314            interval: DeliveryCounters {
2315                accepted_rows: 2,
2316                committed_rows: 1,
2317                unknown_rows: 1,
2318                ..DeliveryCounters::default()
2319            },
2320            reported_dropped_rows: 0,
2321            fatal: None,
2322        };
2323        let report = worker.take_report(1, Instant::now());
2324        assert_eq!(report.outcome, AppendDeliveryOutcome::Partial);
2325        assert_eq!(report.dropped_rows, 1);
2326    }
2327
2328    #[test]
2329    fn backoff_is_capped() {
2330        assert_eq!(
2331            next_backoff(Duration::from_secs(4), Duration::from_secs(5)),
2332            Duration::from_secs(5)
2333        );
2334    }
2335
2336    #[test]
2337    fn retry_after_is_honored_and_capped() {
2338        let error = Error::new(ErrorKind::AppendRowsFailed, "busy")
2339            .set_retry_after(Duration::from_secs(30));
2340        assert_eq!(
2341            retry_delay(Duration::from_millis(100), &error, Duration::from_secs(5)),
2342            Duration::from_secs(5)
2343        );
2344
2345        let error = Error::new(ErrorKind::AppendRowsFailed, "busy")
2346            .set_retry_after(Duration::from_millis(200));
2347        assert_eq!(
2348            retry_delay(Duration::from_secs(1), &error, Duration::from_secs(5)),
2349            Duration::from_secs(1)
2350        );
2351    }
2352
2353    #[tokio::test]
2354    async fn invalid_stream_limits_fail_before_spawning_a_worker() {
2355        let client = Client::builder("https://example.invalid").build().unwrap();
2356        for builder in [
2357            client
2358                .table("events")
2359                .append_stream()
2360                .channel_capacity(usize::MAX),
2361            client
2362                .table("events")
2363                .append_stream()
2364                .initial_backoff(Duration::MAX),
2365            client
2366                .table("events")
2367                .append_stream()
2368                .max_backoff(Duration::MAX),
2369            client.table("..").append_stream(),
2370        ] {
2371            assert_eq!(
2372                builder
2373                    .build()
2374                    .err()
2375                    .expect("expected invalid configuration")
2376                    .kind(),
2377                ErrorKind::ConfigInvalid
2378            );
2379        }
2380    }
2381
2382    #[test]
2383    fn max_batch_rows_cannot_exceed_protocol_limit() {
2384        let result = Client::new("https://example.com", reqwest::Client::new())
2385            .unwrap()
2386            .table("events")
2387            .append_stream()
2388            .max_batch_rows(MAX_APPEND_ROWS + 1)
2389            .build();
2390        let error = match result {
2391            Ok(_) => panic!("invalid max_batch_rows was accepted"),
2392            Err(error) => error,
2393        };
2394        assert_eq!(error.kind(), ErrorKind::ConfigInvalid);
2395    }
2396
2397    #[test]
2398    fn append_stream_enforces_eight_mibibyte_request_limit() {
2399        let result = Client::new("https://example.com", reqwest::Client::new())
2400            .unwrap()
2401            .table("events")
2402            .append_stream()
2403            .target_batch_bytes(MAX_APPEND_BODY_BYTES + 1)
2404            .build();
2405        let error = match result {
2406            Ok(_) => panic!("oversized target_batch_bytes was accepted"),
2407            Err(error) => error,
2408        };
2409        assert_eq!(error.kind(), ErrorKind::ConfigInvalid);
2410
2411        let row = serde_json::json!({"payload": "x".repeat(MAX_APPEND_BODY_BYTES)});
2412        assert!(matches!(
2413            prepare_record_for_try_send(&row),
2414            Err(AppendTrySendError::RecordTooLarge {
2415                limit: MAX_APPEND_BODY_BYTES,
2416                ..
2417            })
2418        ));
2419    }
2420
2421    #[test]
2422    fn error_snapshot_preserves_http_metadata() {
2423        let error = Error::new(ErrorKind::AppendRowsFailed, "busy")
2424            .set_http_status(reqwest::StatusCode::SERVICE_UNAVAILABLE)
2425            .set_request_id("request-123".to_string())
2426            .set_retry_after(Duration::from_secs(2))
2427            .with_context("retries", 3)
2428            .set_source(std::io::Error::new(
2429                std::io::ErrorKind::ConnectionReset,
2430                "lost response",
2431            ))
2432            .set_temporary();
2433        let restored = error.clone();
2434        assert_eq!(
2435            restored.http_status(),
2436            Some(reqwest::StatusCode::SERVICE_UNAVAILABLE)
2437        );
2438        assert_eq!(restored.request_id(), Some("request-123"));
2439        assert_eq!(restored.retry_after(), Some(Duration::from_secs(2)));
2440        assert!(restored.is_temporary());
2441        let cause = StdError::source(&restored)
2442            .unwrap()
2443            .downcast_ref::<std::io::Error>()
2444            .unwrap();
2445        assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
2446        assert!(format!("{restored:?}").contains("retries: 3"));
2447    }
2448
2449    #[tokio::test]
2450    async fn owned_http_client_does_not_follow_redirected_writes() {
2451        for status in [302, 307, 308] {
2452            let server = MockServer::start(move |index, request| {
2453                if index == 0 {
2454                    MockResponse::json(status, "redirect").with_header("Location", "/replayed")
2455                } else {
2456                    MockResponse::committed(request)
2457                }
2458            });
2459            let client = Client::builder(&server.endpoint).build().unwrap();
2460            let error = client
2461                .table("events")
2462                .append(r#"{"id":1}"#)
2463                .await
2464                .unwrap_err();
2465            assert_eq!(error.http_status().unwrap().as_u16(), status);
2466            assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
2467            assert_eq!(server.requests().len(), 1);
2468        }
2469    }
2470
2471    #[tokio::test]
2472    async fn cancelled_admission_and_shutdown_wait_release_capacity_and_settle_rows() {
2473        use std::future::Future;
2474        use std::task::Poll;
2475        let (arrived_tx, arrived_rx) = oneshot::channel();
2476        let arrived_tx = Mutex::new(Some(arrived_tx));
2477        let (release_tx, release_rx) = std::sync::mpsc::channel();
2478        let release_rx = Mutex::new(release_rx);
2479        let server = MockServer::start(move |_, request| {
2480            if let Some(tx) = lock(&arrived_tx).take() {
2481                let _ = tx.send(());
2482            }
2483            lock(&release_rx)
2484                .recv_timeout(Duration::from_secs(3))
2485                .unwrap();
2486            MockResponse::committed(request)
2487        });
2488        let row = serde_json::json!({"id":1});
2489        let bytes = row.to_string().len() + 1;
2490        let stream = server
2491            .client()
2492            .table("events")
2493            .append_stream()
2494            .max_batch_rows(1)
2495            .max_buffered_bytes(bytes)
2496            .build()
2497            .unwrap();
2498        stream.send(&row).await.unwrap();
2499        arrived_rx.await.unwrap();
2500        {
2501            let mut blocked = Box::pin(stream.send(&row));
2502            std::future::poll_fn(|cx| {
2503                assert!(blocked.as_mut().poll(cx).is_pending());
2504                Poll::Ready(())
2505            })
2506            .await;
2507        }
2508        assert_eq!(stream.stats().accepted_rows, 1);
2509        assert_eq!(stream.stats().pending_bytes, bytes);
2510        let mut blocked = Box::pin(stream.send(&row));
2511        std::future::poll_fn(|cx| {
2512            assert!(blocked.as_mut().poll(cx).is_pending());
2513            Poll::Ready(())
2514        })
2515        .await;
2516        {
2517            let mut shutdown = Box::pin(stream.shutdown());
2518            std::future::poll_fn(|cx| {
2519                assert!(shutdown.as_mut().poll(cx).is_pending());
2520                Poll::Ready(())
2521            })
2522            .await;
2523        }
2524        assert_eq!(stream.stats().state, AppendStreamState::Closing);
2525        assert!(stream.send(&row).await.is_err());
2526        release_tx.send(()).unwrap();
2527        let report = stream.shutdown().await.unwrap();
2528        assert!(blocked.await.is_err());
2529        assert_eq!(report.accepted_rows, 1);
2530        assert_eq!(report.committed_rows, 1);
2531        assert_eq!(stream.stats().pending_rows, 0);
2532        assert_eq!(stream.stats().pending_bytes, 0);
2533        assert_eq!(server.requests().len(), 1);
2534    }
2535
2536    #[tokio::test]
2537    async fn direct_append_rejects_excess_rows_before_network_io() {
2538        let server = MockServer::start(|_, request| MockResponse::committed(request));
2539        let error = server
2540            .client()
2541            .table("events")
2542            .append("{}\n".repeat(200_001))
2543            .await
2544            .unwrap_err();
2545        assert_eq!(
2546            error.append_details().unwrap().status,
2547            AppendState::Rejected
2548        );
2549        assert!(server.requests().is_empty());
2550    }
2551
2552    #[tokio::test]
2553    async fn direct_append_uses_default_request_compression() {
2554        let server = MockServer::start(|_, request| MockResponse::committed(request));
2555        let ndjson = "{\"id\":1}\n{\"id\":2}\n";
2556
2557        let result = server
2558            .client()
2559            .table("events")
2560            .append(ndjson)
2561            .await
2562            .unwrap();
2563
2564        assert_eq!(result.num_rows_inserted, 2);
2565        let requests = server.requests();
2566        assert_eq!(requests.len(), 1);
2567        assert_eq!(
2568            requests[0]
2569                .headers
2570                .get("content-encoding")
2571                .map(String::as_str),
2572            Some("zstd")
2573        );
2574        assert_eq!(requests[0].body, ndjson);
2575    }
2576
2577    #[tokio::test]
2578    async fn direct_append_rejects_a_body_over_eight_mibibytes_before_sending() {
2579        let server = MockServer::start(|_, request| MockResponse::committed(request));
2580
2581        let error = server
2582            .client()
2583            .table("events")
2584            .append("x".repeat(MAX_APPEND_BODY_BYTES + 1))
2585            .await
2586            .unwrap_err();
2587
2588        assert_eq!(error.kind(), ErrorKind::AppendRowsFailed);
2589        assert_eq!(
2590            error.append_details().map(|details| details.status),
2591            Some(AppendState::Rejected)
2592        );
2593        assert!(error.message().contains("8388608-byte append limit"));
2594        assert!(server.requests().is_empty());
2595    }
2596
2597    #[tokio::test]
2598    async fn batches_ndjson_and_runs_requests_concurrently() {
2599        let server = MockServer::start(|_, request| {
2600            MockResponse::committed(request).with_delay(Duration::from_millis(30))
2601        });
2602        let stream = server
2603            .client()
2604            .table("events/archive")
2605            .with_database("analytics 2026")
2606            .with_schema("public")
2607            .append_stream()
2608            .target_batch_bytes(1)
2609            .max_concurrent_batches(3)
2610            .build()
2611            .unwrap();
2612
2613        let admitted = stream
2614            .send_all((0..6).map(|id| serde_json::json!({"id": id})))
2615            .await
2616            .unwrap();
2617        let report = stream.shutdown().await.unwrap();
2618
2619        assert_eq!(admitted.accepted_rows, 6);
2620        assert_eq!(report.outcome, AppendDeliveryOutcome::Ok);
2621        assert_eq!(report.accepted_rows, 6);
2622        assert_eq!(report.committed_rows, 6);
2623        assert_eq!(report.committed_batches, 6);
2624        assert!(server.max_active() >= 2);
2625        let requests = server.requests();
2626        assert_eq!(requests.len(), 6);
2627        for request in requests {
2628            assert_eq!(
2629                request.target,
2630                "/v1/catalog/databases/analytics%202026/schemas/public/tables/events%2Farchive/ingest"
2631            );
2632            assert_eq!(
2633                request.headers.get("content-type").map(String::as_str),
2634                Some("application/x-ndjson")
2635            );
2636            assert_eq!(
2637                request.headers.get("content-encoding").map(String::as_str),
2638                Some("zstd")
2639            );
2640            assert_eq!(request.body.lines().count(), 1);
2641            assert!(request.body.starts_with('{'));
2642        }
2643    }
2644
2645    #[tokio::test]
2646    async fn default_and_maximum_batch_targets() {
2647        for target in [None, Some(8 * 1024 * 1024)] {
2648            let server = MockServer::start(|_, request| MockResponse::committed(request));
2649            let mut builder = server
2650                .client()
2651                .table("events")
2652                .append_stream()
2653                .flush_interval(Duration::from_secs(3600))
2654                .max_concurrent_batches(1);
2655            if let Some(target) = target {
2656                builder = builder.target_batch_bytes(target);
2657            }
2658            let stream = builder.build().unwrap();
2659            let row = serde_json::json!({"payload": "x".repeat(2 * 1024 * 1024)});
2660            stream.send(&row).await.unwrap();
2661            stream.send(&row).await.unwrap();
2662            let report = stream.shutdown().await.unwrap();
2663            assert_eq!(report.committed_rows, 2);
2664            let row_counts = server
2665                .requests()
2666                .iter()
2667                .map(|request| request.body.lines().count())
2668                .collect::<Vec<_>>();
2669            assert_eq!(
2670                row_counts,
2671                if target.is_none() {
2672                    vec![1, 1]
2673                } else {
2674                    vec![2]
2675                }
2676            );
2677        }
2678    }
2679
2680    #[tokio::test]
2681    async fn semantic_builder_names_split_batches_by_row_count() {
2682        let server = MockServer::start(|_, request| MockResponse::committed(request));
2683        let stream = server
2684            .client()
2685            .table("events")
2686            .append_stream()
2687            .target_batch_bytes(1024 * 1024)
2688            .max_batch_rows(2)
2689            .max_buffered_bytes(1024 * 1024)
2690            .max_concurrent_batches(1)
2691            .build()
2692            .unwrap();
2693
2694        stream
2695            .send_all((0..5).map(|id| serde_json::json!({"id": id})))
2696            .await
2697            .unwrap();
2698        let report = stream.shutdown().await.unwrap();
2699
2700        assert_eq!(report.committed_rows, 5);
2701        let row_counts = server
2702            .requests()
2703            .iter()
2704            .map(|request| request.body.lines().count())
2705            .collect::<Vec<_>>();
2706        assert_eq!(row_counts, vec![2, 2, 1]);
2707    }
2708
2709    #[tokio::test]
2710    async fn client_builder_attaches_api_key_to_requests() {
2711        let server = MockServer::start(|_, request| MockResponse::committed(request));
2712        let http_client = reqwest::Client::builder().no_proxy().build().unwrap();
2713        let client = Client::builder(&server.endpoint)
2714            .api_key("secret-api-key")
2715            .http_client(http_client)
2716            .build()
2717            .unwrap();
2718
2719        client.table("events").append(r#"{"id":1}"#).await.unwrap();
2720
2721        let requests = server.requests();
2722        assert_eq!(
2723            requests[0].headers.get("authorization").map(String::as_str),
2724            Some("Bearer secret-api-key")
2725        );
2726        assert_eq!(
2727            requests[0]
2728                .headers
2729                .get("content-encoding")
2730                .map(String::as_str),
2731            Some("zstd")
2732        );
2733    }
2734
2735    #[tokio::test]
2736    async fn lost_ack_requires_explicit_replay_and_counts_logical_rows() {
2737        for replay in [false, true] {
2738            for failure_policy in [AppendFailurePolicy::Stop, AppendFailurePolicy::Continue] {
2739                let commits = Arc::new(AtomicUsize::new(0));
2740                let server = MockServer::start({
2741                    let commits = commits.clone();
2742                    move |index, request| {
2743                        commits.fetch_add(1, Ordering::AcqRel);
2744                        let response = MockResponse::committed(request);
2745                        if index == 0 {
2746                            response.without_ack()
2747                        } else {
2748                            response
2749                        }
2750                    }
2751                });
2752                let mut builder = server
2753                    .client()
2754                    .table("events")
2755                    .append_stream()
2756                    .failure_policy(failure_policy)
2757                    .initial_backoff(Duration::ZERO);
2758                if replay {
2759                    builder = builder.retry_unknown(true);
2760                }
2761                let stream = builder.build().unwrap();
2762                stream.send(&serde_json::json!({"id": 1})).await.unwrap();
2763
2764                let result = stream.shutdown().await;
2765                if !replay {
2766                    if failure_policy == AppendFailurePolicy::Stop {
2767                        assert_eq!(
2768                            result.unwrap_err().append_details().unwrap().status,
2769                            AppendState::Unknown
2770                        );
2771                    } else {
2772                        assert_eq!(result.unwrap().outcome, AppendDeliveryOutcome::Unknown);
2773                    }
2774                    let stats = stream.stats();
2775                    assert_eq!(commits.load(Ordering::Acquire), 1);
2776                    assert_eq!(stats.unknown_rows, 1);
2777                    assert_eq!(stats.retries, 0);
2778                    assert_eq!(stats.last_report.unwrap().unknown_rows, 1);
2779                    continue;
2780                }
2781                let report = result.unwrap();
2782
2783                assert_eq!(commits.load(Ordering::Acquire), 2);
2784                let requests = server.requests();
2785                assert_eq!(requests.len(), 2);
2786                assert_eq!(requests[0].body, requests[1].body);
2787                assert_eq!(report.outcome, AppendDeliveryOutcome::Ok);
2788                assert_eq!(report.accepted_rows, 1);
2789                assert_eq!(report.committed_rows, 1);
2790                assert_eq!(report.committed_batches, 1);
2791                assert_eq!(report.unknown_rows, 0);
2792                assert_eq!(report.retries, 1);
2793            }
2794        }
2795    }
2796
2797    #[tokio::test]
2798    async fn unknown_replay_respects_explicit_server_retry_refusal() {
2799        let server = MockServer::start(|_, _| {
2800            MockResponse::json(
2801                503,
2802                r#"{"code":"future_failure","message":"do not replay","retryable":false,"details":{"status":"unknown"}}"#,
2803            )
2804        });
2805        let stream = server
2806            .client()
2807            .table("events")
2808            .append_stream()
2809            .retry_unknown(true)
2810            .initial_backoff(Duration::ZERO)
2811            .build()
2812            .unwrap();
2813        stream.send(&serde_json::json!({"id":1})).await.unwrap();
2814        let error = stream.shutdown().await.unwrap_err();
2815        assert_eq!(error.code(), Some("future_failure"));
2816        assert_eq!(stream.stats().unknown_rows, 1);
2817        assert_eq!(server.requests().len(), 1);
2818    }
2819
2820    #[tokio::test]
2821    async fn direct_append_does_not_retry_lost_ack() {
2822        let server = MockServer::start(|_, request| MockResponse::committed(request).without_ack());
2823
2824        let error = server
2825            .client()
2826            .table("events")
2827            .append(r#"{"id":1}"#)
2828            .await
2829            .unwrap_err();
2830
2831        assert_eq!(server.requests().len(), 1);
2832        assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
2833        assert!(!error.is_retryable());
2834        assert!(StdError::source(&error).unwrap().is::<reqwest::Error>());
2835    }
2836
2837    #[tokio::test]
2838    async fn zero_retry_budget_preserves_transport_cause_at_barrier() {
2839        let server = MockServer::start(|_, request| MockResponse::committed(request).without_ack());
2840        let stream = server
2841            .client()
2842            .table("events")
2843            .append_stream()
2844            .max_retries(0)
2845            .build()
2846            .unwrap();
2847        stream.send(&serde_json::json!({"id": 1})).await.unwrap();
2848
2849        let error = stream.shutdown().await.unwrap_err();
2850
2851        assert_eq!(server.requests().len(), 1);
2852        assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
2853        assert!(StdError::source(&error).unwrap().is::<reqwest::Error>());
2854        assert_eq!(stream.stats().retries, 0);
2855    }
2856
2857    #[tokio::test]
2858    async fn unknown_remains_unknown_after_terminal_rejection() {
2859        for final_status in [422, 503] {
2860            let server = MockServer::start(move |index, _| {
2861                match index {
2862                0 => MockResponse::json(503, r#"{"message":"ambiguous","code":"ingest_failure","details":{"status":"unknown"}}"#),
2863                1 => MockResponse::json(
2864                    final_status,
2865                    r#"{"message":"last rejection","request_id":"last-attempt","code":"ingest_failure","details":{"status":"rejected","row_errors":[{"row_index":0,"column":"id","message":"invalid"}],"row_errors_truncated":true}}"#,
2866                )
2867                .with_header("Retry-After", "2"),
2868                _ => panic!("unexpected retry request {index}"),
2869            }
2870            });
2871            let stream = server
2872                .client()
2873                .table("events")
2874                .append_stream()
2875                .retry_unknown(true)
2876                .max_retries(1)
2877                .initial_backoff(Duration::ZERO)
2878                .build()
2879                .unwrap();
2880            stream.send(&serde_json::json!({"id": 1})).await.unwrap();
2881
2882            let error = stream.shutdown().await.unwrap_err();
2883
2884            assert_eq!(server.requests().len(), 2);
2885            assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
2886            assert!(format!("{error:?}").contains("earlier append attempt may have committed"));
2887            assert_eq!(error.message(), "last rejection");
2888            assert_eq!(error.code(), Some("ingest_failure"));
2889            assert_eq!(error.details().unwrap()["status"], "rejected");
2890            assert_eq!(error.append_details().unwrap().row_errors[0].column, "id");
2891            assert!(error.append_details().unwrap().row_errors_truncated);
2892            assert_eq!(error.http_status().unwrap().as_u16(), final_status);
2893            assert_eq!(error.request_id(), Some("last-attempt"));
2894            assert_eq!(error.retry_after(), Some(Duration::from_secs(2)));
2895            if final_status == 503 {
2896                assert!(format!("{error:?}").contains("retries: 1"));
2897            }
2898            let stats = stream.stats();
2899            assert_eq!(stats.unknown_rows, 1);
2900            assert_eq!(stats.failed_rows, 0);
2901            assert_eq!(stats.retries, 1);
2902        }
2903    }
2904
2905    #[tokio::test]
2906    async fn unknown_retry_budget_preserves_one_unresolved_logical_batch() {
2907        for max_retries in [0, 1] {
2908            for disconnect in [false, true] {
2909                let server = MockServer::start(move |index, request| {
2910                    if disconnect {
2911                        MockResponse::committed(request).without_ack()
2912                    } else {
2913                        MockResponse::json(
2914                            503,
2915                            format!(
2916                                r#"{{
2917                            "message":"still ambiguous",
2918                            "status":"unknown",
2919                            "request_id":"attempt-{index}"
2920                        }}"#
2921                            ),
2922                        )
2923                    }
2924                });
2925                let stream = server
2926                    .client()
2927                    .table("events")
2928                    .append_stream()
2929                    .retry_unknown(true)
2930                    .max_retries(max_retries)
2931                    .initial_backoff(Duration::ZERO)
2932                    .build()
2933                    .unwrap();
2934                stream.send(&serde_json::json!({"id": 1})).await.unwrap();
2935
2936                let error = stream.shutdown().await.unwrap_err();
2937
2938                assert_eq!(server.requests().len(), max_retries + 1);
2939                assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
2940                assert!(format!("{error:?}").contains(&format!("retries: {max_retries}")));
2941                if disconnect {
2942                    assert!(StdError::source(&error).unwrap().is::<reqwest::Error>());
2943                } else {
2944                    assert_eq!(
2945                        error.http_status(),
2946                        Some(reqwest::StatusCode::SERVICE_UNAVAILABLE)
2947                    );
2948                    assert_eq!(
2949                        error.request_id(),
2950                        Some(format!("attempt-{max_retries}").as_str())
2951                    );
2952                }
2953                let stats = stream.stats();
2954                assert_eq!(stats.accepted_rows, 1);
2955                assert_eq!(stats.unknown_rows, 1);
2956                assert_eq!(stats.failed_rows, 0);
2957                assert_eq!(stats.retries, max_retries as u64);
2958                assert_eq!(stats.pending_rows, 0);
2959            }
2960        }
2961    }
2962
2963    #[tokio::test]
2964    async fn explicit_replay_retries_transient_unknown_responses() {
2965        for status in [200, 204, 408, 429, 500, 503] {
2966            let server = MockServer::start(move |index, request| match index {
2967                0 => MockResponse::json(status, "invalid response"),
2968                1 => MockResponse::committed(request),
2969                _ => panic!("unexpected retry request {index}"),
2970            });
2971            let stream = server
2972                .client()
2973                .table("events")
2974                .append_stream()
2975                .retry_unknown(true)
2976                .initial_backoff(Duration::ZERO)
2977                .build()
2978                .unwrap();
2979            stream.send(&serde_json::json!({"id": 1})).await.unwrap();
2980
2981            let report = stream.shutdown().await.unwrap();
2982
2983            assert_eq!(server.requests().len(), 2, "HTTP {status}");
2984            assert_eq!(report.committed_rows, 1);
2985            assert_eq!(report.retries, 1);
2986        }
2987    }
2988
2989    #[tokio::test]
2990    async fn default_delivery_does_not_retry_permanent_unknown_responses() {
2991        for status in [400, 401, 403, 404, 413, 422] {
2992            let server = MockServer::start(move |_, _| {
2993                MockResponse::json(status, r#"{"message":"permanent failure"}"#)
2994            });
2995            let stream = server
2996                .client()
2997                .table("events")
2998                .append_stream()
2999                .build()
3000                .unwrap();
3001            stream.send(&serde_json::json!({"id": 1})).await.unwrap();
3002
3003            let error = stream.shutdown().await.unwrap_err();
3004
3005            assert_eq!(server.requests().len(), 1, "HTTP {status}");
3006            assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
3007            assert_eq!(error.http_status().unwrap().as_u16(), status);
3008            assert_eq!(stream.stats().retries, 0);
3009        }
3010    }
3011
3012    #[tokio::test]
3013    async fn continue_barrier_settles_mixed_results_after_unknown_retry() {
3014        let ambiguous_attempts = AtomicUsize::new(0);
3015        let server = MockServer::start(move |_, request| {
3016            let row: serde_json::Value = serde_json::from_str(&request.body).unwrap();
3017            match row["id"].as_u64().unwrap() {
3018                0 => MockResponse::committed(request),
3019                1 if ambiguous_attempts.fetch_add(1, Ordering::AcqRel) == 0 => MockResponse::json(
3020                    503,
3021                    r#"{"message":"ambiguous","code":"ingest_failure","details":{"status":"unknown"}}"#,
3022                ),
3023                _ => MockResponse::json(
3024                    422,
3025                    r#"{"message":"invalid","code":"ingest_failure","details":{"status":"rejected"}}"#,
3026                ),
3027            }
3028        });
3029        let stream = server
3030            .client()
3031            .table("events")
3032            .append_stream()
3033            .retry_unknown(true)
3034            .failure_policy(AppendFailurePolicy::Continue)
3035            .circuit_breaker(None)
3036            .max_batch_rows(1)
3037            .initial_backoff(Duration::ZERO)
3038            .build()
3039            .unwrap();
3040        stream
3041            .send_all((0..3).map(|id| serde_json::json!({"id": id})))
3042            .await
3043            .unwrap();
3044
3045        let report = stream.flush().await.unwrap();
3046
3047        assert_eq!(server.requests().len(), 4);
3048        assert_eq!(report.outcome, AppendDeliveryOutcome::Partial);
3049        assert_eq!(report.accepted_rows, 3);
3050        assert_eq!(report.committed_rows, 1);
3051        assert_eq!(report.failed_rows, 1);
3052        assert_eq!(report.unknown_rows, 1);
3053        assert_eq!(report.retries, 1);
3054        assert_eq!(
3055            report.accepted_rows,
3056            report.committed_rows + report.failed_rows + report.unknown_rows
3057        );
3058        assert_eq!(stream.stats().pending_rows, 0);
3059        assert_eq!(stream.shutdown().await.unwrap().accepted_rows, 0);
3060        assert_eq!(stream.stats().committed_rows, 1);
3061    }
3062
3063    #[tokio::test]
3064    async fn continue_mode_retries_temporary_rejections_and_unknown_outcomes() {
3065        let server = MockServer::start(|index, request| match index {
3066            0 => MockResponse::json(
3067                503,
3068                r#"{"message":"busy","code":"ingest_failure","details":{"status":"rejected"}}"#,
3069            ),
3070            1 => MockResponse::committed(request),
3071            2 => MockResponse::json(
3072                503,
3073                r#"{"message":"ambiguous","code":"ingest_failure","details":{"status":"unknown"}}"#,
3074            ),
3075            3 | 4 => MockResponse::json(
3076                422,
3077                r#"{"message":"bad row","code":"ingest_failure","details":{"status":"rejected"}}"#,
3078            ),
3079            _ => panic!("unexpected retry request {index}"),
3080        });
3081        let stream = server
3082            .client()
3083            .table("events")
3084            .append_stream()
3085            .retry_unknown(true)
3086            .failure_policy(AppendFailurePolicy::Continue)
3087            .circuit_breaker(None)
3088            .target_batch_bytes(1)
3089            .max_concurrent_batches(1)
3090            .max_retries(4)
3091            .initial_backoff(Duration::ZERO)
3092            .max_backoff(Duration::ZERO)
3093            .build()
3094            .unwrap();
3095
3096        stream
3097            .send_all((0..3).map(|id| serde_json::json!({"id": id})))
3098            .await
3099            .unwrap();
3100        let report = stream.shutdown().await.unwrap();
3101
3102        let requests = server.requests();
3103        assert_eq!(requests.len(), 5);
3104        assert_eq!(requests[0].body, requests[1].body);
3105        assert_eq!(requests[2].body, requests[3].body);
3106        assert_ne!(requests[3].body, requests[4].body);
3107        assert_eq!(report.outcome, AppendDeliveryOutcome::Partial);
3108        assert_eq!(report.accepted_rows, 3);
3109        assert_eq!(report.committed_rows, 1);
3110        assert_eq!(report.failed_rows, 1);
3111        assert_eq!(report.unknown_rows, 1);
3112        assert_eq!(report.committed_batches, 1);
3113        assert_eq!(report.failed_batches, 1);
3114        assert_eq!(report.unknown_batches, 1);
3115        assert_eq!(report.retries, 2);
3116        assert_eq!(
3117            report.accepted_rows,
3118            report.committed_rows + report.failed_rows + report.unknown_rows
3119        );
3120    }
3121
3122    #[tokio::test]
3123    async fn zero_retry_budget_reports_unknown_timeout() {
3124        let server = MockServer::start(|_, request| {
3125            MockResponse::committed(request).with_delay(Duration::from_millis(100))
3126        });
3127        let stream = server
3128            .client()
3129            .table("events")
3130            .append_stream()
3131            .failure_policy(AppendFailurePolicy::Continue)
3132            .circuit_breaker(None)
3133            .target_batch_bytes(1)
3134            .max_retries(0)
3135            .attempt_timeout(Duration::from_millis(10))
3136            .build()
3137            .unwrap();
3138
3139        stream.send(&serde_json::json!({"id": 1})).await.unwrap();
3140        let report = stream.shutdown().await.unwrap();
3141
3142        assert_eq!(server.requests().len(), 1);
3143        assert_eq!(report.outcome, AppendDeliveryOutcome::Unknown);
3144        assert_eq!(report.unknown_rows, 1);
3145        assert_eq!(report.retries, 0);
3146    }
3147
3148    #[tokio::test]
3149    async fn write_timeout_bounds_direct_append_including_response_body() {
3150        for body_delay in [false, true] {
3151            let server = MockServer::start(move |_, _| {
3152                let response =
3153                    MockResponse::json(200, r#"{"status":"committed","num_rows_inserted":1}"#);
3154                if body_delay {
3155                    response.with_body_delay(Duration::from_millis(150))
3156                } else {
3157                    response.with_delay(Duration::from_millis(150))
3158                }
3159            });
3160            let client = server.client_with_write_timeout(Duration::from_millis(50));
3161            let error = client
3162                .table("events")
3163                .append("{\"id\":1}".to_string())
3164                .await
3165                .unwrap_err();
3166            assert!(!error.is_retryable());
3167            assert_write_timeout_contract(&error);
3168            let source = std::error::Error::source(&error).unwrap();
3169            assert!(source.is::<tokio::time::error::Elapsed>());
3170            assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
3171            if body_delay {
3172                assert_eq!(error.http_status(), Some(reqwest::StatusCode::OK));
3173            }
3174            assert_eq!(server.requests().len(), 1);
3175        }
3176    }
3177
3178    #[tokio::test]
3179    async fn append_stream_policies_inherit_the_client_write_timeout() {
3180        for policy in [AppendFailurePolicy::Stop, AppendFailurePolicy::Continue] {
3181            let server = MockServer::start(|_, request| {
3182                MockResponse::committed(request).with_delay(Duration::from_millis(150))
3183            });
3184            let stream = server
3185                .client_with_write_timeout(Duration::from_millis(50))
3186                .table("events")
3187                .append_stream()
3188                .failure_policy(policy)
3189                .build()
3190                .unwrap();
3191            stream.send(&serde_json::json!({"id":1})).await.unwrap();
3192            let result = stream.shutdown().await;
3193            if policy == AppendFailurePolicy::Continue {
3194                let report = result.unwrap();
3195                assert_eq!(report.outcome, AppendDeliveryOutcome::Unknown);
3196            } else {
3197                let error = result.unwrap_err();
3198                assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
3199                assert_write_timeout_contract(&error);
3200                assert_write_timeout_contract(&stream.flush().await.unwrap_err());
3201            }
3202            assert_eq!(stream.stats().unknown_rows, 1);
3203            assert_write_timeout_contract(&stream.stats().last_failure.unwrap().error);
3204            assert_eq!(stream.stats().retries, 0);
3205            assert_eq!(server.requests().len(), 1);
3206        }
3207    }
3208
3209    #[tokio::test]
3210    async fn attempt_timeout_can_exceed_the_client_write_timeout() {
3211        let server = MockServer::start(|_, request| {
3212            MockResponse::committed(request).with_delay(Duration::from_millis(50))
3213        });
3214        let stream = server
3215            .client_with_write_timeout(Duration::from_millis(10))
3216            .table("events")
3217            .append_stream()
3218            .attempt_timeout(Duration::from_secs(1))
3219            .build()
3220            .unwrap();
3221        stream.send(&serde_json::json!({"id":1})).await.unwrap();
3222        assert_eq!(stream.shutdown().await.unwrap().committed_rows, 1);
3223    }
3224
3225    #[tokio::test]
3226    async fn write_timeout_preserves_a_shorter_custom_http_timeout() {
3227        for (body_delay, compressed) in [(false, false), (true, false), (true, true)] {
3228            let server = MockServer::start(move |_, _| {
3229                let mut response =
3230                    MockResponse::json(200, r#"{"status":"committed","num_rows_inserted":1}"#);
3231                if compressed {
3232                    response.body = zstd::stream::encode_all(response.body.as_slice(), 1).unwrap();
3233                    response = response.with_header("Content-Encoding", "zstd");
3234                }
3235                if body_delay {
3236                    response.with_body_delay(Duration::from_millis(150))
3237                } else {
3238                    response.with_delay(Duration::from_millis(150))
3239                }
3240            });
3241            let http = reqwest::Client::builder()
3242                .no_proxy()
3243                .timeout(Duration::from_millis(50))
3244                .build()
3245                .unwrap();
3246            let client = Client::builder(&server.endpoint)
3247                .http_client(http)
3248                .write_timeout(Duration::from_secs(1))
3249                .build()
3250                .unwrap();
3251            let error = client
3252                .table("events")
3253                .append("{\"id\":1}".to_string())
3254                .await
3255                .unwrap_err();
3256            let transport_error = std::error::Error::source(&error)
3257                .and_then(|source| source.downcast_ref::<reqwest::Error>())
3258                .expect("the native transport cause must remain available");
3259            assert!(transport_error.is_timeout());
3260            assert_write_timeout_contract(&error);
3261            if body_delay {
3262                assert_eq!(error.http_status(), Some(reqwest::StatusCode::OK));
3263            }
3264            assert_eq!(server.requests().len(), 1);
3265        }
3266    }
3267
3268    fn assert_write_timeout_contract(error: &crate::Error) {
3269        let details = error.write_details().unwrap();
3270        assert_eq!(details.reason, crate::WriteErrorReason::Timeout);
3271        assert_eq!(details.outcome, AppendState::Unknown);
3272        assert!(!error.is_retryable());
3273        assert_eq!(
3274            serde_json::to_value(details).unwrap(),
3275            serde_json::json!({"reason":"timeout","outcome":"unknown"})
3276        );
3277    }
3278
3279    #[tokio::test]
3280    async fn write_timeout_starts_after_materializing_the_request() {
3281        struct SlowPayload;
3282        impl From<SlowPayload> for String {
3283            fn from(_: SlowPayload) -> Self {
3284                thread::sleep(Duration::from_millis(100));
3285                "{\"id\":1}".to_string()
3286            }
3287        }
3288        let server = MockServer::start(|_, request| MockResponse::committed(request));
3289        let result = server
3290            .client_with_write_timeout(Duration::from_millis(50))
3291            .append_rows("scopedb", "public", "events", SlowPayload)
3292            .await
3293            .unwrap();
3294        assert_eq!(result.num_rows_inserted, 1);
3295        assert_eq!(server.requests().len(), 1);
3296    }
3297
3298    #[tokio::test]
3299    async fn server_timeout_rejection_remains_safely_retryable() {
3300        let server = MockServer::start(|_, _| {
3301            MockResponse::json(
3302                408,
3303                r#"{"message":"admission timed out","retryable":true,"code":"ingest_failure","details":{"status":"rejected"}}"#,
3304            )
3305        });
3306        let error = server
3307            .client()
3308            .table("events")
3309            .append("{\"id\":1}".to_string())
3310            .await
3311            .unwrap_err();
3312        assert!(error.write_details().is_none());
3313        assert_eq!(
3314            error.append_details().unwrap().status,
3315            AppendState::Rejected
3316        );
3317        assert!(error.is_retryable());
3318    }
3319
3320    #[tokio::test]
3321    async fn connection_loss_is_not_replayed_when_append_retries_are_disabled() {
3322        for disconnect in [MockDisconnect::BeforeHeaders, MockDisconnect::DuringBody] {
3323            let server = MockServer::start(move |_, request| {
3324                MockResponse::committed(request).with_disconnect(disconnect)
3325            });
3326            let client = server.client();
3327            let direct = client.table("events").append(r#"{"id":1}"#).await;
3328            let stream = client
3329                .table("events")
3330                .append_stream()
3331                .max_retries(0)
3332                .build()
3333                .unwrap();
3334            stream.send(&serde_json::json!({"id":1})).await.unwrap();
3335            for error in [direct.unwrap_err(), stream.shutdown().await.unwrap_err()] {
3336                assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
3337                assert!(!error.is_retryable());
3338                if matches!(disconnect, MockDisconnect::DuringBody) {
3339                    assert_eq!(error.http_status(), Some(reqwest::StatusCode::OK));
3340                }
3341            }
3342            assert_eq!(server.requests().len(), 2);
3343        }
3344    }
3345
3346    #[tokio::test]
3347    async fn stop_policy_closes_after_a_rejected_batch() {
3348        let server = MockServer::start(|_, _| {
3349            MockResponse::json(
3350                422,
3351                r#"{"message":"invalid","code":"ingest_failure","details":{"status":"rejected"}}"#,
3352            )
3353            .with_delay(Duration::from_millis(20))
3354        });
3355        let stream = server
3356            .client()
3357            .table("events")
3358            .append_stream()
3359            .target_batch_bytes(1)
3360            .max_concurrent_batches(1)
3361            .build()
3362            .unwrap();
3363        stream
3364            .send_all((0..2).map(|id| serde_json::json!({"id": id})))
3365            .await
3366            .unwrap();
3367
3368        let error = stream.shutdown().await.unwrap_err();
3369        assert_eq!(error.kind(), ErrorKind::AppendRowsFailed);
3370        assert_eq!(
3371            error.append_details().map(|details| details.status),
3372            Some(AppendState::Rejected)
3373        );
3374        let stats = stream.stats();
3375        assert_eq!(stats.state, AppendStreamState::Failed);
3376        assert_eq!(stats.accepted_rows, 2);
3377        assert_eq!(stats.failed_rows, 2);
3378        assert_eq!(stats.pending_rows, 0);
3379        let report = stats.last_report.unwrap();
3380        assert_eq!(report.accepted_rows, 2);
3381        assert_eq!(report.failed_rows, 2);
3382        assert!(matches!(
3383            stream.try_send(&serde_json::json!({"id": 3})),
3384            Err(AppendTrySendError::Closed)
3385        ));
3386    }
3387
3388    #[tokio::test]
3389    async fn failed_barriers_report_the_time_spent_waiting() {
3390        for shutdown in [false, true] {
3391            let (arrived_tx, arrived_rx) = oneshot::channel();
3392            let arrived_tx = Mutex::new(Some(arrived_tx));
3393            let (release_tx, release_rx) = std::sync::mpsc::channel();
3394            let release_rx = Mutex::new(release_rx);
3395            let server = MockServer::start(move |_, _| {
3396                lock(&arrived_tx).take().unwrap().send(()).unwrap();
3397                lock(&release_rx)
3398                    .recv_timeout(Duration::from_secs(3))
3399                    .unwrap();
3400                MockResponse::json(
3401                    422,
3402                    r#"{"message":"invalid","code":"invalid_schema","details":{"status":"rejected"}}"#,
3403                )
3404            });
3405            let stream = server
3406                .client()
3407                .table("events")
3408                .append_stream()
3409                .flush_interval(Duration::from_secs(60))
3410                .build()
3411                .unwrap();
3412            stream.send(&serde_json::json!({"id":1})).await.unwrap();
3413            let mut barrier = Box::pin(async {
3414                if shutdown {
3415                    stream.shutdown().await
3416                } else {
3417                    stream.flush().await
3418                }
3419            });
3420            tokio::select! {
3421                _ = &mut barrier => panic!("response must remain gated"),
3422                arrived = arrived_rx => arrived.unwrap(),
3423            }
3424            let waiting = Instant::now();
3425            // Measure an actual wait after the barrier has dispatched its batch.
3426            tokio::time::sleep(Duration::from_millis(25)).await;
3427            let minimum_duration = waiting.elapsed();
3428            release_tx.send(()).unwrap();
3429            assert!(barrier.await.is_err());
3430            let report = stream.stats().last_report.unwrap();
3431            assert_eq!(report.accepted_rows, 1);
3432            assert_eq!(report.failed_rows, 1);
3433            assert!(report.duration >= minimum_duration, "{report:?}");
3434            assert_eq!(server.requests().len(), 1);
3435        }
3436    }
3437
3438    #[tokio::test]
3439    async fn strict_terminal_waits_for_every_in_flight_outcome() {
3440        use std::future::Future;
3441        use std::task::Poll;
3442        let both_started = std::sync::Barrier::new(2);
3443        let (release_tx, release_rx) = std::sync::mpsc::channel();
3444        let release_rx = Mutex::new(release_rx);
3445        let server = MockServer::start(move |index, _| {
3446            both_started.wait();
3447            if index == 0 {
3448                MockResponse::json(
3449                    422,
3450                    r#"{"message":"invalid","code":"invalid_schema","details":{"status":"rejected"}}"#,
3451                )
3452            } else {
3453                lock(&release_rx)
3454                    .recv_timeout(Duration::from_secs(3))
3455                    .unwrap();
3456                MockResponse::json(
3457                    503,
3458                    r#"{"message":"ambiguous","code":"ingest_failure","details":{"status":"unknown"}}"#,
3459                )
3460            }
3461        });
3462        let (failed_tx, failed_rx) = oneshot::channel();
3463        let failed_tx = Mutex::new(Some(failed_tx));
3464        let stream = server
3465            .client()
3466            .table("events")
3467            .append_stream()
3468            .max_batch_rows(1)
3469            .max_concurrent_batches(2)
3470            .max_retries(0)
3471            .on_batch_failure(move |_| {
3472                if let Some(tx) = lock(&failed_tx).take() {
3473                    let _ = tx.send(());
3474                }
3475            })
3476            .build()
3477            .unwrap();
3478        stream
3479            .send_all((0..2).map(|id| serde_json::json!({"id":id})))
3480            .await
3481            .unwrap();
3482        tokio::time::timeout(Duration::from_secs(3), failed_rx)
3483            .await
3484            .unwrap()
3485            .unwrap();
3486        let mut shutdown = Box::pin(stream.shutdown());
3487        std::future::poll_fn(|cx| {
3488            assert!(shutdown.as_mut().poll(cx).is_pending());
3489            Poll::Ready(())
3490        })
3491        .await;
3492        release_tx.send(()).unwrap();
3493        let error = shutdown.await.unwrap_err();
3494        assert_eq!(error.append_details().unwrap().status, AppendState::Unknown);
3495        let stats = stream.stats();
3496        assert_eq!(stats.failed_rows, 1);
3497        assert_eq!(stats.unknown_rows, 1);
3498        assert_eq!(stats.pending_rows, 0);
3499        let report = stats.last_report.unwrap();
3500        assert_eq!(report.accepted_rows, 2);
3501        assert_eq!(report.failed_rows, 1);
3502        assert_eq!(report.unknown_rows, 1);
3503    }
3504
3505    #[tokio::test]
3506    async fn concurrent_flush_and_shutdown_both_settle() {
3507        let server = MockServer::start(|_, request| {
3508            MockResponse::committed(request).with_delay(Duration::from_millis(5))
3509        });
3510        let stream = server
3511            .client()
3512            .table("events")
3513            .append_stream()
3514            .build()
3515            .unwrap();
3516        stream.send(&serde_json::json!({"id": 1})).await.unwrap();
3517
3518        let (flush, shutdown) = tokio::join!(stream.flush(), stream.shutdown());
3519        assert!(flush.is_ok(), "flush failed: {flush:?}");
3520        assert!(shutdown.is_ok(), "shutdown failed: {shutdown:?}");
3521        assert_eq!(stream.stats().state, AppendStreamState::Closed);
3522    }
3523
3524    #[tokio::test]
3525    async fn shutdown_is_idempotent_across_clones() {
3526        let server = MockServer::start(|_, request| MockResponse::committed(request));
3527        let stream = server
3528            .client()
3529            .table("events")
3530            .append_stream()
3531            .build()
3532            .unwrap();
3533        let clone = stream.clone();
3534        stream.send(&serde_json::json!({"id": 1})).await.unwrap();
3535
3536        let (first, second) = tokio::join!(stream.shutdown(), clone.shutdown());
3537        assert_eq!(first.unwrap(), second.unwrap());
3538        assert_eq!(server.requests().len(), 1);
3539        assert!(matches!(
3540            stream.try_send(&serde_json::json!({"id": 2})),
3541            Err(AppendTrySendError::Closed)
3542        ));
3543    }
3544
3545    #[tokio::test]
3546    async fn try_send_rejects_rows_exceeding_buffer_capacity() {
3547        let client = Client::new("http://127.0.0.1:1", reqwest::Client::new()).unwrap();
3548        let stream = client
3549            .table("events")
3550            .append_stream()
3551            .failure_policy(AppendFailurePolicy::Continue)
3552            .max_buffered_bytes(1)
3553            .build()
3554            .unwrap();
3555
3556        assert!(matches!(
3557            stream.try_send(&serde_json::json!({"id": 1})),
3558            Err(AppendTrySendError::RecordTooLarge { bytes: 9, limit: 1 })
3559        ));
3560        let report = stream.shutdown().await.unwrap();
3561        assert_eq!(report.outcome, AppendDeliveryOutcome::Failed);
3562        assert_eq!(report.accepted_rows, 0);
3563        assert_eq!(report.dropped_rows, 1);
3564        assert_eq!(stream.stats().dropped_by_reason.buffer_full, 0);
3565        assert_eq!(stream.stats().dropped_by_reason.record_too_large, 1);
3566    }
3567}