pg_proto/
pipeline.rs

1//! Bounded, payload-free orchestration for proxy request pipelines.
2//!
3//! The ledger in this module records protocol obligations, not wire messages.
4//! Applications retain ownership of decoded messages until [`FrontendAction`] or
5//! [`BackendAction`] tells them to forward, emit, retry, or discard the value.
6//! [`Pipeline::accept_frontend`] and [`Pipeline::accept_backend`] are the canonical
7//! ledger interface. Their typed counterparts additionally dispatch accepted
8//! messages through phase-specific middleware before committing them.
9//! Callers receiving [`crate::demux::SessionItem`] values should first consume
10//! any pooling or attribution evidence they need, then convert the item with
11//! [`crate::demux::SessionItem::into_backend_message`] before backend acceptance.
12
13use std::{collections::VecDeque, convert::Infallible, sync::Arc};
14
15use tokio::sync::Notify;
16
17use crate::{
18    codec::{BackendMessage, FrontendMessage},
19    demux::Demux,
20    grammar::backend,
21    middleware::{
22        AsynchronousBackendMessage, ChainError, MessageMiddleware, Middleware,
23        ReconstructableMessage as _, Then,
24    },
25};
26
27/// Stable identity of one accepted frontend operation.
28///
29/// The identity is opaque: it permits application state to be correlated with
30/// ordered responses without exposing the intermediary's pipeline ledger or
31/// projection controls.
32#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
33pub struct OperationId(u64);
34
35/// Pipeline policy which preserves the historical lock-step behaviour.
36#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
37pub struct NoPipeline;
38
39/// Configuration for a bounded frontend operation pipeline.
40#[derive(Clone, Copy, Debug, Eq, PartialEq)]
41pub struct BoundedPipeline {
42    max_operations: usize,
43}
44
45impl BoundedPipeline {
46    /// Creates a pipeline with a non-zero operation-count limit.
47    ///
48    /// # Errors
49    ///
50    /// Returns an error when `max_operations` is zero.
51    pub fn new(max_operations: usize) -> Result<Self, PipelineConfigError> {
52        if max_operations == 0 {
53            return Err(PipelineConfigError);
54        }
55        Ok(Self { max_operations })
56    }
57
58    /// Returns the maximum number of incomplete operations.
59    #[must_use]
60    pub const fn max_operations(self) -> usize {
61        self.max_operations
62    }
63}
64
65/// A zero operation-count limit is not a usable pipeline configuration.
66#[derive(Clone, Copy, Debug, Eq, PartialEq)]
67pub struct PipelineConfigError;
68
69impl std::fmt::Display for PipelineConfigError {
70    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        formatter.write_str("pipeline operation limit must be non-zero")
72    }
73}
74
75impl std::error::Error for PipelineConfigError {}
76
77mod private {
78    pub trait Sealed {}
79}
80
81/// Sealed configuration accepted by the intermediary builder.
82pub trait PipelinePolicy: private::Sealed + Copy {
83    /// Maximum number of incomplete operation records.
84    fn operation_limit(self) -> usize;
85}
86
87impl private::Sealed for NoPipeline {}
88impl PipelinePolicy for NoPipeline {
89    fn operation_limit(self) -> usize {
90        1
91    }
92}
93
94impl private::Sealed for BoundedPipeline {}
95impl PipelinePolicy for BoundedPipeline {
96    fn operation_limit(self) -> usize {
97        self.max_operations
98    }
99}
100
101/// Whether an accepted frontend operation is locally handled or forwarded.
102#[derive(Clone, Copy, Debug, Eq, PartialEq)]
103pub(crate) enum FrontendHandling {
104    /// Send the returned message to the upstream connection.
105    Forward,
106    /// Do not send the request upstream; application code will synthesize its response.
107    Local,
108}
109
110/// Application action for one successfully projected frontend message.
111#[derive(Debug, Eq, PartialEq)]
112pub(crate) enum FrontendAction {
113    /// Forward the owned message upstream.
114    Forward {
115        /// Accepted operation identity.
116        id: OperationId,
117        /// Original, unretained frontend message.
118        message: FrontendMessage,
119    },
120    /// The operation is locally handled and the message can be discarded.
121    Discard {
122        /// Accepted operation identity.
123        id: OperationId,
124    },
125}
126
127/// Position of a successfully accepted operation.
128#[derive(Debug, Eq, PartialEq)]
129pub(crate) enum FrontendAdmission {
130    /// Nothing earlier prevents this operation's response from being emitted.
131    Immediate(FrontendAction),
132    /// The operation was accepted but an earlier response must be emitted first.
133    Waiting(FrontendAction),
134}
135
136impl FrontendAdmission {
137    /// Returns the application action, discarding only the positional annotation.
138    #[must_use]
139    pub(crate) fn into_action(self) -> FrontendAction {
140        match self {
141            Self::Immediate(action) | Self::Waiting(action) => action,
142        }
143    }
144}
145
146/// Why a frontend message could not be accepted.
147#[derive(Debug, Eq, PartialEq)]
148pub enum FrontendProjectionError {
149    /// The bounded ledger is full; the unchanged message may be retried.
150    Capacity(Box<FrontendMessage>),
151    /// The message is not legal in the projected frontend protocol state.
152    Illegal {
153        /// Projected state at rejection.
154        state: PipelineState,
155        /// Unchanged illegal message.
156        message: Box<FrontendMessage>,
157    },
158}
159
160/// Application action for a backend message.
161#[derive(Debug, Eq, PartialEq)]
162pub(crate) enum BackendAction {
163    /// Emit this owned message to the downstream client now.
164    Emit(BackendMessage),
165    /// An earlier operation must complete; retry this unchanged message later.
166    Deferred(BackendMessage),
167}
168
169/// A backend message was not legal for any outstanding operation.
170#[derive(Debug, Eq, PartialEq)]
171pub struct BackendProjectionError {
172    /// Current response-side state.
173    pub state: PipelineState,
174    /// Unchanged illegal message.
175    pub message: BackendMessage,
176}
177
178#[derive(Debug, Eq, PartialEq)]
179pub(crate) enum BackendSequenceError {
180    Cardinality { expected: usize, actual: usize },
181    Source(BackendMessage),
182    Replacement(BackendMessage),
183    DifferentSpan,
184}
185
186/// Public summary of the pipeline's projected frontend state.
187#[derive(Clone, Copy, Debug, Eq, PartialEq)]
188pub enum PipelineState {
189    /// A simple or extended cycle may begin.
190    Ready,
191    /// Extended-query messages are being accepted.
192    Extended,
193    /// An extended error discards messages through `Sync`.
194    ExtendedError,
195    /// COPY IN accepts frontend data.
196    CopyIn,
197    /// COPY OUT accepts only backend data.
198    CopyOut,
199    /// COPY BOTH accepts data in both directions.
200    CopyBoth,
201    /// The connection has terminated.
202    Terminated,
203}
204
205/// Error returned while dispatching a pipeline message through typed middleware.
206#[derive(Debug)]
207pub(crate) enum PipelineMiddlewareError<MiddlewareError, ProjectionError> {
208    /// Middleware rejected the phase-typed message.
209    Middleware(MiddlewareError),
210    /// The pipeline rejected the original or rewritten message.
211    Projection(ProjectionError),
212}
213
214macro_rules! frontend_pipeline_phases {
215    ($consumer:ident) => {
216        $consumer! {
217            Ready => frontend_ready => backend::ReadyExternalMessage,
218            Building => frontend_building => backend::BuildingExternalMessage,
219            ExtendedError => frontend_extended_error => backend::ExtendedErrorExternalMessage,
220            SimpleCopyIn => frontend_simple_copy_in => backend::SimpleCopyInExternalMessage,
221            ExtendedCopyIn => frontend_extended_copy_in => backend::ExtendedCopyInExternalMessage,
222            SimpleCopyBoth => frontend_simple_copy_both => backend::SimpleCopyBothExternalMessage,
223            ExtendedCopyBoth => frontend_extended_copy_both => backend::ExtendedCopyBothExternalMessage,
224        }
225    };
226}
227
228macro_rules! backend_pipeline_phases {
229    ($consumer:ident) => {
230        $consumer! {
231            Asynchronous => backend_asynchronous => AsynchronousBackendMessage,
232            Simple => backend_simple => backend::SimpleInternalMessage,
233            SimpleError => backend_simple_error => backend::SimpleErrorInternalMessage,
234            ParseResponse => backend_parse_response => backend::ParseResponseInternalMessage,
235            BindResponse => backend_bind_response => backend::BindResponseInternalMessage,
236            DescribeResponse => backend_describe_response => backend::DescribeResponseInternalMessage,
237            ExecuteResponse => backend_execute_response => backend::ExecuteResponseInternalMessage,
238            CloseResponse => backend_close_response => backend::CloseResponseInternalMessage,
239            SyncResponse => backend_sync_response => backend::SyncResponseInternalMessage,
240            FunctionResponse => backend_function_response => backend::FunctionResponseInternalMessage,
241            FunctionReady => backend_function_ready => backend::FunctionReadyInternalMessage,
242            SimpleCopyInDone => backend_simple_copy_in_done => backend::SimpleCopyInDoneInternalMessage,
243            SimpleCopyInFailed => backend_simple_copy_in_failed => backend::SimpleCopyInFailedInternalMessage,
244            SimpleCopyOut => backend_simple_copy_out => backend::SimpleCopyOutInternalMessage,
245            SimpleCopyOutDone => backend_simple_copy_out_done => backend::SimpleCopyOutDoneInternalMessage,
246            SimpleCopyReady => backend_simple_copy_ready => backend::SimpleCopyReadyInternalMessage,
247            ExtendedCopyInDone => backend_extended_copy_in_done => backend::ExtendedCopyInDoneInternalMessage,
248            ExtendedCopyInFailed => backend_extended_copy_in_failed => backend::ExtendedCopyInFailedInternalMessage,
249            ExtendedCopyOut => backend_extended_copy_out => backend::ExtendedCopyOutInternalMessage,
250            ExtendedCopyOutDone => backend_extended_copy_out_done => backend::ExtendedCopyOutDoneInternalMessage,
251            SimpleCopyBoth => backend_simple_copy_both => backend::SimpleCopyBothInternalMessage,
252            SimpleCopyBothClientDone => backend_simple_copy_both_client_done => backend::SimpleCopyBothClientDoneInternalMessage,
253            SimpleCopyBothDone => backend_simple_copy_both_done => backend::SimpleCopyBothDoneInternalMessage,
254            SimpleCopyBothFailed => backend_simple_copy_both_failed => backend::SimpleCopyBothFailedInternalMessage,
255            ExtendedCopyBoth => backend_extended_copy_both => backend::ExtendedCopyBothInternalMessage,
256            ExtendedCopyBothClientDone => backend_extended_copy_both_client_done => backend::ExtendedCopyBothClientDoneInternalMessage,
257            ExtendedCopyBothDone => backend_extended_copy_both_done => backend::ExtendedCopyBothDoneInternalMessage,
258            ExtendedCopyBothFailed => backend_extended_copy_both_failed => backend::ExtendedCopyBothFailedInternalMessage,
259        }
260    };
261}
262
263macro_rules! declare_pipeline_hooks {
264    ($($phase:ident => $method:ident => $message:path),+ $(,)?) => {
265        $(
266            #[doc = concat!("Intercepts backend messages in generated `", stringify!($message), "` phase.")]
267            async fn $method(
268                &mut self,
269                _state: &mut State,
270                message: $message,
271            ) -> Result<$message, Self::Error> {
272                Ok(message)
273            }
274        )+
275    };
276}
277
278/// Async middleware for frontend messages selected from the runtime ledger phase.
279#[allow(async_fn_in_trait)]
280pub(crate) trait FrontendPipelineMiddleware<State> {
281    /// An error which prevents the message from continuing through the pipeline.
282    type Error;
283    frontend_pipeline_phases!(declare_pipeline_hooks);
284}
285
286/// Async middleware for backend messages selected from the runtime ledger phase.
287#[allow(async_fn_in_trait)]
288pub(crate) trait BackendPipelineMiddleware<State> {
289    /// An error which prevents the message from continuing through the pipeline.
290    type Error;
291    backend_pipeline_phases!(declare_pipeline_hooks);
292}
293
294impl<State> FrontendPipelineMiddleware<State> for crate::middleware::Identity {
295    type Error = Infallible;
296}
297
298impl<State> BackendPipelineMiddleware<State> for crate::middleware::Identity {
299    type Error = Infallible;
300}
301
302macro_rules! chained_pipeline_hooks {
303    ($($phase:ident => $method:ident => $message:ty),+ $(,)?) => {
304        $(
305            async fn $method(
306                &mut self,
307                state: &mut State,
308                message: $message,
309            ) -> Result<$message, Self::Error> {
310                let (first, second) = self.parts_mut();
311                let message = first
312                    .$method(state, message)
313                    .await
314                    .map_err(ChainError::First)?;
315                second
316                    .$method(state, message)
317                    .await
318                    .map_err(ChainError::Second)
319            }
320        )+
321    };
322}
323
324impl<State, First, Second> FrontendPipelineMiddleware<State> for Then<First, Second>
325where
326    First: FrontendPipelineMiddleware<State>,
327    Second: FrontendPipelineMiddleware<State>,
328{
329    type Error = ChainError<First::Error, Second::Error>;
330
331    frontend_pipeline_phases!(chained_pipeline_hooks);
332}
333
334impl<State, First, Second> BackendPipelineMiddleware<State> for Then<First, Second>
335where
336    First: BackendPipelineMiddleware<State>,
337    Second: BackendPipelineMiddleware<State>,
338{
339    type Error = ChainError<First::Error, Second::Error>;
340    backend_pipeline_phases!(chained_pipeline_hooks);
341}
342
343/// Adapts direction-wide async middleware to every typed pipeline hook.
344pub(crate) struct PipelineWireAdapter<Handler> {
345    handler: Handler,
346}
347
348impl<Handler> PipelineWireAdapter<Handler> {
349    /// Wraps direction-wide middleware for runtime phase dispatch.
350    pub(crate) const fn new(handler: Handler) -> Self {
351        Self { handler }
352    }
353
354    /// Returns the wrapped direction-wide middleware.
355    pub(crate) fn into_inner(self) -> Handler {
356        self.handler
357    }
358}
359
360/// Failure from direction-wide middleware adapted to typed pipeline dispatch.
361#[derive(Debug)]
362pub(crate) enum FrontendPipelineWireAdapterError<Error> {
363    /// The wrapped middleware rejected a message.
364    Middleware(Error),
365    /// The wrapped middleware returned a frontend message illegal in the selected phase.
366    IllegalFrontend(FrontendMessage),
367}
368
369/// Failure from backend wire middleware adapted to typed pipeline dispatch.
370#[derive(Debug)]
371pub(crate) enum BackendPipelineWireAdapterError<Error> {
372    /// The wrapped middleware rejected a message.
373    Middleware(Error),
374    /// The wrapped middleware returned a message illegal in the selected phase.
375    Illegal(BackendMessage),
376}
377
378macro_rules! pipeline_adapter_frontend_hooks {
379    ($($phase:ident => $method:ident => $message:ty),+ $(,)?) => {
380        $(
381            async fn $method(
382                &mut self,
383                state: &mut State,
384                message: $message,
385            ) -> Result<$message, Self::Error> {
386                let message: FrontendMessage = message.into();
387                let message = self
388                    .handler
389                    .intercept(state, message)
390                    .await
391                    .map_err(FrontendPipelineWireAdapterError::Middleware)?;
392                <$message>::try_from(message)
393                    .map_err(FrontendPipelineWireAdapterError::IllegalFrontend)
394            }
395        )+
396    };
397}
398
399macro_rules! pipeline_adapter_backend_hooks {
400    ($ignored:ident => $async_method:ident => AsynchronousBackendMessage, $($phase:ident => $method:ident => $message:ty),+ $(,)?) => {
401        async fn $async_method(
402            &mut self,
403            state: &mut State,
404            message: AsynchronousBackendMessage,
405        ) -> Result<AsynchronousBackendMessage, Self::Error> {
406            let message = self.handler.intercept(state, message.into_wire()).await
407                .map_err(BackendPipelineWireAdapterError::Middleware)?;
408            AsynchronousBackendMessage::try_from(message)
409                .map_err(BackendPipelineWireAdapterError::Illegal)
410        }
411        $(
412            async fn $method(
413                &mut self,
414                state: &mut State,
415                message: $message,
416            ) -> Result<$message, Self::Error> {
417                let message: BackendMessage = message.into();
418                let message = self
419                    .handler
420                    .intercept(state, message)
421                    .await
422                    .map_err(BackendPipelineWireAdapterError::Middleware)?;
423                <$message>::try_from(message).map_err(BackendPipelineWireAdapterError::Illegal)
424            }
425        )+
426    };
427}
428
429impl<State, Handler> FrontendPipelineMiddleware<State> for PipelineWireAdapter<Handler>
430where
431    Handler: MessageMiddleware<FrontendMessage, State>,
432{
433    type Error = FrontendPipelineWireAdapterError<Handler::Error>;
434    frontend_pipeline_phases!(pipeline_adapter_frontend_hooks);
435}
436
437impl<State, Handler> BackendPipelineMiddleware<State> for PipelineWireAdapter<Handler>
438where
439    Handler: MessageMiddleware<BackendMessage, State>,
440{
441    type Error = BackendPipelineWireAdapterError<Handler::Error>;
442    backend_pipeline_phases!(pipeline_adapter_backend_hooks);
443}
444
445#[derive(Clone, Copy, Debug, Eq, PartialEq)]
446enum RequestState {
447    Ready,
448    Extended { bound: bool },
449    ExtendedError,
450    CopyIn,
451    CopyOut,
452    CopyBoth,
453    Terminated,
454}
455
456impl RequestState {
457    const fn public(self) -> PipelineState {
458        match self {
459            Self::Ready => PipelineState::Ready,
460            Self::Extended { .. } => PipelineState::Extended,
461            Self::ExtendedError => PipelineState::ExtendedError,
462            Self::CopyIn => PipelineState::CopyIn,
463            Self::CopyOut => PipelineState::CopyOut,
464            Self::CopyBoth => PipelineState::CopyBoth,
465            Self::Terminated => PipelineState::Terminated,
466        }
467    }
468}
469
470#[derive(Clone, Copy, Debug, Eq, PartialEq)]
471enum Origin {
472    Forwarded,
473    Local,
474}
475
476#[derive(Clone, Copy, Debug, Eq, PartialEq)]
477enum OperationKind {
478    Query,
479    FunctionCall,
480    Parse,
481    Bind,
482    Describe,
483    Execute,
484    Close,
485    Flush,
486    Sync,
487    CopyData,
488    CopyDone,
489    CopyFail,
490    Terminate,
491}
492
493#[derive(Clone, Copy, Debug, Eq, PartialEq)]
494enum PreparedResponse {
495    Asynchronous,
496    Emit {
497        head: Operation,
498        response_state: backend::RuntimeState,
499    },
500    Deferred,
501    Illegal,
502}
503
504#[derive(Clone, Copy, Debug, Eq, PartialEq)]
505enum FrontendPhase {
506    Ready,
507    Building,
508    ExtendedError,
509    SimpleCopyIn,
510    ExtendedCopyIn,
511    SimpleCopyBoth,
512    ExtendedCopyBoth,
513}
514
515#[derive(Clone, Copy, Debug)]
516struct PreparedFrontend {
517    phase: FrontendPhase,
518    request_state: RequestState,
519}
520
521#[derive(Clone, Copy, Debug, Eq, PartialEq)]
522struct Operation {
523    id: OperationId,
524    kind: OperationKind,
525    origin: Origin,
526    discarded: bool,
527    response_state: backend::RuntimeState,
528}
529
530/// A bounded ledger coordinating independently owned frontend and backend values.
531#[derive(Debug)]
532pub(crate) struct Pipeline<P = NoPipeline> {
533    policy: P,
534    operations: VecDeque<Operation>,
535    request_state: RequestState,
536    response_state: Option<PipelineState>,
537    next_id: u64,
538    changed: Arc<Notify>,
539}
540
541impl Default for Pipeline<NoPipeline> {
542    fn default() -> Self {
543        Self::new(NoPipeline)
544    }
545}
546
547impl<P: PipelinePolicy> Pipeline<P> {
548    /// Returns the identity that will be assigned to the next accepted
549    /// frontend operation.
550    pub(crate) const fn next_operation_id(&self) -> OperationId {
551        OperationId(self.next_id)
552    }
553
554    /// Attributes a backend message to its forwarded frontend operation.
555    /// Asynchronous messages have no operation identity.
556    pub(crate) fn backend_operation_id(&self, message: &BackendMessage) -> Option<OperationId> {
557        if is_asynchronous(message) {
558            return None;
559        }
560        self.operations
561            .iter()
562            .find(|operation| {
563                operation.origin == Origin::Forwarded
564                    && !operation.discarded
565                    && response_fits(**operation, message)
566            })
567            .map(|operation| operation.id)
568    }
569
570    /// Attributes an ordered backend span while projecting each source against
571    /// a private ledger snapshot.
572    pub(crate) fn backend_operation_ids(
573        &self,
574        messages: &[BackendMessage],
575    ) -> Vec<Option<OperationId>> {
576        let mut projection = self.snapshot();
577        messages
578            .iter()
579            .map(|message| {
580                let id = projection.backend_operation_id(message);
581                let _ = projection.accept_backend(message.clone());
582                id
583            })
584            .collect()
585    }
586
587    fn snapshot(&self) -> Self {
588        Self {
589            policy: self.policy,
590            operations: self.operations.clone(),
591            request_state: self.request_state,
592            response_state: self.response_state,
593            next_id: self.next_id,
594            changed: Arc::clone(&self.changed),
595        }
596    }
597
598    pub(crate) fn prepare_backend_replacements(
599        &self,
600        sources: &[BackendMessage],
601        replacements: &[BackendMessage],
602    ) -> Result<Self, BackendSequenceError> {
603        if sources.len() != replacements.len() {
604            return Err(BackendSequenceError::Cardinality {
605                expected: sources.len(),
606                actual: replacements.len(),
607            });
608        }
609        let mut source_projection = self.snapshot();
610        for message in sources.iter().cloned() {
611            match source_projection.accept_backend(message) {
612                Ok(BackendAction::Emit(_)) => {}
613                Ok(BackendAction::Deferred(message))
614                | Err(BackendProjectionError { message, .. }) => {
615                    return Err(BackendSequenceError::Source(message));
616                }
617            }
618        }
619        let mut replacement_projection = self.snapshot();
620        for message in replacements.iter().cloned() {
621            if !message.is_reconstructable() {
622                return Err(BackendSequenceError::Replacement(message));
623            }
624            match replacement_projection.accept_backend(message) {
625                Ok(BackendAction::Emit(_)) => {}
626                Ok(BackendAction::Deferred(message))
627                | Err(BackendProjectionError { message, .. }) => {
628                    return Err(BackendSequenceError::Replacement(message));
629                }
630            }
631        }
632        if source_projection.operations != replacement_projection.operations
633            || source_projection.request_state != replacement_projection.request_state
634            || source_projection.response_state != replacement_projection.response_state
635        {
636            return Err(BackendSequenceError::DifferentSpan);
637        }
638        Ok(replacement_projection)
639    }
640    /// Creates an empty pipeline using `policy`.
641    #[must_use]
642    pub(crate) fn new(policy: P) -> Self {
643        Self {
644            policy,
645            operations: VecDeque::new(),
646            request_state: RequestState::Ready,
647            response_state: None,
648            next_id: 0,
649            changed: Arc::new(Notify::new()),
650        }
651    }
652
653    /// Returns the projected frontend protocol state.
654    #[must_use]
655    pub(crate) fn state(&self) -> PipelineState {
656        self.response_state
657            .unwrap_or_else(|| self.request_state.public())
658    }
659
660    /// Returns the number of incomplete lightweight operation records.
661    #[must_use]
662    pub(crate) fn len(&self) -> usize {
663        self.operations.len()
664    }
665
666    /// Reports whether no operations remain outstanding.
667    #[must_use]
668    pub(crate) fn is_empty(&self) -> bool {
669        self.operations.is_empty()
670    }
671
672    /// Projects and accepts one frontend message without retaining its payload.
673    ///
674    /// Capacity and legality failures return the original owned message. A
675    /// capacity failure does not mutate either projected state or the ledger.
676    ///
677    /// # Errors
678    ///
679    /// Returns the unchanged boxed message when capacity is exhausted or the
680    /// message is illegal in the projected state.
681    pub(crate) fn accept_frontend(
682        &mut self,
683        message: FrontendMessage,
684        handling: FrontendHandling,
685    ) -> Result<FrontendAdmission, FrontendProjectionError> {
686        let prepared = self.prepare_frontend(&message)?;
687        Ok(self.commit_frontend(prepared, message, handling))
688    }
689
690    fn prepare_frontend(
691        &self,
692        message: &FrontendMessage,
693    ) -> Result<PreparedFrontend, FrontendProjectionError> {
694        if self.operations.len() == self.policy.operation_limit()
695            && !matches!(
696                self.request_state,
697                RequestState::CopyIn | RequestState::CopyBoth
698            )
699        {
700            return Err(FrontendProjectionError::Capacity(Box::new(message.clone())));
701        }
702        if project_frontend(self.request_state, message).is_none() {
703            return Err(FrontendProjectionError::Illegal {
704                state: self.state(),
705                message: Box::new(message.clone()),
706            });
707        }
708        Ok(PreparedFrontend {
709            phase: frontend_phase(
710                self.request_state,
711                self.operations.front().map(|operation| operation.kind),
712            ),
713            request_state: self.request_state,
714        })
715    }
716
717    fn commit_frontend(
718        &mut self,
719        prepared: PreparedFrontend,
720        message: FrontendMessage,
721        handling: FrontendHandling,
722    ) -> FrontendAdmission {
723        let (kind, next_state) = classify_frontend(prepared.request_state, &message)
724            .expect("phase-typed frontend replacement has a ledger classification");
725        let waiting = !self.operations.is_empty();
726        let id = OperationId(self.next_id);
727        self.next_id = self.next_id.saturating_add(1);
728        self.request_state = next_state;
729        let origin = match handling {
730            FrontendHandling::Forward => Origin::Forwarded,
731            FrontendHandling::Local => Origin::Local,
732        };
733        if let Some(head) = self.operations.front_mut()
734            && let Some(event) = backend::project_external(head.response_state, &message)
735            && let Some(transition) = backend::transition(head.response_state, event)
736        {
737            head.response_state = transition.target;
738        }
739        self.operations.push_back(Operation {
740            id,
741            kind,
742            origin,
743            discarded: matches!(self.request_state, RequestState::ExtendedError)
744                && kind != OperationKind::Sync,
745            response_state: initial_response_state(kind),
746        });
747        let action = match handling {
748            FrontendHandling::Forward => FrontendAction::Forward { id, message },
749            FrontendHandling::Local => FrontendAction::Discard { id },
750        };
751        let admission = if waiting {
752            FrontendAdmission::Waiting(action)
753        } else {
754            FrontendAdmission::Immediate(action)
755        };
756        self.remove_inert_heads();
757        admission
758    }
759
760    /// Projects, asynchronously intercepts, and accepts one frontend message.
761    ///
762    /// The ledger selects the phase-specific middleware hook at runtime. The
763    /// selected hook can only return a message legal in that same phase.
764    /// Middleware is not invoked when capacity is exhausted.
765    ///
766    /// # Errors
767    ///
768    /// Returns a middleware error, an illegal original or replacement message,
769    /// or the unchanged message when capacity is exhausted.
770    pub(crate) async fn accept_frontend_typed<State, Handler>(
771        &mut self,
772        middleware: &mut Middleware<State, Handler>,
773        message: FrontendMessage,
774        handling: FrontendHandling,
775    ) -> Result<FrontendAdmission, PipelineMiddlewareError<Handler::Error, FrontendProjectionError>>
776    where
777        Handler: FrontendPipelineMiddleware<State>,
778    {
779        let prepared = self
780            .prepare_frontend(&message)
781            .map_err(PipelineMiddlewareError::Projection)?;
782
783        let message = self
784            .intercept_frontend(prepared.phase, middleware, message)
785            .await
786            .map_err(PipelineMiddlewareError::Middleware)?;
787        if !message.is_reconstructable() {
788            return Err(PipelineMiddlewareError::Projection(
789                FrontendProjectionError::Illegal {
790                    state: self.state(),
791                    message: Box::new(message),
792                },
793            ));
794        }
795        Ok(self.commit_frontend(prepared, message, handling))
796    }
797
798    /// Projects one upstream backend message and preserves response order.
799    ///
800    /// # Errors
801    ///
802    /// Returns an unchanged response which cannot belong to any outstanding operation.
803    pub(crate) fn accept_backend(
804        &mut self,
805        message: BackendMessage,
806    ) -> Result<BackendAction, BackendProjectionError> {
807        self.accept_response(None, message)
808    }
809
810    /// Intercepts an emittable backend response through its operation-typed hook.
811    ///
812    /// Responses belonging to a later operation are returned unchanged as
813    /// [`BackendAction::Deferred`] and are intercepted only when retried at the
814    /// response head. Asynchronous messages use their non-advancing hook.
815    ///
816    /// # Errors
817    ///
818    /// Returns a middleware error or an unchanged response which cannot belong
819    /// to any outstanding operation.
820    pub(crate) async fn accept_backend_typed<State, Handler>(
821        &mut self,
822        middleware: &mut Middleware<State, Handler>,
823        message: BackendMessage,
824    ) -> Result<BackendAction, PipelineMiddlewareError<Handler::Error, BackendProjectionError>>
825    where
826        Handler: BackendPipelineMiddleware<State>,
827    {
828        self.accept_response_typed(None, middleware, message).await
829    }
830
831    /// Attempts to register and emit a locally synthesized response.
832    ///
833    /// The message is returned as [`BackendAction::Deferred`] when `id` has not
834    /// reached the head. No backend payload is retained by the ledger.
835    ///
836    /// # Errors
837    ///
838    /// Returns the unchanged message when it is illegal for the named operation.
839    pub(crate) fn try_emit_local(
840        &mut self,
841        id: OperationId,
842        message: BackendMessage,
843    ) -> Result<BackendAction, BackendProjectionError> {
844        self.accept_response(Some(id), message)
845    }
846
847    /// Typed-middleware counterpart to [`Self::try_emit_local`].
848    ///
849    /// Deferred local responses are not intercepted until their operation reaches
850    /// the response head.
851    ///
852    /// # Errors
853    ///
854    /// Returns a middleware error or an illegal response for the named operation.
855    pub(crate) async fn try_emit_local_typed<State, Handler>(
856        &mut self,
857        middleware: &mut Middleware<State, Handler>,
858        id: OperationId,
859        message: BackendMessage,
860    ) -> Result<BackendAction, PipelineMiddlewareError<Handler::Error, BackendProjectionError>>
861    where
862        Handler: BackendPipelineMiddleware<State>,
863    {
864        self.accept_response_typed(Some(id), middleware, message)
865            .await
866    }
867
868    /// Waits until a local operation reaches the response head.
869    ///
870    /// Cancellation is safe: polling this future never reserves or removes a
871    /// ledger entry. The caller should then invoke [`Self::try_emit_local`].
872    pub(crate) async fn wait_until_emittable(&self, id: OperationId) {
873        loop {
874            let notified = self.changed.notified();
875            if self
876                .operations
877                .front()
878                .is_some_and(|operation| operation.id == id)
879            {
880                return;
881            }
882            notified.await;
883        }
884    }
885
886    async fn intercept_frontend<State, Handler>(
887        &self,
888        phase: FrontendPhase,
889        middleware: &mut Middleware<State, Handler>,
890        message: FrontendMessage,
891    ) -> Result<FrontendMessage, Handler::Error>
892    where
893        Handler: FrontendPipelineMiddleware<State>,
894    {
895        let (state, handler) = middleware.parts_mut();
896        macro_rules! dispatch {
897            ($message:expr, $type:path, $handler:ident, $state:ident, $method:ident) => {{
898                let Ok(typed) = <$type>::try_from($message) else {
899                    unreachable!("frontend message was prevalidated for pipeline phase")
900                };
901                $handler.$method($state, typed).await?.into()
902            }};
903        }
904
905        macro_rules! dispatch_catalogue {
906            ($($catalogue_phase:ident => $method:ident => $message_type:path),+ $(,)?) => {
907                match phase {
908                    $(
909                        FrontendPhase::$catalogue_phase =>
910                            dispatch!(message, $message_type, handler, state, $method),
911                    )+
912                }
913            };
914        }
915
916        Ok(frontend_pipeline_phases!(dispatch_catalogue))
917    }
918
919    #[allow(clippy::too_many_lines)]
920    async fn accept_response_typed<State, Handler>(
921        &mut self,
922        local_id: Option<OperationId>,
923        middleware: &mut Middleware<State, Handler>,
924        message: BackendMessage,
925    ) -> Result<BackendAction, PipelineMiddlewareError<Handler::Error, BackendProjectionError>>
926    where
927        Handler: BackendPipelineMiddleware<State>,
928    {
929        let prepared = self.prepare_response(local_id, &message);
930        if matches!(prepared, PreparedResponse::Deferred) {
931            return Ok(BackendAction::Deferred(message));
932        }
933        if matches!(prepared, PreparedResponse::Illegal) {
934            return Err(PipelineMiddlewareError::Projection(
935                BackendProjectionError {
936                    state: self.state(),
937                    message,
938                },
939            ));
940        }
941
942        let (state, handler) = middleware.parts_mut();
943        let message = match prepared {
944            PreparedResponse::Asynchronous => {
945                let Ok(typed) = AsynchronousBackendMessage::try_from(message) else {
946                    unreachable!("asynchronous response was prevalidated")
947                };
948                handler
949                    .backend_asynchronous(state, typed)
950                    .await
951                    .map_err(PipelineMiddlewareError::Middleware)?
952                    .into_wire()
953            }
954            PreparedResponse::Emit { response_state, .. } => {
955                macro_rules! dispatch {
956                    ($message:ty, $method:ident) => {{
957                        let typed = match <$message>::try_from(message) {
958                            Ok(typed) => typed,
959                            Err(message) => {
960                                return Err(PipelineMiddlewareError::Projection(
961                                    BackendProjectionError {
962                                        state: self.state(),
963                                        message,
964                                    },
965                                ));
966                            }
967                        };
968                        handler
969                            .$method(state, typed)
970                            .await
971                            .map_err(PipelineMiddlewareError::Middleware)?
972                            .into_wire()
973                    }};
974                }
975
976                macro_rules! dispatch_catalogue {
977                    ($ignored:ident => $ignored_method:ident => AsynchronousBackendMessage,
978                     $($catalogue_phase:ident => $method:ident => $message_type:path),+ $(,)?) => {
979                        match response_state {
980                            $(
981                                backend::RuntimeState::$catalogue_phase =>
982                                    dispatch!($message_type, $method),
983                            )+
984                            _ => unreachable!("response phase has no backend-selected transition"),
985                        }
986                    };
987                }
988
989                backend_pipeline_phases!(dispatch_catalogue)
990            }
991            PreparedResponse::Deferred | PreparedResponse::Illegal => unreachable!(),
992        };
993
994        if !message.is_reconstructable() {
995            return Err(PipelineMiddlewareError::Projection(
996                BackendProjectionError {
997                    state: self.state(),
998                    message,
999                },
1000            ));
1001        }
1002        self.commit_response(prepared, message)
1003            .map_err(PipelineMiddlewareError::Projection)
1004    }
1005
1006    fn prepare_response(
1007        &self,
1008        local_id: Option<OperationId>,
1009        message: &BackendMessage,
1010    ) -> PreparedResponse {
1011        if is_asynchronous(message) {
1012            return PreparedResponse::Asynchronous;
1013        }
1014        let Some(head) = self.operations.front().copied() else {
1015            return PreparedResponse::Illegal;
1016        };
1017        if let Some(id) = local_id {
1018            if head.id != id {
1019                return PreparedResponse::Deferred;
1020            }
1021            if head.origin != Origin::Local {
1022                return PreparedResponse::Illegal;
1023            }
1024        } else if head.origin == Origin::Local {
1025            return if self.operations.iter().skip(1).any(|operation| {
1026                operation.origin == Origin::Forwarded && response_fits(*operation, message)
1027            }) {
1028                PreparedResponse::Deferred
1029            } else {
1030                PreparedResponse::Illegal
1031            };
1032        }
1033        if head.discarded || !response_fits(head, message) {
1034            return if self
1035                .operations
1036                .iter()
1037                .skip(1)
1038                .any(|operation| response_fits(*operation, message))
1039            {
1040                PreparedResponse::Deferred
1041            } else {
1042                PreparedResponse::Illegal
1043            };
1044        }
1045        PreparedResponse::Emit {
1046            head,
1047            response_state: head.response_state,
1048        }
1049    }
1050
1051    fn accept_response(
1052        &mut self,
1053        local_id: Option<OperationId>,
1054        message: BackendMessage,
1055    ) -> Result<BackendAction, BackendProjectionError> {
1056        let prepared = self.prepare_response(local_id, &message);
1057        self.commit_response(prepared, message)
1058    }
1059
1060    fn commit_response(
1061        &mut self,
1062        prepared: PreparedResponse,
1063        message: BackendMessage,
1064    ) -> Result<BackendAction, BackendProjectionError> {
1065        let PreparedResponse::Emit {
1066            head,
1067            response_state,
1068        } = prepared
1069        else {
1070            return match prepared {
1071                PreparedResponse::Asynchronous => Ok(BackendAction::Emit(message)),
1072                PreparedResponse::Deferred => Ok(BackendAction::Deferred(message)),
1073                PreparedResponse::Illegal => Err(BackendProjectionError {
1074                    state: self.state(),
1075                    message,
1076                }),
1077                PreparedResponse::Emit { .. } => unreachable!(),
1078            };
1079        };
1080        let event = backend::project_internal(response_state, &message)
1081            .expect("response was validated against its generated backend phase");
1082        let next_response_state = backend::transition(response_state, event)
1083            .expect("projected backend event has a generated transition")
1084            .target;
1085        let terminal = response_is_terminal(head.kind, &message);
1086        let error = matches!(message, BackendMessage::ErrorResponse(_));
1087        let copy_state = response_copy_state(next_response_state);
1088        if terminal {
1089            self.operations.pop_front();
1090            if error && is_extended_kind(head.kind) {
1091                self.enter_extended_error();
1092            }
1093        } else if let Some(head) = self.operations.front_mut() {
1094            head.response_state = next_response_state;
1095        }
1096        if !terminal {
1097            self.response_state = copy_state.map(RequestState::public);
1098        }
1099        if let Some(state) = copy_state
1100            && matches!(state, RequestState::CopyIn | RequestState::CopyBoth)
1101        {
1102            self.request_state = state;
1103        }
1104        if terminal {
1105            self.response_state = None;
1106        }
1107        self.remove_inert_heads();
1108        self.changed.notify_waiters();
1109        Ok(BackendAction::Emit(message))
1110    }
1111
1112    fn enter_extended_error(&mut self) {
1113        self.response_state = None;
1114        let mut sync_accepted = false;
1115        for operation in &mut self.operations {
1116            if operation.kind == OperationKind::Sync {
1117                sync_accepted = true;
1118                break;
1119            }
1120            operation.discarded = true;
1121        }
1122        if !sync_accepted {
1123            self.request_state = RequestState::ExtendedError;
1124        }
1125    }
1126
1127    fn remove_inert_heads(&mut self) {
1128        let previous_len = self.operations.len();
1129        while self.operations.front().is_some_and(|operation| {
1130            operation.kind == OperationKind::Flush
1131                || operation.kind == OperationKind::CopyData
1132                || operation.kind == OperationKind::CopyDone
1133                || operation.kind == OperationKind::CopyFail
1134                || operation.kind == OperationKind::Terminate
1135                || operation.discarded
1136        }) {
1137            self.operations.pop_front();
1138        }
1139        if self.operations.len() != previous_len {
1140            self.changed.notify_waiters();
1141        }
1142    }
1143}
1144
1145fn project_frontend(
1146    state: RequestState,
1147    message: &FrontendMessage,
1148) -> Option<(OperationKind, RequestState)> {
1149    use RequestState as S;
1150    let generated_state = match state {
1151        S::Ready => backend::RuntimeState::Ready,
1152        S::Extended { .. } => backend::RuntimeState::Building,
1153        S::ExtendedError => backend::RuntimeState::ExtendedError,
1154        S::CopyIn => backend::RuntimeState::ExtendedCopyIn,
1155        S::CopyOut => backend::RuntimeState::ExtendedCopyOut,
1156        S::CopyBoth => backend::RuntimeState::ExtendedCopyBoth,
1157        S::Terminated => backend::RuntimeState::Terminated,
1158    };
1159    backend::project_external(generated_state, message)?;
1160    classify_frontend(state, message)
1161}
1162
1163fn classify_frontend(
1164    state: RequestState,
1165    message: &FrontendMessage,
1166) -> Option<(OperationKind, RequestState)> {
1167    use FrontendMessage as F;
1168    use OperationKind as O;
1169    use RequestState as S;
1170
1171    match (state, message) {
1172        (S::Ready, F::Query(_)) => Some((O::Query, S::Ready)),
1173        (S::Ready, F::FunctionCall(_)) => Some((O::FunctionCall, S::Ready)),
1174        (S::Ready, F::Parse(_)) => Some((O::Parse, S::Extended { bound: false })),
1175        (S::Ready | S::Extended { .. }, F::Bind(_)) => Some((O::Bind, S::Extended { bound: true })),
1176        (S::Ready, F::Describe(_)) => Some((O::Describe, S::Extended { bound: false })),
1177        (S::Ready | S::Extended { .. }, F::Execute(_)) => {
1178            Some((O::Execute, S::Extended { bound: true }))
1179        }
1180        (S::Ready, F::Close(_)) => Some((O::Close, S::Extended { bound: false })),
1181        (S::Ready, F::Terminate) => Some((O::Terminate, S::Terminated)),
1182        (S::Extended { bound }, F::Parse(_)) => Some((O::Parse, S::Extended { bound })),
1183        (S::Extended { bound }, F::Describe(_)) => Some((O::Describe, S::Extended { bound })),
1184        (S::Extended { bound }, F::Close(_)) => Some((O::Close, S::Extended { bound })),
1185        (S::Extended { bound }, F::Flush) => Some((O::Flush, S::Extended { bound })),
1186        (S::Extended { .. } | S::ExtendedError, F::Sync) => Some((O::Sync, S::Ready)),
1187        (S::ExtendedError, _) => Some((classify_discard(message)?, S::ExtendedError)),
1188        (S::CopyIn, F::CopyData(_)) => Some((O::CopyData, S::CopyIn)),
1189        (S::CopyIn, F::CopyDone) => Some((O::CopyDone, S::Extended { bound: true })),
1190        (S::CopyIn, F::CopyFail(_)) => Some((O::CopyFail, S::ExtendedError)),
1191        (S::CopyBoth, F::CopyData(_)) => Some((O::CopyData, S::CopyBoth)),
1192        (S::CopyBoth, F::CopyDone) => Some((O::CopyDone, S::CopyBoth)),
1193        _ => None,
1194    }
1195}
1196
1197fn frontend_phase(state: RequestState, response_head: Option<OperationKind>) -> FrontendPhase {
1198    match (state, response_head) {
1199        (RequestState::Ready, _) => FrontendPhase::Ready,
1200        (RequestState::Extended { .. }, _) => FrontendPhase::Building,
1201        (RequestState::ExtendedError, _) => FrontendPhase::ExtendedError,
1202        (RequestState::CopyIn, Some(OperationKind::Query)) => FrontendPhase::SimpleCopyIn,
1203        (RequestState::CopyIn, Some(OperationKind::Execute)) => FrontendPhase::ExtendedCopyIn,
1204        (RequestState::CopyBoth, Some(OperationKind::Query)) => FrontendPhase::SimpleCopyBoth,
1205        (RequestState::CopyBoth, Some(OperationKind::Execute)) => FrontendPhase::ExtendedCopyBoth,
1206        (RequestState::CopyIn | RequestState::CopyBoth, _) => {
1207            unreachable!("COPY phase must belong to Query or Execute")
1208        }
1209        (RequestState::CopyOut | RequestState::Terminated, _) => {
1210            unreachable!("non-accepting frontend phase cannot be prepared")
1211        }
1212    }
1213}
1214
1215fn classify_discard(message: &FrontendMessage) -> Option<OperationKind> {
1216    Some(match message {
1217        FrontendMessage::Parse(_) => OperationKind::Parse,
1218        FrontendMessage::Bind(_) => OperationKind::Bind,
1219        FrontendMessage::Describe(_) => OperationKind::Describe,
1220        FrontendMessage::Execute(_) => OperationKind::Execute,
1221        FrontendMessage::Close(_) => OperationKind::Close,
1222        FrontendMessage::Flush => OperationKind::Flush,
1223        FrontendMessage::Query(_) => OperationKind::Query,
1224        FrontendMessage::FunctionCall(_) => OperationKind::FunctionCall,
1225        FrontendMessage::CopyData(_) => OperationKind::CopyData,
1226        FrontendMessage::CopyDone => OperationKind::CopyDone,
1227        FrontendMessage::CopyFail(_) => OperationKind::CopyFail,
1228        FrontendMessage::Terminate => OperationKind::Terminate,
1229        FrontendMessage::PasswordResponse(_) => return None,
1230        FrontendMessage::Sync => unreachable!("Sync is classified before discard"),
1231    })
1232}
1233
1234const fn initial_response_state(kind: OperationKind) -> backend::RuntimeState {
1235    use OperationKind as O;
1236    match kind {
1237        O::Query => backend::RuntimeState::Simple,
1238        O::FunctionCall => backend::RuntimeState::FunctionResponse,
1239        O::Parse => backend::RuntimeState::ParseResponse,
1240        O::Bind => backend::RuntimeState::BindResponse,
1241        O::Describe => backend::RuntimeState::DescribeResponse,
1242        O::Execute => backend::RuntimeState::ExecuteResponse,
1243        O::Close => backend::RuntimeState::CloseResponse,
1244        O::Sync => backend::RuntimeState::SyncResponse,
1245        O::Flush | O::CopyData | O::CopyDone | O::CopyFail | O::Terminate => {
1246            backend::RuntimeState::Terminated
1247        }
1248    }
1249}
1250
1251fn response_fits(operation: Operation, message: &BackendMessage) -> bool {
1252    backend::project_internal(operation.response_state, message).is_some()
1253}
1254
1255fn response_is_terminal(kind: OperationKind, message: &BackendMessage) -> bool {
1256    use BackendMessage as B;
1257    use OperationKind as O;
1258    match kind {
1259        O::Query | O::FunctionCall | O::Sync => matches!(message, B::ReadyForQuery(_)),
1260        O::Parse => matches!(message, B::ParseComplete | B::ErrorResponse(_)),
1261        O::Bind => matches!(message, B::BindComplete | B::ErrorResponse(_)),
1262        O::Describe => matches!(
1263            message,
1264            B::RowDescription(_) | B::NoData | B::ErrorResponse(_)
1265        ),
1266        O::Execute => matches!(
1267            message,
1268            B::CommandComplete(_) | B::PortalSuspended | B::ErrorResponse(_)
1269        ),
1270        O::Close => matches!(message, B::CloseComplete | B::ErrorResponse(_)),
1271        O::CopyDone => matches!(
1272            message,
1273            B::CopyDone | B::CommandComplete(_) | B::ErrorResponse(_)
1274        ),
1275        O::CopyFail => matches!(message, B::ErrorResponse(_)),
1276        O::Flush | O::CopyData | O::Terminate => true,
1277    }
1278}
1279
1280fn is_extended_kind(kind: OperationKind) -> bool {
1281    !matches!(
1282        kind,
1283        OperationKind::Query | OperationKind::FunctionCall | OperationKind::Terminate
1284    )
1285}
1286
1287fn response_copy_state(state: backend::RuntimeState) -> Option<RequestState> {
1288    use backend::RuntimeState as S;
1289    match state {
1290        S::SimpleCopyIn | S::ExtendedCopyIn => Some(RequestState::CopyIn),
1291        S::SimpleCopyOut | S::ExtendedCopyOut => Some(RequestState::CopyOut),
1292        S::SimpleCopyBoth
1293        | S::SimpleCopyBothClientDone
1294        | S::SimpleCopyBothServerDone
1295        | S::ExtendedCopyBoth
1296        | S::ExtendedCopyBothClientDone
1297        | S::ExtendedCopyBothServerDone => Some(RequestState::CopyBoth),
1298        _ => None,
1299    }
1300}
1301
1302fn is_asynchronous(message: &BackendMessage) -> bool {
1303    Demux::is_asynchronous(message)
1304}
pg_proto/pipeline.rs

pg_proto/
pipeline.rs
