s2n_quic_core/
query.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! This module provides `Query` and `QueryMut` traits, which are used for querying
5//! and executing functions on several different providers. This includes
6//! [`Subscriber::ConnectionContext`](crate::event::Subscriber::ConnectionContext)
7//! on a Subscriber and the [`Sender`](crate::datagram::Sender) and
8//! [`Receiver`](crate::datagram::Receiver) types on a datagram [`Endpoint`](crate::datagram::Endpoint).
9
10use core::marker::PhantomData;
11
12pub trait Query {
13    fn execute(&mut self, context: &dyn core::any::Any) -> ControlFlow;
14}
15
16pub trait QueryMut {
17    fn execute_mut(&mut self, context: &mut dyn core::any::Any) -> ControlFlow;
18}
19
20/// Used to tell a query whether it should exit early or go on as usual.
21#[derive(Debug, Clone, Copy, PartialEq)]
22pub enum ControlFlow {
23    Continue,
24    Break,
25}
26
27impl ControlFlow {
28    #[inline]
29    #[must_use]
30    pub fn and_then(self, f: impl FnOnce() -> Self) -> Self {
31        match self {
32            Self::Continue => f(),
33            Self::Break => Self::Break,
34        }
35    }
36}
37
38/// A type that implements Query and QueryMut traits and only executes once.
39///
40/// This will execute and short-circuit on the first match.
41pub struct Once<F, Context, Outcome> {
42    query: Option<F>,
43    result: Option<Outcome>,
44    context: PhantomData<Context>,
45}
46
47impl<F, Context, Outcome> From<Once<F, Context, Outcome>> for Result<Outcome, Error> {
48    #[inline]
49    fn from(query: Once<F, Context, Outcome>) -> Self {
50        query.result.ok_or(Error::ContextTypeMismatch)
51    }
52}
53
54impl<F, ConnectionContext, Outcome> Once<F, ConnectionContext, Outcome>
55where
56    F: FnOnce(&ConnectionContext) -> Outcome,
57    ConnectionContext: 'static,
58{
59    #[inline]
60    pub fn new(query: F) -> Self {
61        Self {
62            query: Some(query),
63            result: None,
64            context: PhantomData,
65        }
66    }
67}
68
69impl<F, ConnectionContext, Outcome> Once<F, ConnectionContext, Outcome>
70where
71    F: FnOnce(&mut ConnectionContext) -> Outcome,
72    ConnectionContext: 'static,
73{
74    #[inline]
75    pub fn new_mut(query: F) -> Self {
76        Self {
77            query: Some(query),
78            result: None,
79            context: PhantomData,
80        }
81    }
82}
83
84impl<F, Context, Outcome> Query for Once<F, Context, Outcome>
85where
86    F: FnOnce(&Context) -> Outcome,
87    Context: 'static,
88{
89    fn execute(&mut self, context: &dyn core::any::Any) -> ControlFlow {
90        match context.downcast_ref::<Context>() {
91            Some(context) => {
92                let query = self.query.take().expect("can only match once");
93                self.result = Some(query(context));
94                ControlFlow::Break
95            }
96            None => ControlFlow::Continue,
97        }
98    }
99}
100
101impl<F, Context, Outcome> QueryMut for Once<F, Context, Outcome>
102where
103    F: FnOnce(&mut Context) -> Outcome,
104    Context: 'static,
105{
106    fn execute_mut(&mut self, context: &mut dyn core::any::Any) -> ControlFlow {
107        match context.downcast_mut::<Context>() {
108            Some(context) => {
109                let query = self.query.take().expect("can only match once");
110                self.result = Some(query(context));
111                ControlFlow::Break
112            }
113            None => ControlFlow::Continue,
114        }
115    }
116}
117
118#[non_exhaustive]
119#[derive(Debug, Clone)]
120/// Reason for the failed query.
121pub enum Error {
122    /// The connection lock is poisoned and the connection unusable.
123    ConnectionLockPoisoned,
124
125    /// The expected query type failed to match any of the configured types.
126    ContextTypeMismatch,
127}
128
129impl core::fmt::Display for Error {
130    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
131        write!(f, "{self:?}")
132    }
133}
134
135#[cfg(feature = "std")]
136impl std::error::Error for Error {}
s2n_quic_core/query.rs

s2n_quic_core/
query.rs
