1use std::{collections::VecDeque, sync::Arc};
12
13use tokio::sync::Notify;
14
15use crate::{
16 codec::{BackendMessage, FrontendMessage},
17 demux::{Demux, SessionItem},
18 grammar::backend,
19};
20
21#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
23pub struct OperationId(u64);
24
25#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
27pub struct NoPipeline;
28
29#[derive(Clone, Copy, Debug, Eq, PartialEq)]
31pub struct BoundedPipeline {
32 max_operations: usize,
33}
34
35impl BoundedPipeline {
36 pub fn new(max_operations: usize) -> Result<Self, PipelineConfigError> {
42 if max_operations == 0 {
43 return Err(PipelineConfigError);
44 }
45 Ok(Self { max_operations })
46 }
47
48 #[must_use]
50 pub const fn max_operations(self) -> usize {
51 self.max_operations
52 }
53}
54
55#[derive(Clone, Copy, Debug, Eq, PartialEq)]
57pub struct PipelineConfigError;
58
59impl std::fmt::Display for PipelineConfigError {
60 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61 formatter.write_str("pipeline operation limit must be non-zero")
62 }
63}
64
65impl std::error::Error for PipelineConfigError {}
66
67mod private {
68 pub trait Sealed {}
69}
70
71pub trait PipelinePolicy: private::Sealed + Copy {
73 fn operation_limit(self) -> usize;
75}
76
77impl private::Sealed for NoPipeline {}
78impl PipelinePolicy for NoPipeline {
79 fn operation_limit(self) -> usize {
80 1
81 }
82}
83
84impl private::Sealed for BoundedPipeline {}
85impl PipelinePolicy for BoundedPipeline {
86 fn operation_limit(self) -> usize {
87 self.max_operations
88 }
89}
90
91#[derive(Clone, Copy, Debug, Eq, PartialEq)]
93pub enum FrontendHandling {
94 Forward,
96 Local,
98}
99
100#[derive(Debug, Eq, PartialEq)]
102pub enum FrontendAction {
103 Forward {
105 id: OperationId,
107 message: FrontendMessage,
109 },
110 Discard {
112 id: OperationId,
114 },
115 Backpressure(FrontendMessage),
117}
118
119#[derive(Debug, Eq, PartialEq)]
121pub enum FrontendAdmission {
122 Immediate(FrontendAction),
124 Waiting(FrontendAction),
126}
127
128impl FrontendAdmission {
129 #[must_use]
131 pub fn into_action(self) -> FrontendAction {
132 match self {
133 Self::Immediate(action) | Self::Waiting(action) => action,
134 }
135 }
136}
137
138#[derive(Debug, Eq, PartialEq)]
140pub enum FrontendProjectionError {
141 Capacity(Box<FrontendMessage>),
143 Illegal {
145 state: PipelineState,
147 message: Box<FrontendMessage>,
149 },
150}
151
152#[derive(Debug, Eq, PartialEq)]
154pub enum BackendAction {
155 Emit(BackendMessage),
157 Deferred(BackendMessage),
159}
160
161#[derive(Debug, Eq, PartialEq)]
163pub struct BackendProjectionError {
164 pub state: PipelineState,
166 pub message: BackendMessage,
168}
169
170#[derive(Clone, Copy, Debug, Eq, PartialEq)]
172pub enum PipelineState {
173 Ready,
175 Extended,
177 ExtendedError,
179 CopyIn,
181 CopyOut,
183 CopyBoth,
185 Terminated,
187}
188
189#[derive(Clone, Copy, Debug, Eq, PartialEq)]
190enum RequestState {
191 Ready,
192 Extended { bound: bool },
193 ExtendedError,
194 CopyIn,
195 CopyOut,
196 CopyBoth,
197 Terminated,
198}
199
200impl RequestState {
201 const fn public(self) -> PipelineState {
202 match self {
203 Self::Ready => PipelineState::Ready,
204 Self::Extended { .. } => PipelineState::Extended,
205 Self::ExtendedError => PipelineState::ExtendedError,
206 Self::CopyIn => PipelineState::CopyIn,
207 Self::CopyOut => PipelineState::CopyOut,
208 Self::CopyBoth => PipelineState::CopyBoth,
209 Self::Terminated => PipelineState::Terminated,
210 }
211 }
212}
213
214#[derive(Clone, Copy, Debug, Eq, PartialEq)]
215enum Origin {
216 Forwarded,
217 Local,
218}
219
220#[derive(Clone, Copy, Debug, Eq, PartialEq)]
221enum OperationKind {
222 Query,
223 FunctionCall,
224 Parse,
225 Bind,
226 Describe,
227 Execute,
228 Close,
229 Flush,
230 Sync,
231 CopyData,
232 CopyDone,
233 CopyFail,
234 Terminate,
235}
236
237#[derive(Clone, Copy, Debug, Eq, PartialEq)]
238struct Operation {
239 id: OperationId,
240 kind: OperationKind,
241 origin: Origin,
242 discarded: bool,
243}
244
245#[derive(Debug)]
247pub struct Pipeline<P = NoPipeline> {
248 policy: P,
249 operations: VecDeque<Operation>,
250 request_state: RequestState,
251 response_state: Option<PipelineState>,
252 next_id: u64,
253 changed: Arc<Notify>,
254}
255
256impl Default for Pipeline<NoPipeline> {
257 fn default() -> Self {
258 Self::new(NoPipeline)
259 }
260}
261
262impl<P: PipelinePolicy> Pipeline<P> {
263 #[must_use]
265 pub fn new(policy: P) -> Self {
266 Self {
267 policy,
268 operations: VecDeque::new(),
269 request_state: RequestState::Ready,
270 response_state: None,
271 next_id: 0,
272 changed: Arc::new(Notify::new()),
273 }
274 }
275
276 #[must_use]
278 pub fn state(&self) -> PipelineState {
279 self.response_state
280 .unwrap_or_else(|| self.request_state.public())
281 }
282
283 #[must_use]
285 pub fn len(&self) -> usize {
286 self.operations.len()
287 }
288
289 #[must_use]
291 pub fn is_empty(&self) -> bool {
292 self.operations.is_empty()
293 }
294
295 pub fn accept_frontend(
305 &mut self,
306 message: FrontendMessage,
307 handling: FrontendHandling,
308 ) -> Result<FrontendAdmission, FrontendProjectionError> {
309 self.remove_inert_heads();
310 if self.operations.len() == self.policy.operation_limit() {
311 return Err(FrontendProjectionError::Capacity(Box::new(message)));
312 }
313
314 let Some((kind, next_state)) = project_frontend(self.request_state, &message) else {
315 return Err(FrontendProjectionError::Illegal {
316 state: self.state(),
317 message: Box::new(message),
318 });
319 };
320 let waiting = !self.operations.is_empty();
321 let id = OperationId(self.next_id);
322 self.next_id = self.next_id.saturating_add(1);
323 self.request_state = next_state;
324 let origin = match handling {
325 FrontendHandling::Forward => Origin::Forwarded,
326 FrontendHandling::Local => Origin::Local,
327 };
328 self.operations.push_back(Operation {
329 id,
330 kind,
331 origin,
332 discarded: matches!(self.request_state, RequestState::ExtendedError)
333 && kind != OperationKind::Sync,
334 });
335 let action = match handling {
336 FrontendHandling::Forward => FrontendAction::Forward { id, message },
337 FrontendHandling::Local => FrontendAction::Discard { id },
338 };
339 Ok(if waiting {
340 FrontendAdmission::Waiting(action)
341 } else {
342 FrontendAdmission::Immediate(action)
343 })
344 }
345
346 pub fn project_frontend(
352 &mut self,
353 message: FrontendMessage,
354 handling: FrontendHandling,
355 ) -> Result<FrontendAdmission, FrontendProjectionError> {
356 self.accept_frontend(message, handling)
357 }
358
359 pub fn frontend_action(
369 &mut self,
370 message: FrontendMessage,
371 handling: FrontendHandling,
372 ) -> Result<FrontendAction, FrontendProjectionError> {
373 match self.accept_frontend(message, handling) {
374 Ok(admission) => Ok(admission.into_action()),
375 Err(FrontendProjectionError::Capacity(message)) => {
376 Ok(FrontendAction::Backpressure(*message))
377 }
378 Err(error @ FrontendProjectionError::Illegal { .. }) => Err(error),
379 }
380 }
381
382 pub fn accept_backend(
388 &mut self,
389 message: BackendMessage,
390 ) -> Result<BackendAction, BackendProjectionError> {
391 self.accept_response(None, message)
392 }
393
394 pub fn accept_session_item(
405 &mut self,
406 item: SessionItem,
407 ) -> Result<BackendAction, BackendProjectionError> {
408 let message = match item {
409 SessionItem::Message(message) => message,
410 SessionItem::ReadyForQuery { status, .. } => BackendMessage::ReadyForQuery(status),
411 SessionItem::CommandComplete { tag, .. } => BackendMessage::CommandComplete(tag),
412 };
413 self.accept_backend(message)
414 }
415
416 pub fn try_emit_local(
425 &mut self,
426 id: OperationId,
427 message: BackendMessage,
428 ) -> Result<BackendAction, BackendProjectionError> {
429 self.accept_response(Some(id), message)
430 }
431
432 pub async fn wait_until_emittable(&self, id: OperationId) {
437 loop {
438 let notified = self.changed.notified();
439 if self
440 .operations
441 .front()
442 .is_some_and(|operation| operation.id == id)
443 {
444 return;
445 }
446 notified.await;
447 }
448 }
449
450 fn accept_response(
451 &mut self,
452 local_id: Option<OperationId>,
453 message: BackendMessage,
454 ) -> Result<BackendAction, BackendProjectionError> {
455 if is_asynchronous(&message) {
456 return Ok(BackendAction::Emit(message));
457 }
458 self.remove_inert_heads();
459 let Some(head) = self.operations.front().copied() else {
460 return Err(BackendProjectionError {
461 state: self.state(),
462 message,
463 });
464 };
465 if let Some(id) = local_id {
466 if head.id != id {
467 return Ok(BackendAction::Deferred(message));
468 }
469 if head.origin != Origin::Local {
470 return Err(BackendProjectionError {
471 state: self.state(),
472 message,
473 });
474 }
475 } else if head.origin == Origin::Local {
476 if self.operations.iter().skip(1).any(|operation| {
477 operation.origin == Origin::Forwarded && response_fits(operation.kind, &message)
478 }) {
479 return Ok(BackendAction::Deferred(message));
480 }
481 return Err(BackendProjectionError {
482 state: self.state(),
483 message,
484 });
485 }
486 if head.discarded || !response_fits(head.kind, &message) {
487 if self
488 .operations
489 .iter()
490 .skip(1)
491 .any(|operation| response_fits(operation.kind, &message))
492 {
493 return Ok(BackendAction::Deferred(message));
494 }
495 return Err(BackendProjectionError {
496 state: self.state(),
497 message,
498 });
499 }
500
501 let terminal = response_is_terminal(head.kind, &message);
502 let error = matches!(message, BackendMessage::ErrorResponse(_));
503 let copy_state = copy_state(&message);
504 if terminal {
505 self.operations.pop_front();
506 if error && is_extended_kind(head.kind) {
507 self.enter_extended_error();
508 }
509 }
510 if let Some(state) = copy_state {
511 self.response_state = Some(state.public());
512 if matches!(state, RequestState::CopyIn | RequestState::CopyBoth) {
513 self.request_state = state;
514 }
515 }
516 if terminal {
517 self.response_state = None;
518 match head.kind {
519 OperationKind::Sync | OperationKind::Query => {
520 self.request_state = RequestState::Ready;
521 }
522 OperationKind::Execute if !error => {
523 self.request_state = RequestState::Extended { bound: true };
524 }
525 _ => {}
526 }
527 }
528 self.remove_inert_heads();
529 self.changed.notify_waiters();
530 Ok(BackendAction::Emit(message))
531 }
532
533 fn enter_extended_error(&mut self) {
534 self.request_state = RequestState::ExtendedError;
535 self.response_state = None;
536 for operation in &mut self.operations {
537 if operation.kind == OperationKind::Sync {
538 break;
539 }
540 operation.discarded = true;
541 }
542 }
543
544 fn remove_inert_heads(&mut self) {
545 let previous_len = self.operations.len();
546 while self.operations.front().is_some_and(|operation| {
547 operation.kind == OperationKind::Flush
548 || operation.kind == OperationKind::CopyData
549 || operation.kind == OperationKind::CopyDone
550 || operation.kind == OperationKind::CopyFail
551 || operation.kind == OperationKind::Terminate
552 || operation.discarded
553 }) {
554 self.operations.pop_front();
555 }
556 if self.operations.len() != previous_len {
557 self.changed.notify_waiters();
558 }
559 }
560}
561
562fn project_frontend(
563 state: RequestState,
564 message: &FrontendMessage,
565) -> Option<(OperationKind, RequestState)> {
566 use FrontendMessage as F;
567 use OperationKind as O;
568 use RequestState as S;
569 let generated_state = match state {
570 S::Ready => backend::RuntimeState::Ready,
571 S::Extended { .. } => backend::RuntimeState::Building,
572 S::ExtendedError => backend::RuntimeState::ExtendedError,
573 S::CopyIn => backend::RuntimeState::ExtendedCopyIn,
574 S::CopyOut => backend::RuntimeState::ExtendedCopyOut,
575 S::CopyBoth => backend::RuntimeState::ExtendedCopyBoth,
576 S::Terminated => backend::RuntimeState::Terminated,
577 };
578 backend::project_external(generated_state, message)?;
579
580 match (state, message) {
581 (S::Ready, F::Query(_)) => Some((O::Query, S::Ready)),
582 (S::Ready, F::FunctionCall(_)) => Some((O::FunctionCall, S::Ready)),
583 (S::Ready, F::Parse(_)) => Some((O::Parse, S::Extended { bound: false })),
584 (S::Ready | S::Extended { .. }, F::Bind(_)) => Some((O::Bind, S::Extended { bound: true })),
585 (S::Ready, F::Describe(_)) => Some((O::Describe, S::Extended { bound: false })),
586 (S::Ready | S::Extended { .. }, F::Execute(_)) => {
587 Some((O::Execute, S::Extended { bound: true }))
588 }
589 (S::Ready, F::Close(_)) => Some((O::Close, S::Extended { bound: false })),
590 (S::Ready, F::Terminate) => Some((O::Terminate, S::Terminated)),
591 (S::Extended { bound }, F::Parse(_)) => Some((O::Parse, S::Extended { bound })),
592 (S::Extended { bound }, F::Describe(_)) => Some((O::Describe, S::Extended { bound })),
593 (S::Extended { bound }, F::Close(_)) => Some((O::Close, S::Extended { bound })),
594 (S::Extended { bound }, F::Flush) => Some((O::Flush, S::Extended { bound })),
595 (S::Extended { .. } | S::ExtendedError, F::Sync) => Some((O::Sync, S::Ready)),
596 (S::ExtendedError, _) => Some((classify_discard(message)?, S::ExtendedError)),
597 (S::CopyIn, F::CopyData(_)) => Some((O::CopyData, S::CopyIn)),
598 (S::CopyIn, F::CopyDone) => Some((O::CopyDone, S::Extended { bound: true })),
599 (S::CopyIn, F::CopyFail(_)) => Some((O::CopyFail, S::ExtendedError)),
600 (S::CopyBoth, F::CopyData(_)) => Some((O::CopyData, S::CopyBoth)),
601 (S::CopyBoth, F::CopyDone) => Some((O::CopyDone, S::CopyBoth)),
602 _ => None,
603 }
604}
605
606fn classify_discard(message: &FrontendMessage) -> Option<OperationKind> {
607 Some(match message {
608 FrontendMessage::Parse(_) => OperationKind::Parse,
609 FrontendMessage::Bind(_) => OperationKind::Bind,
610 FrontendMessage::Describe(_) => OperationKind::Describe,
611 FrontendMessage::Execute(_) => OperationKind::Execute,
612 FrontendMessage::Close(_) => OperationKind::Close,
613 FrontendMessage::Flush => OperationKind::Flush,
614 FrontendMessage::Query(_) => OperationKind::Query,
615 FrontendMessage::FunctionCall(_) => OperationKind::FunctionCall,
616 FrontendMessage::CopyData(_) => OperationKind::CopyData,
617 FrontendMessage::CopyDone => OperationKind::CopyDone,
618 FrontendMessage::CopyFail(_) => OperationKind::CopyFail,
619 FrontendMessage::Terminate => OperationKind::Terminate,
620 FrontendMessage::PasswordResponse(_) => return None,
621 FrontendMessage::Sync => unreachable!("Sync is classified before discard"),
622 })
623}
624
625fn response_fits(kind: OperationKind, message: &BackendMessage) -> bool {
626 use BackendMessage as B;
627 use OperationKind as O;
628 match kind {
629 O::Query => matches!(
630 message,
631 B::RowDescription(_)
632 | B::DataRow(_)
633 | B::CommandComplete(_)
634 | B::EmptyQueryResponse
635 | B::CopyInResponse(_)
636 | B::CopyOutResponse(_)
637 | B::CopyBothResponse(_)
638 | B::CopyData(_)
639 | B::CopyDone
640 | B::ErrorResponse(_)
641 | B::ReadyForQuery(_)
642 ),
643 O::FunctionCall => matches!(
644 message,
645 B::FunctionCallResponse(_) | B::ErrorResponse(_) | B::ReadyForQuery(_)
646 ),
647 O::Parse => matches!(message, B::ParseComplete | B::ErrorResponse(_)),
648 O::Bind => matches!(message, B::BindComplete | B::ErrorResponse(_)),
649 O::Describe => matches!(
650 message,
651 B::ParameterDescription(_) | B::RowDescription(_) | B::NoData | B::ErrorResponse(_)
652 ),
653 O::Execute => matches!(
654 message,
655 B::RowDescription(_)
656 | B::DataRow(_)
657 | B::EmptyQueryResponse
658 | B::CommandComplete(_)
659 | B::PortalSuspended
660 | B::CopyInResponse(_)
661 | B::CopyOutResponse(_)
662 | B::CopyBothResponse(_)
663 | B::CopyData(_)
664 | B::CopyDone
665 | B::ErrorResponse(_)
666 ),
667 O::Close => matches!(message, B::CloseComplete | B::ErrorResponse(_)),
668 O::Sync => matches!(message, B::ReadyForQuery(_)),
669 O::CopyDone => matches!(
670 message,
671 B::CopyDone | B::CommandComplete(_) | B::ErrorResponse(_)
672 ),
673 O::CopyFail => matches!(message, B::ErrorResponse(_)),
674 O::Flush | O::CopyData | O::Terminate => false,
675 }
676}
677
678fn response_is_terminal(kind: OperationKind, message: &BackendMessage) -> bool {
679 use BackendMessage as B;
680 use OperationKind as O;
681 match kind {
682 O::Query | O::FunctionCall | O::Sync => matches!(message, B::ReadyForQuery(_)),
683 O::Parse => matches!(message, B::ParseComplete | B::ErrorResponse(_)),
684 O::Bind => matches!(message, B::BindComplete | B::ErrorResponse(_)),
685 O::Describe => matches!(
686 message,
687 B::RowDescription(_) | B::NoData | B::ErrorResponse(_)
688 ),
689 O::Execute => matches!(
690 message,
691 B::CommandComplete(_) | B::PortalSuspended | B::ErrorResponse(_)
692 ),
693 O::Close => matches!(message, B::CloseComplete | B::ErrorResponse(_)),
694 O::CopyDone => matches!(
695 message,
696 B::CopyDone | B::CommandComplete(_) | B::ErrorResponse(_)
697 ),
698 O::CopyFail => matches!(message, B::ErrorResponse(_)),
699 O::Flush | O::CopyData | O::Terminate => true,
700 }
701}
702
703fn is_extended_kind(kind: OperationKind) -> bool {
704 !matches!(
705 kind,
706 OperationKind::Query | OperationKind::FunctionCall | OperationKind::Terminate
707 )
708}
709
710fn copy_state(message: &BackendMessage) -> Option<RequestState> {
711 match message {
712 BackendMessage::CopyInResponse(_) => Some(RequestState::CopyIn),
713 BackendMessage::CopyOutResponse(_) => Some(RequestState::CopyOut),
714 BackendMessage::CopyBothResponse(_) => Some(RequestState::CopyBoth),
715 _ => None,
716 }
717}
718
719fn is_asynchronous(message: &BackendMessage) -> bool {
720 Demux::is_asynchronous(message)
721}