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
14type FlushFuture = Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'static>>;
18
19pub 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 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 #[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 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 fn capacity(&self) -> usize {
88 self.config.buffer_size().max(1)
89 }
90
91 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 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 fn poll_flush_inner(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
116 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 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 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 #[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 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 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 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 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 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}