1use std::future::Future;
19use std::pin::Pin;
20use std::sync::Arc;
21
22use a2a_protocol_types::error::{A2aError, A2aResult};
23use a2a_protocol_types::events::StreamResponse;
24use tokio::sync::{broadcast, mpsc};
25
26use super::{EventQueueReader, EventQueueWriter};
27
28struct CountingWriter(usize);
34
35impl std::io::Write for CountingWriter {
36 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
37 self.0 += buf.len();
38 Ok(buf.len())
39 }
40
41 fn flush(&mut self) -> std::io::Result<()> {
42 Ok(())
43 }
44}
45
46#[derive(Clone)]
58pub struct InMemoryQueueWriter {
59 tx: broadcast::Sender<A2aResult<StreamResponse>>,
60 persistence_tx: Option<mpsc::Sender<A2aResult<StreamResponse>>>,
67 max_event_size: usize,
69 write_timeout: std::time::Duration,
74 metrics: Option<Arc<dyn crate::metrics::Metrics>>,
77}
78
79impl std::fmt::Debug for InMemoryQueueWriter {
80 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81 f.debug_struct("InMemoryQueueWriter")
82 .field("persistence", &self.persistence_tx.is_some())
83 .field("max_event_size", &self.max_event_size)
84 .field("write_timeout", &self.write_timeout)
85 .field("metrics", &self.metrics.is_some())
86 .finish_non_exhaustive()
87 }
88}
89
90impl InMemoryQueueWriter {
91 pub(super) const fn new(
93 tx: broadcast::Sender<A2aResult<StreamResponse>>,
94 max_event_size: usize,
95 write_timeout: std::time::Duration,
96 ) -> Self {
97 Self {
98 tx,
99 persistence_tx: None,
100 max_event_size,
101 write_timeout,
102 metrics: None,
103 }
104 }
105
106 pub(super) const fn new_with_persistence(
108 tx: broadcast::Sender<A2aResult<StreamResponse>>,
109 persistence_tx: mpsc::Sender<A2aResult<StreamResponse>>,
110 max_event_size: usize,
111 write_timeout: std::time::Duration,
112 ) -> Self {
113 Self {
114 tx,
115 persistence_tx: Some(persistence_tx),
116 max_event_size,
117 write_timeout,
118 metrics: None,
119 }
120 }
121
122 #[must_use]
125 pub(crate) fn with_metrics(mut self, metrics: Arc<dyn crate::metrics::Metrics>) -> Self {
126 self.metrics = Some(metrics);
127 self
128 }
129
130 #[must_use]
135 pub fn subscribe(&self) -> InMemoryQueueReader {
136 InMemoryQueueReader::new(self.tx.subscribe())
137 }
138
139 pub(crate) fn raw_subscribe(&self) -> broadcast::Receiver<A2aResult<StreamResponse>> {
144 self.tx.subscribe()
145 }
146}
147
148#[allow(clippy::manual_async_fn)]
149impl EventQueueWriter for InMemoryQueueWriter {
150 fn write<'a>(
151 &'a self,
152 event: StreamResponse,
153 ) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
154 Box::pin(async move {
155 let serialized_size = {
160 let mut counter = CountingWriter(0);
161 serde_json::to_writer(&mut counter, &event)
162 .map_err(|e| A2aError::internal(format!("event serialization failed: {e}")))?;
163 counter.0
164 };
165 if serialized_size > self.max_event_size {
166 return Err(A2aError::internal(format!(
167 "event size {serialized_size} bytes exceeds maximum {} bytes",
168 self.max_event_size
169 )));
170 }
171 if let Some(ref persistence_tx) = self.persistence_tx {
184 match persistence_tx
185 .send_timeout(Ok(event.clone()), self.write_timeout)
186 .await
187 {
188 Ok(()) => {}
189 Err(mpsc::error::SendTimeoutError::Closed(_)) => {
190 trace_warn!("persistence channel closed, event not persisted");
191 if let Some(metrics) = &self.metrics {
196 metrics.on_persistence_error(
197 crate::metrics::persistence_operation::QUEUE_HANDOFF,
198 crate::metrics::queue_handoff_error::CHANNEL_CLOSED,
199 );
200 }
201 }
202 Err(mpsc::error::SendTimeoutError::Timeout(_)) => {
203 trace_warn!(
204 timeout_ms =
205 u64::try_from(self.write_timeout.as_millis()).unwrap_or(u64::MAX),
206 "persistence channel full; background processor is not draining"
207 );
208 return Err(A2aError::internal(format!(
209 "event queue: the persistence channel was still full after {:?}; \
210 the background processor is not draining events",
211 self.write_timeout
212 )));
213 }
214 }
215 }
216 match self.tx.send(Ok(event)) {
225 Ok(_) => Ok(()),
226 Err(_) if self.persistence_tx.is_some() => {
227 trace_warn!("no live event subscribers; event persisted only");
228 Ok(())
229 }
230 Err(_) => Err(A2aError::internal("event queue: no active receivers")),
231 }
232 })
233 }
234
235 fn close<'a>(&'a self) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
236 Box::pin(async move {
237 Ok(())
240 })
241 }
242}
243
244pub struct InMemoryQueueReader {
255 rx: broadcast::Receiver<A2aResult<StreamResponse>>,
256 pending_first: Option<A2aResult<StreamResponse>>,
257 reattach: Option<ReattachFn>,
259 saw_terminal: bool,
263}
264
265#[allow(clippy::missing_fields_in_debug)]
269impl std::fmt::Debug for InMemoryQueueReader {
270 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
271 f.debug_struct("InMemoryQueueReader")
272 .field("pending_first", &self.pending_first.is_some())
273 .field("reattach", &self.reattach.is_some())
274 .field("saw_terminal", &self.saw_terminal)
275 .finish()
276 }
277}
278
279#[allow(clippy::large_enum_variant)]
283pub enum Reattached {
284 Channel(broadcast::Receiver<A2aResult<StreamResponse>>),
286 Final(StreamResponse),
291 End,
293}
294
295pub type ReattachFn =
298 Arc<dyn Fn() -> Pin<Box<dyn Future<Output = Reattached> + Send>> + Send + Sync>;
299
300const fn carries_terminal_state(event: &StreamResponse) -> bool {
302 match event {
303 StreamResponse::Task(t) => t.status.state.is_terminal(),
304 StreamResponse::StatusUpdate(u) => u.status.state.is_terminal(),
305 _ => false,
306 }
307}
308
309impl InMemoryQueueReader {
310 pub(crate) fn with_reattach(mut self, reattach: ReattachFn) -> Self {
325 self.reattach = Some(reattach);
326 self
327 }
328
329 pub(crate) const fn new(rx: broadcast::Receiver<A2aResult<StreamResponse>>) -> Self {
331 Self {
332 rx,
333 pending_first: None,
334 reattach: None,
335 saw_terminal: false,
336 }
337 }
338
339 pub fn set_first_event(&mut self, event: StreamResponse) {
341 self.pending_first = Some(Ok(event));
342 }
343
344 pub(crate) const fn with_first_event(
346 rx: broadcast::Receiver<A2aResult<StreamResponse>>,
347 first: StreamResponse,
348 ) -> Self {
349 Self {
350 rx,
351 pending_first: Some(Ok(first)),
352 reattach: None,
353 saw_terminal: false,
354 }
355 }
356
357 pub(crate) fn snapshot_then_end(first: StreamResponse) -> Self {
364 let (tx, rx) = broadcast::channel(1);
367 drop(tx);
368 Self {
369 rx,
370 pending_first: Some(Ok(first)),
371 reattach: None,
372 saw_terminal: false,
373 }
374 }
375}
376
377fn lag_error(dropped: u64) -> a2a_protocol_types::error::A2aError {
393 a2a_protocol_types::error::A2aError::stream_lagged(dropped)
394}
395
396#[allow(clippy::redundant_pub_crate)] pub(crate) fn is_lag_error(err: &a2a_protocol_types::error::A2aError) -> bool {
400 err.is_stream_lagged()
401}
402
403impl EventQueueReader for InMemoryQueueReader {
404 fn read(
405 &mut self,
406 ) -> Pin<Box<dyn Future<Output = Option<A2aResult<StreamResponse>>> + Send + '_>> {
407 Box::pin(async move {
408 if let Some(first) = self.pending_first.take() {
411 if let Ok(ref ev) = first {
412 self.saw_terminal |= carries_terminal_state(ev);
413 }
414 return Some(first);
415 }
416 loop {
417 match self.rx.recv().await {
418 Ok(event) => {
419 if let Ok(ref ev) = event {
420 self.saw_terminal |= carries_terminal_state(ev);
421 }
422 return Some(event);
423 }
424 Err(broadcast::error::RecvError::Lagged(n)) => {
425 trace_warn!(
426 dropped_events = n,
427 "event queue reader lagged, {n} events dropped"
428 );
429 return Some(Err(lag_error(n)));
430 }
431 Err(broadcast::error::RecvError::Closed) => {
432 if self.saw_terminal {
437 return None;
438 }
439 let reattach = self.reattach.as_ref()?;
440 match reattach().await {
441 Reattached::Channel(rx) => self.rx = rx,
442 Reattached::Final(event) => {
443 self.saw_terminal = true;
444 return Some(Ok(event));
445 }
446 Reattached::End => return None,
447 }
448 }
449 }
450 }
451 })
452 }
453}
454
455#[cfg(test)]
456mod tests {
457 use super::*;
458 use crate::streaming::event_queue::{
459 DEFAULT_MAX_EVENT_SIZE, DEFAULT_WRITE_TIMEOUT, new_in_memory_queue,
460 new_in_memory_queue_with_options, new_in_memory_queue_with_persistence,
461 };
462 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
463 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
464
465 fn make_status_event(task_id: &str, state: TaskState) -> StreamResponse {
467 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
468 task_id: TaskId::new(task_id),
469 context_id: ContextId::new("ctx-test"),
470 status: TaskStatus {
471 state,
472 message: None,
473 timestamp: None,
474 },
475 metadata: None,
476 })
477 }
478
479 fn counting_reattach(flag: &Arc<std::sync::atomic::AtomicBool>) -> ReattachFn {
491 let flag = Arc::clone(flag);
492 Arc::new(move || {
493 let flag = Arc::clone(&flag);
494 Box::pin(async move {
495 flag.store(true, std::sync::atomic::Ordering::SeqCst);
496 Reattached::End
497 })
498 })
499 }
500
501 #[tokio::test]
506 async fn terminal_status_update_ends_the_stream_without_reattaching() {
507 use std::sync::atomic::{AtomicBool, Ordering};
508
509 let (writer, reader) = new_in_memory_queue();
510 let called = Arc::new(AtomicBool::new(false));
511 let mut reader = reader.with_reattach(counting_reattach(&called));
512
513 writer
514 .write(make_status_event("t-term", TaskState::Completed))
515 .await
516 .unwrap();
517 drop(writer);
518
519 assert!(reader.read().await.is_some(), "the terminal frame arrives");
520 assert!(
521 reader.read().await.is_none(),
522 "a stream that delivered a terminal state ends at close"
523 );
524 assert!(
525 !called.load(Ordering::SeqCst),
526 "the reattach hook must not be consulted once a terminal state has been seen"
527 );
528 }
529
530 #[tokio::test]
534 async fn terminal_task_snapshot_ends_the_stream_without_reattaching() {
535 use a2a_protocol_types::task::Task;
536 use std::sync::atomic::{AtomicBool, Ordering};
537
538 let (writer, reader) = new_in_memory_queue();
539 let called = Arc::new(AtomicBool::new(false));
540 let mut reader = reader.with_reattach(counting_reattach(&called));
541 reader.set_first_event(StreamResponse::Task(Task {
542 id: TaskId::new("t-snap"),
543 context_id: ContextId::new("ctx-test"),
544 status: TaskStatus {
545 state: TaskState::Completed,
546 message: None,
547 timestamp: None,
548 },
549 history: None,
550 artifacts: None,
551 metadata: None,
552 }));
553 drop(writer);
554
555 assert!(reader.read().await.is_some(), "the snapshot arrives first");
556 assert!(
557 reader.read().await.is_none(),
558 "a terminal Task snapshot ends the stream at close"
559 );
560 assert!(
561 !called.load(Ordering::SeqCst),
562 "the reattach hook must not be consulted after a terminal Task snapshot"
563 );
564 }
565
566 #[tokio::test]
571 async fn non_terminal_event_still_consults_the_reattach_hook() {
572 use std::sync::atomic::{AtomicBool, Ordering};
573
574 let (writer, reader) = new_in_memory_queue();
575 let called = Arc::new(AtomicBool::new(false));
576 let mut reader = reader.with_reattach(counting_reattach(&called));
577
578 writer
579 .write(make_status_event("t-working", TaskState::Working))
580 .await
581 .unwrap();
582 drop(writer);
583
584 assert!(reader.read().await.is_some(), "the working frame arrives");
585 assert!(reader.read().await.is_none(), "the hook here returns End");
586 assert!(
587 called.load(Ordering::SeqCst),
588 "without a terminal state the reader must ask the hook whether the task is done"
589 );
590 }
591
592 #[tokio::test]
598 async fn event_of_exactly_max_size_is_accepted() {
599 let event = make_status_event("t-exact", TaskState::Working);
600 let exact = serde_json::to_vec(&event).expect("serializes").len();
601
602 let (writer, _reader) = new_in_memory_queue_with_options(16, exact, DEFAULT_WRITE_TIMEOUT);
603 assert!(
604 writer.write(event).await.is_ok(),
605 "an event of exactly max_event_size ({exact} bytes) is within the \
606 limit and must be accepted"
607 );
608
609 let event = make_status_event("t-exact", TaskState::Working);
612 let (writer, _reader) =
613 new_in_memory_queue_with_options(16, exact - 1, DEFAULT_WRITE_TIMEOUT);
614 assert!(
615 writer.write(event).await.is_err(),
616 "one byte over the limit must still be rejected"
617 );
618 }
619
620 #[tokio::test]
624 async fn writer_debug_reports_the_decision_carrying_fields() {
625 let (writer, _reader) = new_in_memory_queue();
626 let rendered = format!("{writer:?}");
627 assert!(
628 rendered.contains("InMemoryQueueWriter"),
629 "the type name must appear: {rendered}"
630 );
631 assert!(
632 rendered.contains("persistence: false"),
633 "the persistence hand-off must be visible: {rendered}"
634 );
635 assert!(
636 rendered.contains("metrics: false"),
637 "the metrics hook must be visible: {rendered}"
638 );
639 }
640
641 #[tokio::test]
646 async fn reader_debug_reports_the_decision_carrying_fields() {
647 let (_writer, mut reader) = new_in_memory_queue();
648 reader.set_first_event(make_status_event("t-dbg", TaskState::Working));
649
650 let rendered = format!("{reader:?}");
651 assert!(
652 rendered.contains("InMemoryQueueReader"),
653 "the type name must appear: {rendered}"
654 );
655 assert!(
656 rendered.contains("pending_first: true"),
657 "a pending snapshot must be visible: {rendered}"
658 );
659 assert!(
660 rendered.contains("saw_terminal: false"),
661 "terminal tracking must be visible: {rendered}"
662 );
663 }
664
665 #[tokio::test]
673 async fn write_with_no_subscribers_succeeds_when_persistence_attached() {
674 let (writer, reader, mut persistence_rx) =
675 crate::streaming::event_queue::new_in_memory_queue_with_persistence(
676 8,
677 1024 * 1024,
678 std::time::Duration::from_secs(1),
679 );
680 drop(reader); writer
683 .write(make_status_event("t1", TaskState::Working))
684 .await
685 .expect("write must succeed with persistence attached");
686
687 let persisted = persistence_rx
688 .recv()
689 .await
690 .expect("persistence channel should have the event")
691 .expect("event should be Ok");
692 match persisted {
693 StreamResponse::StatusUpdate(evt) => {
694 assert_eq!(evt.status.state, TaskState::Working);
695 }
696 other => panic!("expected StatusUpdate, got: {other:?}"),
697 }
698 }
699
700 #[tokio::test]
704 async fn write_with_no_subscribers_fails_without_persistence() {
705 let (writer, reader) = new_in_memory_queue();
706 drop(reader);
707
708 let result = writer
709 .write(make_status_event("t1", TaskState::Working))
710 .await;
711 assert!(
712 result.is_err(),
713 "sync-mode write with no receivers must fail"
714 );
715 }
716
717 #[tokio::test]
718 async fn write_then_read_single_event() {
719 let (writer, mut reader) = new_in_memory_queue();
720 let event = make_status_event("t1", TaskState::Working);
721
722 writer.write(event).await.expect("write should succeed");
723 drop(writer);
724
725 let received = reader.read().await;
726 assert!(received.is_some(), "reader should return the written event");
727 let result = received.unwrap();
728 let event = result.expect("event should be Ok");
729 match &event {
730 StreamResponse::StatusUpdate(evt) => {
731 assert_eq!(
732 evt.status.state,
733 TaskState::Working,
734 "should be Working event"
735 );
736 }
737 other => panic!("expected StatusUpdate, got: {other:?}"),
738 }
739
740 let eof = reader.read().await;
742 assert!(
743 eof.is_none(),
744 "reader should return None after writer is dropped"
745 );
746 }
747
748 #[tokio::test]
749 async fn write_multiple_events_read_in_order() {
750 let (writer, mut reader) = new_in_memory_queue();
751
752 let e1 = make_status_event("t1", TaskState::Working);
753 let e2 = make_status_event("t1", TaskState::Completed);
754
755 writer.write(e1).await.expect("first write should succeed");
756 writer.write(e2).await.expect("second write should succeed");
757 drop(writer);
758
759 let r1 = reader.read().await.expect("should read first event");
761 let sr1 = r1.expect("first event should be Ok");
762 match &sr1 {
763 StreamResponse::StatusUpdate(evt) => {
764 assert_eq!(
765 evt.status.state,
766 TaskState::Working,
767 "first event should be Working"
768 );
769 }
770 other => panic!("expected StatusUpdate, got: {other:?}"),
771 }
772
773 let r2 = reader.read().await.expect("should read second event");
775 let sr2 = r2.expect("second event should be Ok");
776 match &sr2 {
777 StreamResponse::StatusUpdate(evt) => {
778 assert_eq!(
779 evt.status.state,
780 TaskState::Completed,
781 "second event should be Completed"
782 );
783 }
784 other => panic!("expected StatusUpdate, got: {other:?}"),
785 }
786
787 assert!(
789 reader.read().await.is_none(),
790 "should be EOF after all events"
791 );
792 }
793
794 #[tokio::test]
797 async fn read_returns_none_on_empty_closed_queue() {
798 let (writer, mut reader) = new_in_memory_queue();
799 drop(writer); let result = reader.read().await;
802 assert!(
803 result.is_none(),
804 "reading from an empty closed queue should return None"
805 );
806 }
807
808 #[tokio::test]
809 async fn write_after_all_readers_dropped_returns_error() {
810 let (writer, reader) = new_in_memory_queue();
811 drop(reader);
812
813 let result = writer
814 .write(make_status_event("t1", TaskState::Working))
815 .await;
816 assert!(
817 result.is_err(),
818 "writing with no active receivers should return an error"
819 );
820 }
821
822 #[tokio::test]
823 async fn close_is_no_op_and_succeeds() {
824 let (writer, _reader) = new_in_memory_queue();
825 let result = writer.close().await;
826 assert!(result.is_ok(), "close() should succeed");
827 }
828
829 #[tokio::test]
832 async fn subscribe_creates_independent_reader() {
833 let (writer, mut reader1) = new_in_memory_queue();
834 let mut reader2 = writer.subscribe();
835
836 let event = make_status_event("t1", TaskState::Working);
837 writer.write(event).await.expect("write should succeed");
838 drop(writer);
839
840 let r1 = reader1.read().await;
842 assert!(r1.is_some(), "reader1 should receive the event");
843
844 let r2 = reader2.read().await;
845 assert!(r2.is_some(), "reader2 should receive the event");
846
847 assert!(reader1.read().await.is_none(), "reader1 should see EOF");
849 assert!(reader2.read().await.is_none(), "reader2 should see EOF");
850 }
851
852 #[tokio::test]
853 async fn subscriber_only_sees_events_after_subscribe() {
854 let (writer, mut reader1) = new_in_memory_queue();
855
856 writer
858 .write(make_status_event("t1", TaskState::Submitted))
859 .await
860 .expect("write should succeed");
861
862 let mut reader2 = writer.subscribe();
864
865 writer
867 .write(make_status_event("t1", TaskState::Working))
868 .await
869 .expect("write should succeed");
870 drop(writer);
871
872 let r1a = reader1
874 .read()
875 .await
876 .expect("reader1 should see first event");
877 let evt1a = r1a.expect("first event should be Ok");
878 assert!(
879 matches!(&evt1a, StreamResponse::StatusUpdate(e) if e.status.state == TaskState::Submitted),
880 "reader1 first event should be Submitted"
881 );
882 let r1b = reader1
883 .read()
884 .await
885 .expect("reader1 should see second event");
886 let evt_1b = r1b.expect("second event should be Ok");
887 assert!(
888 matches!(&evt_1b, StreamResponse::StatusUpdate(e) if e.status.state == TaskState::Working),
889 "reader1 second event should be Working"
890 );
891 assert!(reader1.read().await.is_none());
892
893 let r2a = reader2
895 .read()
896 .await
897 .expect("reader2 should see second event");
898 let evt2a = r2a.expect("event should be Ok");
899 assert!(
900 matches!(&evt2a, StreamResponse::StatusUpdate(e) if e.status.state == TaskState::Working),
901 "reader2 should see Working event"
902 );
903 assert!(
904 reader2.read().await.is_none(),
905 "reader2 should see EOF after the one event it received"
906 );
907 }
908
909 #[tokio::test]
912 async fn oversized_event_is_rejected() {
913 let (writer, _reader) = new_in_memory_queue_with_options(
915 16,
916 10, DEFAULT_WRITE_TIMEOUT,
918 );
919
920 let event = make_status_event("t1", TaskState::Working);
921 let result = writer.write(event).await;
922 assert!(
923 result.is_err(),
924 "event exceeding max_event_size should be rejected"
925 );
926 let err = result.unwrap_err();
927 let msg = format!("{err}");
928 assert!(
929 msg.contains("exceeds maximum"),
930 "error message should mention size limit, got: {msg}"
931 );
932 }
933
934 #[test]
936 fn counting_writer_flush_is_noop() {
937 use std::io::Write;
938 let mut cw = super::CountingWriter(0);
939 cw.write_all(b"hello").unwrap();
940 assert_eq!(cw.0, 5);
941 cw.flush().unwrap();
943 assert_eq!(cw.0, 5, "flush should not change the count");
944 }
945
946 #[tokio::test]
947 async fn event_within_size_limit_is_accepted() {
948 let (writer, mut reader) =
950 new_in_memory_queue_with_options(16, DEFAULT_MAX_EVENT_SIZE, DEFAULT_WRITE_TIMEOUT);
951
952 let event = make_status_event("t1", TaskState::Working);
953 writer
954 .write(event)
955 .await
956 .expect("event within size limit should be accepted");
957 drop(writer);
958
959 let r = reader.read().await;
960 assert!(r.is_some(), "reader should receive the event");
961 }
962
963 #[tokio::test]
976 async fn a_full_persistence_channel_fails_the_write_within_the_timeout() {
977 let timeout = std::time::Duration::from_millis(150);
978 let (writer, _sse, _held_receiver) =
979 new_in_memory_queue_with_persistence(1, DEFAULT_MAX_EVENT_SIZE, timeout);
980
981 let mut accepted = 0_usize;
987 let start = std::time::Instant::now();
988 loop {
989 let attempt = tokio::time::timeout(
990 timeout * 10,
991 writer.write(make_status_event("t1", TaskState::Working)),
992 )
993 .await
994 .expect("write parked past its own deadline: write_timeout is not being applied");
995 match attempt {
996 Ok(()) => accepted += 1,
997 Err(e) => {
998 assert!(
999 e.to_string().contains("persistence channel"),
1000 "the error must name the channel that is full, got: {e}"
1001 );
1002 break;
1003 }
1004 }
1005 assert!(
1006 accepted <= 4096,
1007 "write never reported a full channel after {accepted} events"
1008 );
1009 }
1010 assert_eq!(
1011 accepted, 1024,
1012 "the persistence channel's capacity is max(capacity * 16, 1024)"
1013 );
1014 assert!(
1017 start.elapsed() < timeout * 10,
1018 "write should fail after roughly one timeout, took {:?}",
1019 start.elapsed()
1020 );
1021 }
1022
1023 #[tokio::test]
1026 async fn the_failing_write_waits_about_one_write_timeout() {
1027 let timeout = std::time::Duration::from_millis(200);
1028 let (writer, _sse, _held_receiver) =
1029 new_in_memory_queue_with_persistence(1, DEFAULT_MAX_EVENT_SIZE, timeout);
1030 for _ in 0..1024 {
1031 writer
1032 .write(make_status_event("t1", TaskState::Working))
1033 .await
1034 .expect("the channel has room until it is full");
1035 }
1036
1037 let start = std::time::Instant::now();
1038 let err = tokio::time::timeout(
1039 timeout * 10,
1040 writer.write(make_status_event("t1", TaskState::Working)),
1041 )
1042 .await
1043 .expect("write parked past its own deadline: write_timeout is not being applied")
1044 .expect_err("the channel is full");
1045 let waited = start.elapsed();
1046
1047 assert!(waited >= timeout, "returned early after {waited:?}");
1048 assert!(
1049 waited < timeout * 4,
1050 "waited far past the deadline: {waited:?}"
1051 );
1052 assert!(
1053 err.to_string().contains("not draining"),
1054 "the message must say what is wrong, got: {err}"
1055 );
1056 }
1057
1058 #[tokio::test]
1062 async fn a_closed_persistence_channel_is_still_not_an_error() {
1063 let (writer, _sse, persistence_rx) = new_in_memory_queue_with_persistence(
1064 16,
1065 DEFAULT_MAX_EVENT_SIZE,
1066 std::time::Duration::from_millis(50),
1067 );
1068 drop(persistence_rx);
1069
1070 let start = std::time::Instant::now();
1071 tokio::time::timeout(
1072 std::time::Duration::from_secs(5),
1073 writer.write(make_status_event("t1", TaskState::Working)),
1074 )
1075 .await
1076 .expect("a closed channel must be observed at once, never waited out")
1077 .expect("a closed persistence channel must not fail the write");
1078 assert!(
1079 start.elapsed() < std::time::Duration::from_millis(50),
1080 "a closed channel must be detected immediately, not waited out"
1081 );
1082 }
1083
1084 #[tokio::test]
1089 async fn a_dropped_event_moves_the_persistence_error_counter() {
1090 #[derive(Default)]
1091 struct Seen(std::sync::Mutex<Vec<(String, String)>>);
1092 impl crate::metrics::Metrics for Seen {
1093 fn on_persistence_error(&self, operation: &str, error_kind: &str) {
1094 self.0
1095 .lock()
1096 .unwrap()
1097 .push((operation.to_owned(), error_kind.to_owned()));
1098 }
1099 }
1100 let seen = Arc::new(Seen::default());
1101 let (writer, _sse, persistence_rx) = new_in_memory_queue_with_persistence(
1102 16,
1103 DEFAULT_MAX_EVENT_SIZE,
1104 std::time::Duration::from_millis(50),
1105 );
1106 let writer = writer.with_metrics(Arc::clone(&seen) as Arc<dyn crate::metrics::Metrics>);
1107
1108 writer
1109 .write(make_status_event("t1", TaskState::Working))
1110 .await
1111 .expect("an open channel takes the event");
1112 assert!(
1113 seen.0.lock().unwrap().is_empty(),
1114 "nothing is dropped while the processor is alive"
1115 );
1116
1117 drop(persistence_rx);
1118 writer
1119 .write(make_status_event("t1", TaskState::Working))
1120 .await
1121 .expect("a closed channel still does not fail the write");
1122 assert_eq!(
1123 *seen.0.lock().unwrap(),
1124 vec![(
1125 crate::metrics::persistence_operation::QUEUE_HANDOFF.to_owned(),
1126 crate::metrics::queue_handoff_error::CHANNEL_CLOSED.to_owned()
1127 )],
1128 "the dropped event is counted, once, with the bounded labels"
1129 );
1130 }
1131}