Skip to main content

apalis_diesel_postgres/
sink.rs

1use std::{
2    future::Future,
3    marker::PhantomData,
4    pin::Pin,
5    sync::Mutex,
6    task::{Context, Poll},
7};
8
9use apalis_codec::json::JsonCodec;
10use futures::{FutureExt, Sink};
11
12use crate::{CompactType, Config, Error, PgPool, PgTask, PostgresStorage, queries};
13
14// Wrapped in `Mutex` upstream so `PgSink: Sync` even when the inner future
15// isn't (ntex's `BlockingResult` is `Send`-only). `Mutex::get_mut` keeps the
16// hot path lock-free.
17type FlushFuture = Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'static>>;
18
19/// Buffered task sink used by [`PostgresStorage`].
20pub struct PgSink<Args, Codec = JsonCodec<CompactType>> {
21    pool: PgPool,
22    config: Config,
23    buffer: Vec<PgTask<CompactType>>,
24    flush_future: Mutex<Option<FlushFuture>>,
25    _marker: PhantomData<(Args, Codec)>,
26}
27
28impl<Args, Codec> std::fmt::Debug for PgSink<Args, Codec> {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        f.debug_struct("PgSink")
31            .field("config", &self.config)
32            .field("buffer_len", &self.buffer.len())
33            .finish_non_exhaustive()
34    }
35}
36
37impl<Args, Codec> Clone for PgSink<Args, Codec> {
38    /// Returns a fresh sink sharing the same pool/config; the buffer and any
39    /// in-flight flush are intentionally **not** cloned. Each `PgSink` owns its
40    /// pipeline state: cloning a sink that holds buffered tasks would either
41    /// silently duplicate (double-insert) or silently drop them on flush. The
42    /// clone starts empty, so callers responsible for pending work should
43    /// flush before cloning.
44    fn clone(&self) -> Self {
45        Self {
46            pool: self.pool.clone(),
47            config: self.config.clone(),
48            buffer: Vec::new(),
49            flush_future: Mutex::new(None),
50            _marker: PhantomData,
51        }
52    }
53}
54
55impl<Args, Codec> PgSink<Args, Codec> {
56    /// Create a sink for the given pool and config.
57    #[must_use]
58    pub fn new(pool: &PgPool, config: &Config) -> Self {
59        Self {
60            pool: pool.clone(),
61            config: config.clone(),
62            buffer: Vec::new(),
63            flush_future: Mutex::new(None),
64            _marker: PhantomData,
65        }
66    }
67}
68
69impl<Args, Codec> PgSink<Args, Codec> {
70    /// Re-type the sink's `Args`/`Codec` markers while carrying over the
71    /// buffered tasks and any in-flight flush. The buffer stores
72    /// codec-independent `PgTask<CompactType>` values, so a codec swap via
73    /// [`crate::PostgresStorage::with_codec`] must not silently drop pending
74    /// work — only the phantom markers change.
75    pub(crate) fn retype<NewArgs, NewCodec>(self) -> PgSink<NewArgs, NewCodec> {
76        PgSink {
77            pool: self.pool,
78            config: self.config,
79            buffer: self.buffer,
80            flush_future: self.flush_future,
81            _marker: PhantomData,
82        }
83    }
84
85    /// Buffer capacity from the underlying config (clamped to ≥1 so a
86    /// misconfigured `buffer_size(0)` does not deadlock the sink).
87    fn capacity(&self) -> usize {
88        self.config.buffer_size().max(1)
89    }
90
91    /// Whether `poll_ready` must drive a flush before accepting more work —
92    /// either a flush is already in flight, or the buffer is at capacity.
93    fn needs_flush_before_ready(&mut self) -> bool {
94        self.flush_future
95            .get_mut()
96            .expect("flush_future mutex poisoned")
97            .is_some()
98            || self.buffer.len() >= self.capacity()
99    }
100
101    /// Try to enqueue a single task into the buffer, returning
102    /// `Error::SinkBufferFull` when capacity has been reached.
103    fn try_push(&mut self, item: PgTask<CompactType>) -> Result<(), Error> {
104        let cap = self.capacity();
105        if self.buffer.len() >= cap {
106            return Err(Error::SinkBufferFull(cap));
107        }
108        self.buffer.push(item);
109        Ok(())
110    }
111
112    /// Drive the buffered batch toward completion. Starts a new flush future
113    /// when none is in flight and the buffer is non-empty; otherwise polls the
114    /// existing future and clears it once it resolves.
115    fn poll_flush_inner(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
116        // `&mut self` makes `Mutex::get_mut` infallible-by-borrow — no lock
117        // acquisition, just unique-borrow projection. The mutex exists purely
118        // to satisfy `PgSink: Sync` when the inner future is not `Sync` (ntex).
119        let flush_future = self
120            .flush_future
121            .get_mut()
122            .expect("flush_future mutex poisoned");
123
124        if flush_future.is_none() && self.buffer.is_empty() {
125            return Poll::Ready(Ok(()));
126        }
127
128        if flush_future.is_none() {
129            let pool = self.pool.clone();
130            let config = self.config.clone();
131            let buffer = std::mem::take(&mut self.buffer);
132            *flush_future = Some(Box::pin(queries::push_tasks(pool, config, buffer)));
133        }
134
135        let Some(future) = flush_future.as_mut() else {
136            return Poll::Ready(Ok(()));
137        };
138
139        match future.poll_unpin(cx) {
140            Poll::Pending => Poll::Pending,
141            Poll::Ready(result) => {
142                *flush_future = None;
143                Poll::Ready(result)
144            }
145        }
146    }
147}
148
149impl<Args, Encode, Fetcher> Sink<PgTask<CompactType>> for PostgresStorage<Args, Encode, Fetcher>
150where
151    Fetcher: Unpin,
152{
153    type Error = Error;
154
155    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
156        let this = self.get_mut();
157        if this.sink.needs_flush_before_ready() {
158            this.sink.poll_flush_inner(cx)
159        } else {
160            Poll::Ready(Ok(()))
161        }
162    }
163
164    fn start_send(self: Pin<&mut Self>, item: PgTask<CompactType>) -> Result<(), Self::Error> {
165        self.get_mut().sink.try_push(item)
166    }
167
168    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
169        self.get_mut().sink.poll_flush_inner(cx)
170    }
171
172    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
173        self.poll_flush(cx)
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use std::{
180        pin::Pin,
181        task::{Context, Poll},
182    };
183
184    use diesel::{
185        PgConnection,
186        r2d2::{ConnectionManager, Pool},
187    };
188    use futures::{Sink, future, task::noop_waker_ref};
189    use lets_expect::{AssertionError, AssertionResult, *};
190
191    use super::*;
192
193    fn unchecked_pool() -> PgPool {
194        let manager = ConnectionManager::<PgConnection>::new("postgres://127.0.0.1:1/not-used");
195        Pool::builder()
196            .max_size(1)
197            .connection_timeout(std::time::Duration::from_millis(10))
198            .build_unchecked(manager)
199    }
200
201    fn task() -> PgTask<CompactType> {
202        PgTask::new(b"payload".to_vec())
203    }
204
205    fn sink(buffer_size: usize) -> PgSink<Vec<u8>> {
206        PgSink::new(
207            &unchecked_pool(),
208            &Config::new("sink-unit").set_buffer_size(buffer_size),
209        )
210    }
211
212    fn storage(buffer_size: usize) -> PostgresStorage<Vec<u8>> {
213        let pool = unchecked_pool();
214        let config = Config::new("sink-unit").set_buffer_size(buffer_size);
215        PostgresStorage::<Vec<u8>>::new_with_config(&pool, &config)
216    }
217
218    /// `start_send_via_storage` exercises the public `Sink` impl. The returned
219    /// `len` is the buffer length after the final send (only set on success).
220    fn start_send_via_storage(buffer_size: usize, existing_items: usize) -> Result<usize, Error> {
221        let mut storage = storage(buffer_size);
222        for _ in 0..existing_items {
223            storage.sink.buffer.push(task());
224        }
225        Pin::new(&mut storage).start_send(task())?;
226        Ok(storage.sink.buffer.len())
227    }
228
229    fn poll_ready_via_storage(
230        buffer_size: usize,
231        existing_items: usize,
232    ) -> Poll<Result<(), Error>> {
233        let mut storage = storage(buffer_size);
234        for _ in 0..existing_items {
235            storage.sink.buffer.push(task());
236        }
237        let mut cx = Context::from_waker(noop_waker_ref());
238        Pin::new(&mut storage).poll_ready(&mut cx)
239    }
240
241    struct ReadyObservation {
242        poll: Poll<Result<(), Error>>,
243        buffer_len: usize,
244        has_flush_future: bool,
245    }
246
247    fn poll_ready_in_flight() -> ReadyObservation {
248        let mut storage = storage(2);
249        storage.sink.flush_future =
250            Mutex::new(Some(Box::pin(future::pending::<Result<(), Error>>())));
251        let mut cx = Context::from_waker(noop_waker_ref());
252        let poll = Pin::new(&mut storage).poll_ready(&mut cx);
253        let has_flush_future = storage
254            .sink
255            .flush_future
256            .get_mut()
257            .expect("flush_future mutex poisoned")
258            .is_some();
259        ReadyObservation {
260            poll,
261            buffer_len: storage.sink.buffer.len(),
262            has_flush_future,
263        }
264    }
265
266    /// `poll_flush_observation` captures the state of `poll_flush_sink` after a
267    /// single poll: the poll result and whether the in-flight future was cleared.
268    struct FlushObservation {
269        poll: Poll<Result<(), Error>>,
270        future_cleared: bool,
271        buffer_len: usize,
272    }
273
274    fn poll_flush_sink_with_state(
275        buffer_size: usize,
276        buffered: usize,
277        future: Option<FlushFuture>,
278    ) -> FlushObservation {
279        let mut sink = sink(buffer_size);
280        for _ in 0..buffered {
281            sink.buffer.push(task());
282        }
283        sink.flush_future = Mutex::new(future);
284        let mut cx = Context::from_waker(noop_waker_ref());
285        let poll = sink.poll_flush_inner(&mut cx);
286        let future_cleared = sink
287            .flush_future
288            .get_mut()
289            .expect("flush_future mutex poisoned")
290            .is_none();
291        FlushObservation {
292            poll,
293            future_cleared,
294            buffer_len: sink.buffer.len(),
295        }
296    }
297
298    fn poll_flush_idle() -> FlushObservation {
299        poll_flush_sink_with_state(1, 0, None)
300    }
301
302    fn poll_flush_in_flight_ready(result: Result<(), Error>) -> FlushObservation {
303        poll_flush_sink_with_state(1, 0, Some(Box::pin(future::ready(result))))
304    }
305
306    fn poll_flush_in_flight_pending() -> FlushObservation {
307        poll_flush_sink_with_state(1, 0, Some(Box::pin(future::pending())))
308    }
309
310    /// `poll_flush_creates_future` exercises the `flush_future.is_none() &&
311    /// !buffer.is_empty()` branch: the function builds a new flush future from
312    /// the buffer and immediately polls it. The buffer is drained into the
313    /// future regardless of the poll outcome. Because `push_tasks` runs the
314    /// blocking diesel work via `spawn_blocking`, the first poll typically
315    /// returns `Poll::Pending` (the connection attempt is still in flight); it
316    /// only later resolves to `Ready(Err(...))` once the unreachable pool's
317    /// connect times out. Either observation confirms the drain happened.
318    #[cfg_attr(not(feature = "tokio"), allow(dead_code))]
319    fn poll_flush_creates_future() -> FlushObservation {
320        poll_flush_sink_with_state(2, 1, None)
321    }
322
323    fn poll_close_via_storage(buffered: usize) -> Poll<Result<(), Error>> {
324        let mut storage = storage(2);
325        for _ in 0..buffered {
326            storage.sink.buffer.push(task());
327        }
328        let mut cx = Context::from_waker(noop_waker_ref());
329        Pin::new(&mut storage).poll_close(&mut cx)
330    }
331
332    fn cloned_sink_buffer_len(buffered_items: usize) -> usize {
333        let mut sink = sink(3);
334        for _ in 0..buffered_items {
335            sink.buffer.push(task());
336        }
337        sink.clone().buffer.len()
338    }
339
340    fn cloned_sink_state_drops_flush_future() -> bool {
341        let mut sink = sink(3);
342        sink.buffer.push(task());
343        sink.flush_future = Mutex::new(Some(Box::pin(future::pending::<Result<(), Error>>())));
344        sink.clone()
345            .flush_future
346            .get_mut()
347            .expect("flush_future mutex poisoned")
348            .is_none()
349    }
350
351    fn cloned_sink_buffer_size(buffer_size: usize) -> usize {
352        sink(buffer_size).clone().config.buffer_size()
353    }
354
355    /// State of the sink after `PostgresStorage::with_codec` re-types the
356    /// storage. The buffer holds already-encoded `PgTask<CompactType>`
357    /// values, so — unlike `clone`, which deliberately starts empty — a codec
358    /// swap must carry both the buffer and any in-flight flush over.
359    struct RetypedObservation {
360        buffer_len: usize,
361        kept_in_flight_flush: bool,
362        first_payload: Option<CompactType>,
363    }
364
365    const RETYPE_SENTINEL: &[u8] = b"retype-sentinel";
366
367    fn with_codec_observation(buffered_items: usize, flush_in_flight: bool) -> RetypedObservation {
368        let mut storage = storage(3);
369        for _ in 0..buffered_items {
370            storage
371                .sink
372                .buffer
373                .push(PgTask::new(RETYPE_SENTINEL.to_vec()));
374        }
375        if flush_in_flight {
376            storage.sink.flush_future =
377                Mutex::new(Some(Box::pin(future::pending::<Result<(), Error>>())));
378        }
379        let mut retyped = storage.with_codec::<()>();
380        let kept_in_flight_flush = retyped
381            .sink
382            .flush_future
383            .get_mut()
384            .expect("flush_future mutex poisoned")
385            .is_some();
386        RetypedObservation {
387            buffer_len: retyped.sink.buffer.len(),
388            kept_in_flight_flush,
389            first_payload: retyped.sink.buffer.first().map(|t| t.args.clone()),
390        }
391    }
392
393    fn carried_buffer_len(expected: usize) -> impl Fn(&RetypedObservation) -> AssertionResult {
394        move |obs| {
395            if obs.buffer_len == expected {
396                Ok(())
397            } else {
398                Err(AssertionError::new(vec![format!(
399                    "expected {expected} buffered task(s) to survive with_codec, got {}",
400                    obs.buffer_len
401                )]))
402            }
403        }
404    }
405
406    fn kept_the_in_flight_flush(obs: &RetypedObservation) -> AssertionResult {
407        if obs.kept_in_flight_flush {
408            Ok(())
409        } else {
410            Err(AssertionError::new(vec![
411                "expected with_codec to carry the in-flight flush over, but it was dropped"
412                    .to_owned(),
413            ]))
414        }
415    }
416
417    fn carried_the_exact_task_bytes(obs: &RetypedObservation) -> AssertionResult {
418        match obs.first_payload.as_deref() {
419            Some(RETYPE_SENTINEL) => Ok(()),
420            other => Err(AssertionError::new(vec![format!(
421                "expected the buffered task's payload bytes to survive with_codec verbatim, got {other:?}"
422            )])),
423        }
424    }
425
426    fn sink_debug(buffered_items: usize) -> String {
427        let mut sink = sink(3);
428        for _ in 0..buffered_items {
429            sink.buffer.push(task());
430        }
431        format!("{sink:?}")
432    }
433
434    fn sink_buffer_full_at(expected: usize) -> impl Fn(&Result<usize, Error>) -> AssertionResult {
435        move |result| match result {
436            Err(Error::SinkBufferFull(c)) if *c == expected => Ok(()),
437            other => Err(AssertionError::new(vec![format!(
438                "expected sink buffer full at capacity {expected}, got {other:?}"
439            )])),
440        }
441    }
442
443    fn poll_ready_ok(result: &Poll<Result<(), Error>>) -> AssertionResult {
444        match result {
445            Poll::Ready(Ok(())) => Ok(()),
446            other => Err(AssertionError::new(vec![format!(
447                "expected ready ok, got {other:?}"
448            )])),
449        }
450    }
451
452    #[cfg_attr(not(feature = "tokio"), allow(dead_code))]
453    fn poll_started_flush(result: &Poll<Result<(), Error>>) -> AssertionResult {
454        match result {
455            Poll::Pending | Poll::Ready(Err(_)) => Ok(()),
456            other => Err(AssertionError::new(vec![format!(
457                "expected backpressure to start flushing, got {other:?}"
458            )])),
459        }
460    }
461
462    fn observation_is_idle_ok(obs: &FlushObservation) -> AssertionResult {
463        match (&obs.poll, obs.future_cleared, obs.buffer_len) {
464            (Poll::Ready(Ok(())), true, 0) => Ok(()),
465            other => Err(AssertionError::new(vec![format!(
466                "expected idle Ready(Ok), got {other:?}"
467            )])),
468        }
469    }
470
471    fn observation_is_ready_ok_and_cleared(obs: &FlushObservation) -> AssertionResult {
472        match (&obs.poll, obs.future_cleared) {
473            (Poll::Ready(Ok(())), true) => Ok(()),
474            other => Err(AssertionError::new(vec![format!(
475                "expected Ready(Ok) with cleared future, got {other:?}"
476            )])),
477        }
478    }
479
480    fn observation_is_ready_err_and_cleared(obs: &FlushObservation) -> AssertionResult {
481        match (&obs.poll, obs.future_cleared) {
482            (Poll::Ready(Err(_)), true) => Ok(()),
483            other => Err(AssertionError::new(vec![format!(
484                "expected Ready(Err) with cleared future, got {other:?}"
485            )])),
486        }
487    }
488
489    fn observation_stays_pending(obs: &FlushObservation) -> AssertionResult {
490        match (&obs.poll, obs.future_cleared) {
491            (Poll::Pending, false) => Ok(()),
492            other => Err(AssertionError::new(vec![format!(
493                "expected Pending with future retained, got {other:?}"
494            )])),
495        }
496    }
497
498    #[cfg_attr(not(feature = "tokio"), allow(dead_code))]
499    fn observation_drained_buffer_into_future(obs: &FlushObservation) -> AssertionResult {
500        if obs.buffer_len != 0 {
501            return Err(AssertionError::new(vec![format!(
502                "expected buffer to be drained into the flush future, got {} items",
503                obs.buffer_len
504            )]));
505        }
506        // The flush future is created from the buffer. Either it is still
507        // running (Pending + future retained) or it has resolved (Ready +
508        // future cleared). Both observations confirm the drain happened; we
509        // reject the inconsistent combinations explicitly so the test cannot
510        // pass with stale state.
511        match (&obs.poll, obs.future_cleared) {
512            (Poll::Pending, false) => Ok(()),
513            (Poll::Ready(_), true) => Ok(()),
514            (Poll::Pending, true) => Err(AssertionError::new(vec![
515                "flush returned Pending but the future was cleared".to_owned(),
516            ])),
517            (Poll::Ready(_), false) => Err(AssertionError::new(vec![
518                "flush returned Ready but the future was retained".to_owned(),
519            ])),
520        }
521    }
522
523    fn keeps_in_flight_flush(observation: &ReadyObservation) -> AssertionResult {
524        match (
525            &observation.poll,
526            observation.buffer_len,
527            observation.has_flush_future,
528        ) {
529            (Poll::Pending, 0, true) => Ok(()),
530            other => Err(AssertionError::new(vec![format!(
531                "expected pending in-flight flush, got {other:?}"
532            )])),
533        }
534    }
535
536    fn debug_mentions_public_fields(result: &String) -> AssertionResult {
537        if result.contains("pool") {
538            return Err(AssertionError::new(vec![format!(
539                "debug output unexpectedly exposes the pool, got {result}"
540            )]));
541        }
542        if result.contains("PgSink") && result.contains("config") && result.contains("buffer_len") {
543            Ok(())
544        } else {
545            Err(AssertionError::new(vec![format!(
546                "expected sink debug output with public fields, got {result}"
547            )]))
548        }
549    }
550
551    lets_expect! {
552        expect(start_send_via_storage(buffer_size, existing_items)) {
553            let buffer_size = 2;
554            let existing_items = 0;
555
556            when buffer_has_room_below_capacity {
557                to buffers_the_task { be_ok_and equal(1) }
558            }
559
560            when buffer_has_exactly_one_slot_left {
561                // len == capacity - 1: the boundary push that exactly fills the
562                // buffer must still be accepted (pins the `>= cap` guard, not
563                // `> cap` or `>= cap - 1`).
564                let buffer_size = 2;
565                let existing_items = 1;
566                to accepts_the_final_task_and_fills_the_buffer { be_ok_and equal(2) }
567            }
568
569            when buffer_is_at_capacity_already {
570                let buffer_size = 1;
571                let existing_items = 1;
572                to rejects_the_send_carrying_the_effective_capacity { sink_buffer_full_at(1) }
573            }
574
575            when buffer_is_at_a_non_minimum_capacity_already {
576                let buffer_size = 2;
577                let existing_items = 2;
578                to rejects_the_send_carrying_the_effective_capacity { sink_buffer_full_at(2) }
579            }
580
581            when configured_capacity_is_zero_and_minimum_one_is_full {
582                let buffer_size = 0;
583                let existing_items = 1;
584                to rejects_the_send_via_the_minimum_capacity { sink_buffer_full_at(1) }
585            }
586        }
587
588        expect(poll_ready_via_storage(buffer_size, existing_items)) {
589            let buffer_size = 2;
590            let existing_items = 0;
591
592            when buffer_is_below_capacity_and_no_flush_is_in_flight {
593                to returns_ready_without_flushing { poll_ready_ok }
594            }
595
596            when buffer_has_exactly_one_slot_left {
597                // len == capacity - 1 with no in-flight flush: one free slot
598                // remains, so there is no backpressure and poll_ready must
599                // return Ready(Ok) without driving a flush (pins the `>= cap`
600                // guard against a `>= cap - 1` off-by-one).
601                let buffer_size = 2;
602                let existing_items = 1;
603                to returns_ready_without_flushing { poll_ready_ok }
604            }
605        }
606
607        expect(poll_ready_in_flight()) {
608            when an_earlier_flush_is_still_in_flight {
609                to waits_for_the_flush_to_complete { keeps_in_flight_flush }
610            }
611        }
612
613        expect(poll_flush_idle()) {
614            when there_is_neither_a_pending_flush_nor_buffered_work {
615                to completes_immediately_without_touching_the_database {
616                    observation_is_idle_ok
617                }
618            }
619        }
620
621        expect(poll_flush_in_flight_ready(result)) {
622            let result = Ok(());
623
624            when the_in_flight_flush_resolves_successfully {
625                to returns_ready_ok_and_clears_the_future {
626                    observation_is_ready_ok_and_cleared
627                }
628            }
629
630            when the_in_flight_flush_resolves_with_an_error {
631                let result = Err(Error::SinkBufferFull(1));
632                to surfaces_the_error_and_clears_the_future {
633                    observation_is_ready_err_and_cleared
634                }
635            }
636        }
637
638        expect(poll_flush_in_flight_pending()) {
639            when the_in_flight_flush_is_still_pending {
640                to stays_pending_and_keeps_the_future {
641                    observation_stays_pending
642                }
643            }
644        }
645
646        expect(poll_close_via_storage(buffered)) {
647            let buffered = 0;
648
649            when the_sink_is_already_drained {
650                to delegates_to_flush_and_completes { poll_ready_ok }
651            }
652        }
653
654        expect(cloned_sink_buffer_len(buffered_items)) {
655            let buffered_items = 2;
656
657            when the_original_sink_has_buffered_tasks {
658                to starts_the_clone_with_an_empty_buffer { equal(0) }
659            }
660        }
661
662        expect(cloned_sink_state_drops_flush_future()) {
663            when the_original_sink_has_an_in_flight_flush {
664                to does_not_share_the_in_flight_flush_future { equal(true) }
665            }
666        }
667
668        expect(cloned_sink_buffer_size(buffer_size)) {
669            let buffer_size = 4;
670
671            when the_original_sink_has_custom_capacity {
672                to keeps_the_capacity_configuration { equal(4) }
673            }
674        }
675
676        expect(sink_debug(buffered_items)) {
677            let buffered_items = 2;
678
679            when the_sink_has_buffered_items {
680                to describes_the_sink_without_exposing_the_pool {
681                    debug_mentions_public_fields
682                }
683            }
684        }
685
686        expect(with_codec_observation(buffered_items, flush_in_flight)) {
687            let buffered_items = 1;
688            let flush_in_flight = false;
689
690            when the_sink_holds_buffered_tasks {
691                to carries_the_buffer_into_the_retyped_storage { carried_buffer_len(1) }
692                to carries_the_exact_task_bytes { carried_the_exact_task_bytes }
693            }
694
695            when a_flush_is_in_flight {
696                let flush_in_flight = true;
697                to keeps_the_in_flight_flush { kept_the_in_flight_flush }
698            }
699
700            when the_buffer_already_drained_into_an_in_flight_flush {
701                // Realistic mid-flush state: `poll_flush_inner` has taken the
702                // buffer into the future, nothing is left behind it.
703                let buffered_items = 0;
704                let flush_in_flight = true;
705                to keeps_the_in_flight_flush { kept_the_in_flight_flush }
706                to has_no_buffered_tasks_left { carried_buffer_len(0) }
707            }
708
709            when the_sink_is_empty {
710                let buffered_items = 0;
711                to starts_the_retyped_storage_with_an_empty_buffer { carried_buffer_len(0) }
712            }
713        }
714    }
715
716    #[cfg(feature = "tokio")]
717    mod tokio_tests {
718        use super::*;
719
720        lets_expect! { #tokio_test
721            expect(poll_ready_via_storage(buffer_size, existing_items)) {
722                let buffer_size = 1;
723                let existing_items = 1;
724
725                when buffer_is_at_capacity_without_a_flush_in_flight {
726                    to starts_flushing_before_accepting_more_work { poll_started_flush }
727                }
728            }
729        }
730
731        lets_expect! { #tokio_test
732            expect(poll_flush_sink_with_state(buffer_size, buffered, None).poll) {
733                let buffer_size = 2;
734                let buffered = 1;
735
736                when poll_flush_runs_on_a_real_runtime_with_buffered_work {
737                    to starts_flushing_against_the_unreachable_pool { poll_started_flush }
738                }
739            }
740        }
741
742        lets_expect! { #tokio_test
743            expect(poll_flush_creates_future()) {
744                when there_is_no_in_flight_flush_but_the_buffer_has_work {
745                    to drains_the_buffer_into_a_new_flush_future {
746                        observation_drained_buffer_into_future
747                    }
748                }
749            }
750
751            expect(poll_close_via_storage(1)) {
752                when there_is_buffered_work_to_flush_before_closing {
753                    to starts_flushing_the_buffered_work_before_completing { poll_started_flush }
754                }
755            }
756        }
757    }
758}