1use 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#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
57#[non_exhaustive]
58pub enum AppendFailurePolicy {
59 #[default]
61 Stop,
62 Continue,
64}
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68#[non_exhaustive]
69pub enum AppendStreamState {
70 Open,
71 Closing,
72 Closed,
73 Failed,
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78#[non_exhaustive]
79pub enum AppendCircuitState {
80 Closed,
81 Open,
82 HalfOpen,
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87#[non_exhaustive]
88pub enum AppendDeliveryOutcome {
89 Ok,
91 Partial,
93 Failed,
95 Unknown,
97}
98
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub struct AppendAdmissionResult {
102 pub accepted_rows: u64,
103}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
109pub struct AppendDeliveryReport {
110 pub outcome: AppendDeliveryOutcome,
111 pub accepted_rows: u64,
112 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 pub duration: Duration,
125}
126
127#[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#[derive(Debug, Clone)]
139pub struct AppendLastFailure {
140 pub at: SystemTime,
141 pub error: Error,
142}
143
144#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
165#[non_exhaustive]
166pub enum AppendBatchFailureAction {
167 Continuing,
168 CircuitOpened,
169 Stopped,
170}
171
172pub 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#[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#[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
265pub 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 pub fn failure_policy(mut self, failure_policy: AppendFailurePolicy) -> Self {
345 self.failure_policy = failure_policy;
346 self
347 }
348
349 pub fn target_batch_bytes(mut self, target_batch_bytes: usize) -> Self {
352 self.batch_bytes = target_batch_bytes;
353 self
354 }
355
356 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 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 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 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 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 pub fn attempt_timeout(mut self, attempt_timeout: Duration) -> Self {
411 self.attempt_timeout = Some(attempt_timeout);
412 self
413 }
414
415 pub fn circuit_breaker(mut self, circuit_breaker: Option<AppendCircuitBreakerOptions>) -> Self {
419 self.circuit_breaker = circuit_breaker;
420 self
421 }
422
423 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#[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 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 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 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 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 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 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 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 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 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}