Skip to main content

rama_core/
extensions.rs

1#![allow(clippy::disallowed_types)]
2//! Extensions passed to and between services
3//!
4//! # State
5//!
6//! [`rama`] supports two kinds of states:
7//!
8//! 1. static state: this state can be a part of the service struct or captured by a closure
9//! 2. dynamic state: these can be injected as [`Extensions`]s in Requests/Responses/Connections if it [`ExtensionsRef`]
10//!
11//! Any state that is optional, and especially optional state injected by middleware, can be inserted using extensions.
12//! It is however important to try as much as possible to then also consume this state in an approach that deals
13//! gracefully with its absence. Good examples of this are header-related inputs. Headers might be set or not,
14//! and so absence of [`Extensions`]s that might be created as a result of these might reasonably not exist.
15//! It might of course still mean the app returns an error response when it is absent, but it should not unwrap/panic.
16//!
17//! [`rama`]: crate
18//!
19//! # Examples
20//!
21//! ## Example: Extensions
22//! ```
23//! use rama_core::extensions::{Extensions, Extension};
24//!
25//! #[derive(Debug, Clone, PartialEq)]
26//! struct MyExt(i32);
27//! impl Extension for MyExt {}
28//!
29//! let mut ext = Extensions::default();
30//! ext.insert(MyExt(5));
31//! assert_eq!(ext.get_ref::<MyExt>(), Some(&MyExt(5)));
32//! ```
33
34use core::any::{Any, TypeId};
35use core::fmt;
36use core::pin::Pin;
37
38use crate::std::{boxed::Box, sync::Arc, vec::Vec};
39
40pub use rama_macros::{Extension, FromExtensions};
41use rama_utils::collections::AppendOnlyVec;
42use rama_utils::macros::impl_deref;
43
44#[derive(Clone, Default)]
45/// A type map of protocol extensions.
46///
47/// [`Extension`]s are internally stored in a type erased [`Arc`]. Since values
48/// are stored in an [`Arc`] there are extra methods exposed that build on top
49/// of this and leverage characteristics of an [`Arc`] to expose things like
50/// cheap cloning of the Arc.
51///
52/// [`Extensions`] may have an optional [`parent`][Self::parent]: the
53/// [`Extensions`] this one was forked from. Lookups walk the parent chain when
54/// the local [`Extension`]s don't have the requested type. The parent relationship is
55/// best described as "I'm layered on top of that, but I'm not exactly the same":
56/// - Retry attempts fork from the original request (don't leak failed extensions)
57/// - Responses fork the request (response != request)
58/// - H2 streams fork from underlying H2 connection (nested connection with isolated properties)
59///
60/// Connection's who logically map one-to-one we don't fork and we just pass the [`Extensions`]
61/// up, examples are:
62/// - TLS layered on top of TCP
63/// - HTTP layered on top of TLS
64/// - ...
65pub struct Extensions {
66    extensions: Arc<AppendOnlyVec<TypeErasedExtension, 12, 3>>,
67    parent: Option<Box<Self>>,
68}
69
70impl Extensions {
71    /// Create an empty [`Extensions`] store with no parent.
72    #[inline(always)]
73    #[must_use]
74    pub fn new() -> Self {
75        Self::default()
76    }
77
78    /// Create a fresh child [`Extensions`] whose parent is this [`Extensions`] store.
79    ///
80    /// The child has its own empty top-level storage. Lookups that miss
81    /// locally walk into the parent (and recursively up the chain). Inserts
82    /// land on this child only, parents are never mutated through the child.
83    #[must_use]
84    pub fn fork(&self) -> Self {
85        Self {
86            extensions: Arc::new(AppendOnlyVec::new()),
87            parent: Some(Box::new(self.clone())),
88        }
89    }
90
91    /// The parent [`Extensions`] this blob was forked from, if any.
92    #[inline(always)]
93    #[must_use]
94    pub fn parent(&self) -> Option<&Self> {
95        self.parent.as_deref()
96    }
97
98    /// Return a view of this [`Extensions`] chain layered on top of `base`.
99    ///
100    /// The toplevel extensions are still the same and are editable.
101    /// Lookups resolve through this chain first falling back to `base`
102    /// only when nothing matches. `base` is all the way at the bottom.
103    ///
104    /// Use this when you have an Extensions (chain) to which you want to apply
105    /// defaults, e.g. `request_extensions().with_base(&base_extension_config)`
106    #[must_use]
107    pub fn with_base(&self, base: &Self) -> Self {
108        let parent = match &self.parent {
109            Some(parent) => parent.with_base(base),
110            None => base.clone(),
111        };
112        Self {
113            extensions: Arc::clone(&self.extensions),
114            parent: Some(Box::new(parent)),
115        }
116    }
117
118    /// Insert a type `T` into this [`Extensions`] store.
119    ///
120    /// This method returns a reference to the just inserted value.
121    ///
122    /// If the value you are inserting is an `Arc<T>`, prefer using
123    /// [`Self::insert_arc`] to prevent the double indirection of storing
124    /// an `Arc<Arc<T>>`.
125    pub fn insert<T: Extension>(&self, val: T) -> &T {
126        let extension = TypeErasedExtension::new(val);
127        let idx = self.extensions.push(extension);
128
129        #[expect(
130            clippy::unwrap_used,
131            reason = "`downcast_ref` can only be none if TypeId doesn't match, but we just inserted this type"
132        )]
133        self.extensions[idx].downcast_ref::<T>().unwrap()
134    }
135
136    /// Insert a type `Arc<T>` into this [`Extensions]` store.
137    ///
138    /// This method returns a a cloned Arc of the value just inserted
139    ///
140    /// If the value you are inserting is not an `Arc<T>` or you don't
141    /// need a cloned `Arc<T>` prefer using [`Self::insert()`]
142    pub fn insert_arc<T: Extension>(&self, val: Arc<T>) -> Arc<T> {
143        let extension = TypeErasedExtension::new_arc(val);
144        let idx = self.extensions.push(extension);
145
146        #[expect(
147            clippy::unwrap_used,
148            reason = "`cloned_downcast` can only be none if TypeId doesn't match, but we just inserted this type"
149        )]
150        self.extensions[idx].cloned_downcast::<T>().unwrap()
151    }
152
153    /// Extend this [`Extensions`] store with the other [`Extensions`].
154    ///
155    /// The other [`Extensions`]s will be appended behind the current ones
156    pub fn extend(&self, other: &Self) {
157        for ext in other.extensions.iter() {
158            self.extensions.push(ext.clone());
159        }
160    }
161
162    /// Returns true if the [`Extensions`] store contains the given type.
163    ///
164    /// This function is recursive and will traverse multiple nested [`Extensions`]
165    /// stores to find the correct item. See [`Extensions::get_ref()`] for how this works.
166    ///
167    /// If you don't want any of this special logic and you just want to check this
168    /// [`Extensions`] store, use [`Extensions::self_contains()`] instead.
169    #[must_use]
170    pub fn contains<T: Extension>(&self) -> bool {
171        self.get_ref::<T>().is_some()
172    }
173
174    /// Returns true if this [`Extensions`] store contains the given type
175    ///
176    /// This only checks this [`Extensions`] store
177    #[must_use]
178    pub fn self_contains<T: Extension>(&self) -> bool {
179        let type_id = TypeId::of::<T>();
180        self.extensions
181            .iter()
182            .rev()
183            .any(|item| item.type_id == type_id)
184    }
185
186    #[must_use]
187    /// Get a reference to `T`. Walks the parent chain and the connection
188    /// wrappers ([`Egress`] / [`Ingress`]) if not found locally.
189    ///
190    /// Search rule (single pass, newest insertion wins):
191    ///
192    /// 1. Iterate local entries newest -> oldest. For each entry:
193    ///    - if its type matches `T`, return it,
194    ///    - if it is an [`Egress<Extensions>`] or [`Ingress<Extensions>`]
195    ///      wrapper, recurse into the wrapped blob with the same rule
196    ///      (its own local first, then its wrappers, then its parent)
197    ///      and return any match found,
198    ///    - otherwise skip.
199    /// 2. If still not found, recurse into the parent (same rule applied).
200    ///
201    /// Wrappers are spliced into the local scan in insertion order, so a
202    /// connection-pointer detour inserted on this blob is treated as part of
203    /// "local" for ordering purposes: a directly-inserted `T` and a wrapper
204    /// containing `T` both compete by insertion time, newest wins. Parent is
205    /// only consulted after the entire local scan (including wrapper
206    /// recursion) finishes empty. The wrappers themselves can still be
207    /// retrieved directly (or via [`Self::egress`] / [`Self::ingress`]).
208    ///
209    /// For a raw flat lookup, use [`Self::self_get_ref`].
210    ///
211    /// Returns the most recently inserted match, for the oldest, see [`Self::self_first_ref`].
212    pub fn get_ref<T: Extension>(&self) -> Option<&T> {
213        let target = TypeId::of::<T>();
214        let egress_id = TypeId::of::<Egress<Self>>();
215        let ingress_id = TypeId::of::<Ingress<Self>>();
216        for ext in self.extensions.iter().rev() {
217            if ext.type_id == target {
218                if let Some(v) = ext.downcast_ref::<T>() {
219                    return Some(v);
220                }
221            } else if ext.type_id == egress_id
222                && let Some(eg) = ext.downcast_ref::<Egress<Self>>()
223                && let Some(v) = eg.0.get_ref::<T>()
224            {
225                return Some(v);
226            } else if ext.type_id == ingress_id
227                && let Some(ig) = ext.downcast_ref::<Ingress<Self>>()
228                && let Some(v) = ig.0.get_ref::<T>()
229            {
230                return Some(v);
231            }
232        }
233        self.parent().and_then(|p| p.get_ref::<T>())
234    }
235
236    /// Raw flat [`Self::get_ref`]: returns the most recently inserted `T`, for the oldest, see [`Self::self_first_ref`].
237    ///
238    /// This only checks this [`Extensions`] store
239    #[must_use]
240    pub fn self_get_ref<T: Extension>(&self) -> Option<&T> {
241        let type_id = TypeId::of::<T>();
242        self.extensions
243            .iter()
244            .rev()
245            .find(|item| item.type_id == type_id)
246            .and_then(|ext| ext.downcast_ref())
247    }
248
249    #[must_use]
250    /// Get an owned `Arc<T>`. Walks the parent chain and the structural
251    /// connection wrappers if not found locally.
252    ///
253    /// See [`Self::get_ref`] for the search order.
254    ///
255    /// For a raw flat lookup (top-level only), use [`Self::self_get_arc`].
256    pub fn get_arc<T: Extension>(&self) -> Option<Arc<T>> {
257        let target = TypeId::of::<T>();
258        let egress_id = TypeId::of::<Egress<Self>>();
259        let ingress_id = TypeId::of::<Ingress<Self>>();
260        for ext in self.extensions.iter().rev() {
261            if ext.type_id == target {
262                if let Some(v) = ext.cloned_downcast::<T>() {
263                    return Some(v);
264                }
265            } else if ext.type_id == egress_id
266                && let Some(eg) = ext.downcast_ref::<Egress<Self>>()
267                && let Some(v) = eg.0.get_arc::<T>()
268            {
269                return Some(v);
270            } else if ext.type_id == ingress_id
271                && let Some(ig) = ext.downcast_ref::<Ingress<Self>>()
272                && let Some(v) = ig.0.get_arc::<T>()
273            {
274                return Some(v);
275            }
276        }
277        self.parent().and_then(|p| p.get_arc::<T>())
278    }
279
280    /// Raw flat [`Self::get_arc`]: returns the most recently inserted `T`
281    ///
282    /// This only checks this [`Extensions`] store
283    #[must_use]
284    pub fn self_get_arc<T: Extension>(&self) -> Option<Arc<T>> {
285        let type_id = TypeId::of::<T>();
286        self.extensions
287            .iter()
288            .rev()
289            .find(|item| item.type_id == type_id)
290            .and_then(|ext| ext.cloned_downcast())
291    }
292
293    /// Fetch several extension types in a single pass, returning a tuple of
294    /// `Option<&T>` mirroring the requested tuple. Each slot uses the same
295    /// search rule as [`Self::get_ref`] (newest-wins, walking the
296    /// [`Egress`] / [`Ingress`] wrappers and the parent chain), but the whole
297    /// structure is traversed only once instead of once per type.
298    ///
299    /// See [`Self::get_many_arc`] for an owned-[`Arc`] variant.
300    pub fn get_many_ref<'a, T: GetManyRef<'a>>(&'a self) -> T::Output {
301        T::get_many_ref(self)
302    }
303
304    /// Like [`Self::get_many_ref`], but returns a tuple of `Option<Arc<T>>`
305    /// (cheap, owned [`Arc`] clones) instead of borrowed references, the
306    /// multi-fetch counterpart of [`Self::get_arc`]. Same single-pass traversal.
307    pub fn get_many_arc<T: GetManyArc>(&self) -> T::Output {
308        T::get_many_arc(self)
309    }
310
311    /// Single iteration logic behind [`Self::get_many_ref`], [`Self::get_many_arc`]
312    /// and `#[derive(FromExtensions)]`. Hidden and not part of the public
313    /// API surface: call `get_many_ref`/`get_many_arc` or the derive instead.
314    ///
315    /// For each requested [`TypeId`] in `targets`, fill the matching slot in
316    /// `out` with the newest matching entry, walking wrappers and the parent
317    /// chain like [`Self::get_ref`]. Slots already `Some` are left untouched, so
318    /// callers may pre-fill or reuse the buffer. `targets[i]` fills `out[i]`.
319    ///
320    /// Alongside the entry, each filled slot records its rank: the position of
321    /// the entry in this single newest to oldest traversal. The newest entry
322    /// visited is rank `0` and the count grows by one for every entry walked
323    /// back (across wrappers and the parent chain). Comparing the ranks of two
324    /// filled slots tells you which value was inserted more recently (0 = newest).
325    ///
326    /// A rank is a completely opaque type and should only be used to compare positions,
327    /// it does not tell anything about the absolute position.
328    #[doc(hidden)]
329    pub fn get_many_erased<'a, const N: usize>(
330        &'a self,
331        targets: &[TypeId; N],
332        out: &mut [Option<(&'a TypeErasedExtension, usize)>; N],
333    ) {
334        let mut rank = 0;
335        self.get_many_erased_ranked(targets, out, &mut rank);
336    }
337
338    fn get_many_erased_ranked<'a, const N: usize>(
339        &'a self,
340        targets: &[TypeId; N],
341        out: &mut [Option<(&'a TypeErasedExtension, usize)>; N],
342        rank: &mut usize,
343    ) {
344        let egress_id = TypeId::of::<Egress<Self>>();
345        let ingress_id = TypeId::of::<Ingress<Self>>();
346        let mut remaining = out.iter().filter(|slot| slot.is_none()).count();
347        if remaining == 0 {
348            return;
349        }
350        for ext in self.extensions.iter().rev() {
351            let current = *rank;
352            *rank += 1;
353            for (i, &tid) in targets.iter().enumerate() {
354                if out[i].is_none() && ext.type_id == tid {
355                    out[i] = Some((ext, current));
356                    remaining -= 1;
357
358                    if remaining == 0 {
359                        return;
360                    }
361                }
362            }
363
364            if ext.type_id == egress_id
365                && let Some(eg) = ext.downcast_ref::<Egress<Self>>()
366            {
367                eg.0.get_many_erased_ranked(targets, out, rank);
368                remaining = out.iter().filter(|slot| slot.is_none()).count();
369                if remaining == 0 {
370                    return;
371                }
372            } else if ext.type_id == ingress_id
373                && let Some(ig) = ext.downcast_ref::<Ingress<Self>>()
374            {
375                ig.0.get_many_erased_ranked(targets, out, rank);
376                remaining = out.iter().filter(|slot| slot.is_none()).count();
377                if remaining == 0 {
378                    return;
379                }
380            }
381        }
382        if remaining != 0
383            && let Some(parent) = self.parent()
384        {
385            parent.get_many_erased_ranked(targets, out, rank);
386        }
387    }
388
389    /// Recursive find-or-create: return `&T` if one exists anywhere in this
390    /// this [`Extensions`] store (using [`Self::get_ref`] dispatch), otherwise
391    /// insert the value produced by `create_fn` at the top level and return
392    /// a reference to it.
393    ///
394    /// Useful when a type conceptually belongs to the scope (e.g. `ConnectionHealth`
395    /// on a connection chain) and you want to reuse an existing instance rather
396    /// than create a duplicate at every layer. For strict "ensure local exists",
397    /// use [`Self::self_get_ref_or_insert`].
398    pub fn get_ref_or_insert<T, F>(&self, create_fn: F) -> &T
399    where
400        T: Extension,
401        F: FnOnce() -> T,
402    {
403        self.get_ref().unwrap_or_else(|| self.insert(create_fn()))
404    }
405
406    /// Recursive find-or-create returning an [`Arc<T>`]: see [`Self::get_ref_or_insert`].
407    pub fn get_arc_or_insert<T, F>(&self, create_fn: F) -> Arc<T>
408    where
409        T: Extension,
410        F: FnOnce() -> Arc<T>,
411    {
412        self.get_arc()
413            .unwrap_or_else(|| self.insert_arc(create_fn()))
414    }
415
416    /// Raw flat find-or-create: return `&T` if one exists at the top level of
417    /// this [`Extensions`] store, otherwise insert the value produced by
418    /// `create_fn` at the top level and return a reference to it.
419    ///
420    /// Does not follow the parent chain. Useful when you want strict "ensure T
421    /// exists on THIS blob" (e.g. materializing a direction wrapper
422    /// like [`Ingress<Connection<Extensions>>`] at the outer blob).
423    pub fn self_get_ref_or_insert<T, F>(&self, create_fn: F) -> &T
424    where
425        T: Extension,
426        F: FnOnce() -> T,
427    {
428        self.self_get_ref()
429            .unwrap_or_else(|| self.insert(create_fn()))
430    }
431
432    /// Raw flat find-or-create returning an [`Arc<T>`]: see [`Self::self_get_ref_or_insert`].
433    pub fn self_get_arc_or_insert<T, F>(&self, create_fn: F) -> Arc<T>
434    where
435        T: Extension,
436        F: FnOnce() -> Arc<T>,
437    {
438        self.self_get_arc()
439            .unwrap_or_else(|| self.insert_arc(create_fn()))
440    }
441
442    /// Raw flat reference to the oldest inserted `T` at the top level of this
443    /// [`Extensions`] store, does not walk structural wrappers.
444    ///
445    /// In most cases you want [`Self::get_ref`] (newest, scope-aware). Use this
446    /// only when you specifically need insertion order access inside this [`Extensions`]
447    /// store.
448    ///
449    /// Currently we don't provide a recursive variant of this method since we don't have
450    /// a use case for it, and it's not exactly clear what would be considered "first".
451    #[must_use]
452    pub fn self_first_ref<T: Extension>(&self) -> Option<&T> {
453        let type_id = TypeId::of::<T>();
454        self.extensions
455            .iter()
456            .find(|item| item.type_id == type_id)
457            .and_then(|ext| ext.downcast_ref())
458    }
459
460    /// Raw flat [`Arc<T>`] to the oldest inserted `T` at the top level, see
461    /// [`Self::self_first_ref`] for caveats.
462    #[must_use]
463    pub fn self_first_arc<T: Extension>(&self) -> Option<Arc<T>> {
464        let type_id = TypeId::of::<T>();
465        self.extensions
466            .iter()
467            .find(|item| item.type_id == type_id)
468            .and_then(|ext| ext.cloned_downcast())
469    }
470
471    /// Raw flat iteration over all inserted items of type `T` at the top level
472    /// of this [`Extensions`] store, newest to oldest.
473    ///
474    /// The order matches [`Self::self_get_ref`] (newest-first), so
475    /// `self_iter_ref::<T>().next() == self_get_ref::<T>()`.
476    pub fn self_iter_ref<T: Extension>(&self) -> impl Iterator<Item = &T> {
477        let type_id = TypeId::of::<T>();
478
479        self.extensions
480            .iter()
481            .rev()
482            .filter(move |item| item.type_id == type_id)
483            .filter_map(TypeErasedExtension::downcast_ref::<T>)
484    }
485
486    /// Raw flat iteration over all inserted items of type `T` at the top level
487    /// as cloned [`Arc`] values, newest to oldest.
488    ///
489    /// The order matches [`Self::self_get_arc`] (newest-first), so
490    /// `self_iter_arc::<T>().next() == self_get_arc::<T>()`.
491    pub fn self_iter_arc<T: Extension>(&self) -> impl Iterator<Item = Arc<T>> {
492        let type_id = TypeId::of::<T>();
493
494        self.extensions
495            .iter()
496            .rev()
497            .filter(move |item| item.type_id == type_id)
498            .filter_map(TypeErasedExtension::cloned_downcast::<T>)
499    }
500
501    /// Raw flat iteration over all [`TypeErasedExtension`] entries at the top
502    /// level of this [`Extensions`] store.
503    ///
504    /// Use to efficiently combine different types of [`Extension`]s in a single
505    /// iteration. [`TypeErasedExtension`] exposes methods to convert back to
506    /// type `T` when it matches the erased type.
507    pub fn self_iter_all(&self) -> impl Iterator<Item = &TypeErasedExtension> {
508        self.extensions.iter()
509    }
510
511    /// Iterate over all inserted items of type `T`, walking the parent chain
512    /// and the structural [`Egress`] / [`Ingress`] connection wrappers.
513    ///
514    /// Yield order matches [`Self::get_ref`] preference (so
515    /// `iter_ref::<T>().next() == get_ref::<T>()`):
516    ///
517    /// At each level, iterate the local entries newest -> oldest. For each
518    /// entry: yield it if its type matches `T`, if it is an
519    /// [`Egress<Extensions>`] or [`Ingress<Extensions>`] wrapper, recurse into
520    /// the wrapped blob (same rule applied) and yield its results inline. Then
521    /// recurse into the parent.
522    ///
523    /// For a flat top-level-only iteration use [`Self::self_iter_ref`].
524    ///
525    /// The iterator type is left opaque (`impl Iterator`) so the internal
526    /// representation can change without breaking callers.
527    pub fn iter_ref<T: Extension>(&self) -> impl Iterator<Item = &T> + '_ {
528        self.iter_ref_inner::<T>()
529    }
530
531    /// Iteration yielding cloned [`Arc<T>`] values, see [`Self::iter_ref`].
532    pub fn iter_arc<T: Extension>(&self) -> impl Iterator<Item = Arc<T>> + '_ {
533        self.iter_arc_inner::<T>()
534    }
535
536    // TODO replace this later with a custom Iterator to avoid boxing
537    fn iter_ref_inner<T: Extension>(&self) -> Box<dyn Iterator<Item = &T> + '_> {
538        let target = TypeId::of::<T>();
539        let egress_id = TypeId::of::<Egress<Self>>();
540        let ingress_id = TypeId::of::<Ingress<Self>>();
541        let local = self.extensions.iter().rev().flat_map(
542            move |ext| -> Box<dyn Iterator<Item = &T> + '_> {
543                if ext.type_id == target {
544                    match ext.downcast_ref::<T>() {
545                        Some(v) => Box::new(core::iter::once(v)),
546                        None => Box::new(core::iter::empty()),
547                    }
548                } else if ext.type_id == egress_id {
549                    match ext.downcast_ref::<Egress<Self>>() {
550                        Some(e) => e.0.iter_ref_inner::<T>(),
551                        None => Box::new(core::iter::empty()),
552                    }
553                } else if ext.type_id == ingress_id {
554                    match ext.downcast_ref::<Ingress<Self>>() {
555                        Some(i) => i.0.iter_ref_inner::<T>(),
556                        None => Box::new(core::iter::empty()),
557                    }
558                } else {
559                    Box::new(core::iter::empty())
560                }
561            },
562        );
563        let parent: Box<dyn Iterator<Item = &T>> = match self.parent() {
564            Some(p) => p.iter_ref_inner::<T>(),
565            None => Box::new(core::iter::empty()),
566        };
567        Box::new(local.chain(parent))
568    }
569
570    // TODO replace this later with a custom Iterator to avoid boxing
571    fn iter_arc_inner<T: Extension>(&self) -> Box<dyn Iterator<Item = Arc<T>> + '_> {
572        let target = TypeId::of::<T>();
573        let egress_id = TypeId::of::<Egress<Self>>();
574        let ingress_id = TypeId::of::<Ingress<Self>>();
575        let local = self.extensions.iter().rev().flat_map(
576            move |ext| -> Box<dyn Iterator<Item = Arc<T>> + '_> {
577                if ext.type_id == target {
578                    match ext.cloned_downcast::<T>() {
579                        Some(v) => Box::new(core::iter::once(v)),
580                        None => Box::new(core::iter::empty()),
581                    }
582                } else if ext.type_id == egress_id {
583                    match ext.downcast_ref::<Egress<Self>>() {
584                        Some(e) => e.0.iter_arc_inner::<T>(),
585                        None => Box::new(core::iter::empty()),
586                    }
587                } else if ext.type_id == ingress_id {
588                    match ext.downcast_ref::<Ingress<Self>>() {
589                        Some(i) => i.0.iter_arc_inner::<T>(),
590                        None => Box::new(core::iter::empty()),
591                    }
592                } else {
593                    Box::new(core::iter::empty())
594                }
595            },
596        );
597        let parent: Box<dyn Iterator<Item = Arc<T>>> = match self.parent() {
598            Some(p) => p.iter_arc_inner::<T>(),
599            None => Box::new(core::iter::empty()),
600        };
601        Box::new(local.chain(parent))
602    }
603
604    /// Get a reference to the [`Ingress<Extensions>`] wrapper if one exists
605    /// on this blob or anywhere reachable through the parent chain.
606    ///
607    /// Returns `None` when no ingress wrapper has been set up. For correctly
608    /// constructed server-side requests this is always `Some`, server
609    /// stacks insert the wrapper at the boundary where the connection becomes
610    /// visible to a request, so a `None` here is almost always a framework
611    /// setup bug rather than a normal state.
612    ///
613    /// This is just a shortcut for `extensions.get_ref::<Ingress<Extensions>>()`
614    #[must_use]
615    pub fn ingress(&self) -> Option<&Ingress<Self>> {
616        self.get_ref::<Ingress<Self>>()
617    }
618
619    /// Get a reference to the [`Egress<Extensions>`] wrapper if one exists
620    /// on this blob or anywhere reachable through the parent chain.
621    /// See [`Self::ingress`] for semantics.
622    ///
623    /// This is just a shortcut for `extensions.get_ref::<Egress<Extensions>>()`
624    #[must_use]
625    pub fn egress(&self) -> Option<&Egress<Self>> {
626        self.get_ref::<Egress<Self>>()
627    }
628
629    /// Clone an [`Extension`] from [`Extensions`] to another and get a `Arc` clone of it.
630    pub fn clone_to<T: Extension>(&self, target: &Self) -> Option<Arc<T>> {
631        let item = self.get_arc();
632        if let Some(item) = item.clone() {
633            target.insert_arc(item);
634        };
635
636        item
637    }
638
639    /// Clone an [`Extension`] from [`Extensions`] to another, if and only if the other
640    /// [`Extensions`] store does not already contain this [`Extension`] `T`.
641    ///
642    /// If the other [`Extensions`] store already contains it, we return [`Extension`] `T`
643    /// from the other store, otherwise we return the `T` that we transferred over.
644    ///
645    /// This is mainly used in connectors where [`Extension`]s that become part of the connection
646    /// should be transfer from the input to the connection.
647    pub fn clone_to_if_absent<T: Extension>(&self, target: &Self) -> Option<Arc<T>> {
648        let item = target.get_arc::<T>();
649        if item.is_some() {
650            return item;
651        }
652
653        self.clone_to(target)
654    }
655}
656
657/// Macro-support trait that lets a `#[derive(FromExtensions)]` group (e.g. an
658/// any-of enum) be folded into a parent's single pass instead of running its own
659/// traversal. Implemented by the derive, hidden and not part of the public API.
660///
661/// A group occupies [`Self::TARGETS`] consecutive slots in the shared buffers:
662/// [`Self::from_ext_targets`] writes its candidate [`TypeId`]s at `offset`, and
663/// after one [`Extensions::get_many_erased`] pass [`Self::from_ext_slots`] reads
664/// the same window back. Because ranks are global across the pass, newest-wins
665/// selection still holds across the whole struct.
666#[doc(hidden)]
667pub trait FromExtensionsGroup<'a>: Sized {
668    /// Number of slots this group occupies in the shared single pass.
669    const TARGETS: usize;
670    /// Write the group's candidate `TypeId`s into `targets[offset..offset + TARGETS]`.
671    fn from_ext_targets(targets: &mut [TypeId], offset: usize);
672    /// Build the group value from the filled slots `out[offset..offset + TARGETS]`,
673    /// or `None` if no candidate is present.
674    fn from_ext_slots(
675        out: &[Option<(&'a TypeErasedExtension, usize)>],
676        offset: usize,
677    ) -> Option<Self>;
678}
679
680/// Helper trait powering [`Extensions::get_many_ref`]: implemented for tuples of
681/// [`Extension`] types (up to arity 12).
682pub trait GetManyRef<'a>: Sized {
683    /// Tuple of `Option<&'a T>` mirroring the requested tuple of types.
684    type Output;
685
686    #[doc(hidden)]
687    const SEALED: seal::Seal;
688    #[doc(hidden)]
689    fn get_many_ref(ext: &'a Extensions) -> Self::Output;
690}
691
692/// Helper trait powering [`Extensions::get_many_arc`]; the owned-[`Arc`]
693/// counterpart of [`GetManyRef`].
694pub trait GetManyArc: Sized {
695    /// Tuple of `Option<Arc<T>>` mirroring the requested tuple of types.
696    type Output;
697
698    #[doc(hidden)]
699    const SEALED: seal::Seal;
700    #[doc(hidden)]
701    fn get_many_arc(ext: &Extensions) -> Self::Output;
702}
703
704macro_rules! impl_get_many {
705    ($n:literal; $($T:ident => $idx:tt),+ $(,)?) => {
706        impl<'a, $($T: Extension),+> GetManyRef<'a> for ($($T,)+) {
707            type Output = ($(Option<&'a $T>,)+);
708
709            const SEALED: seal::Seal = seal::Seal;
710            fn get_many_ref(ext: &'a Extensions) -> Self::Output {
711                let targets = [$(TypeId::of::<$T>()),+];
712                let mut out: [Option<(&'a TypeErasedExtension, usize)>; $n] = [None; $n];
713                ext.get_many_erased(&targets, &mut out);
714                ($(out[$idx].and_then(|(e, _)| e.downcast_ref::<$T>()),)+)
715            }
716        }
717
718        impl<$($T: Extension),+> GetManyArc for ($($T,)+) {
719            type Output = ($(Option<Arc<$T>>,)+);
720
721            const SEALED: seal::Seal = seal::Seal;
722            fn get_many_arc(ext: &Extensions) -> Self::Output {
723                let targets = [$(TypeId::of::<$T>()),+];
724                let mut out: [Option<(&TypeErasedExtension, usize)>; $n] = [None; $n];
725                ext.get_many_erased(&targets, &mut out);
726                ($(out[$idx].and_then(|(e, _)| e.cloned_downcast::<$T>()),)+)
727            }
728        }
729    };
730}
731
732impl_get_many!(1; A => 0);
733impl_get_many!(2; A => 0, B => 1);
734impl_get_many!(3; A => 0, B => 1, C => 2);
735impl_get_many!(4; A => 0, B => 1, C => 2, D => 3);
736impl_get_many!(5; A => 0, B => 1, C => 2, D => 3, E => 4);
737impl_get_many!(6; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5);
738impl_get_many!(7; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5, G => 6);
739impl_get_many!(8; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5, G => 6, H => 7);
740impl_get_many!(9; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5, G => 6, H => 7, I => 8);
741impl_get_many!(10; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5, G => 6, H => 7, I => 8, J => 9);
742impl_get_many!(11; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5, G => 6, H => 7, I => 8, J => 9, K => 10);
743impl_get_many!(12; A => 0, B => 1, C => 2, D => 3, E => 4, F => 5, G => 6, H => 7, I => 8, J => 9, K => 10, L => 11);
744
745impl fmt::Debug for Extensions {
746    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
747        let mut s = f.debug_struct("Extensions");
748        if let Some(parent) = self.parent() {
749            s.field("parent", parent);
750        }
751        s.field(
752            "entries",
753            &self.extensions.iter().map(|e| &e.value).collect::<Vec<_>>(),
754        );
755
756        s.finish()
757    }
758}
759
760#[derive(Clone, Debug)]
761/// A [`TypeErasedExtension`] is a type erased item which can be stored in an [`Extensions`]
762///
763/// Internally the value is stored inside an `Arc` so this is cheap to clone
764pub struct TypeErasedExtension {
765    type_id: TypeId,
766    value: Arc<dyn Extension>,
767}
768
769impl TypeErasedExtension {
770    /// Create a new [`TypeErasedExtension`] for `T`
771    ///
772    /// If the value you are inserting is an `Arc<T>`, prefer using
773    /// [`Self::new_arc()`] to prevent the double indirection of storing
774    /// an `Arc<Arc<T>>`. This happens because internally we use a type erased
775    /// Arc to store the actual value.
776    pub fn new<T: Extension>(value: T) -> Self {
777        Self {
778            type_id: TypeId::of::<T>(),
779            value: Arc::new(value),
780        }
781    }
782
783    /// Create a new [`TypeErasedExtension`] for `Arc<T>`
784    ///
785    ///
786    /// If the value you are inserting is not an `Arc<T>` prefer using
787    /// [`Self::new()`] instead.
788    pub fn new_arc<T: Extension>(value: Arc<T>) -> Self {
789        Self {
790            type_id: TypeId::of::<T>(),
791            value,
792        }
793    }
794
795    /// Get the [`TypeId`] for the internally stored type `Arc<T>`
796    pub fn type_id(&self) -> TypeId {
797        self.type_id
798    }
799
800    /// Get a cloned `Arc<T>` of the internally stored type `Arc<T>`
801    ///
802    /// This method will return `None`, if the internally stored
803    /// type `S` doesn't match the requested type `T`
804    pub fn cloned_downcast<T: Extension>(&self) -> Option<Arc<T>> {
805        let any = self.value.clone() as Arc<dyn Any + Send + Sync>;
806        any.downcast::<T>().ok()
807    }
808
809    /// Get a reference `&T` of the internally stored type `T`
810    ///
811    /// This method will return `None`, if the internally stored
812    /// type `S` doesn't match the requested type `T`
813    pub fn downcast_ref<T: Extension>(&self) -> Option<&T> {
814        let inner_any = self.value.as_ref() as &dyn Any;
815        (inner_any).downcast_ref::<T>()
816    }
817}
818
819#[derive(Debug, Clone, Extension)]
820/// Ingress connection wrapper use by servers
821pub struct Ingress<T>(pub T);
822
823impl_deref!(Ingress);
824
825#[derive(Debug, Clone, Extension)]
826/// Egress connection wrapper use by client
827pub struct Egress<T>(pub T);
828
829impl_deref!(Egress);
830
831// We use this syntax: [`TlsExtension`] — TLS and secure transport
832// Instead of [`TlsExtension`]: TLS and secure transport
833// Because otherwise we hit `link definitions are not shown in rendered documentation`
834
835/// [`Extension`] is type which can be stored inside an [`Extensions`] store
836///
837/// This is has to be manually implement or can be implemented using `#[derive(Extension)]`
838///
839/// We have not implemented this for any container types:
840/// - `Arc<T>`: sounds nice, but by not implement it, it has become impossible to misuse `Extensions::insert()`
841///   with `Extensions::insert_arc()`. Otherwise this is very tricky and error prone
842/// - `Vec<T>`: Collections should use the new type pattern to give it a meaningfull name, and to prevent collisions
843///
844/// There might be valid use cases for implementing it for other type of containers, so in case you run into these
845/// open a Github issue and we can see if implementing it makes sense
846///
847/// # Extension Tags
848///
849/// Extensions can be tagged with one or more categories using the `#[extension(tags(tag1, tag2))]`
850/// attribute on the derive macro. This generates implementations for the corresponding
851/// marker traits below, which groups them in rust docs
852///
853/// - [`TlsExtension`] — TLS and secure transport
854/// - [`HttpExtension`] — HTTP protocol
855/// - [`NetExtension`] — Network and connection level
856/// - [`UaExtension`] — User-agent emulation
857/// - [`ProxyExtension`] — Proxy
858/// - [`WsExtension`] — WebSocket
859/// - [`DnsExtension`] — DNS resolution
860/// - [`GrpcExtension`] — gRPC
861///
862/// ```rust,ignore
863/// #[derive(Debug, Clone, Extension)]
864/// #[extension(tags(tls, net))]
865/// pub struct SecureTransport(..);
866/// ```
867///
868/// Types that implement [`Extension`] manually can opt into tagged docs by implementing
869/// the marker trait(s) directly:
870///
871/// ```rust,ignore
872/// impl Extension for MyType {}
873/// impl HttpExtension for MyType {}
874/// ```
875pub trait Extension: Any + Send + Sync + core::fmt::Debug + 'static {}
876
877/// TLS and secure transport related extensions.
878///
879/// Derive with `#[extension(tags(tls))]`
880pub trait TlsExtension: Extension {}
881
882/// HTTP protocol related extensions.
883///
884/// Derive with `#[extension(tags(http))]`
885pub trait HttpExtension: Extension {}
886
887/// Network and connection level extensions.
888///
889/// Derive with `#[extension(tags(net))]`
890pub trait NetExtension: Extension {}
891
892/// User-agent emulation related extensions.
893///
894/// Derive with `#[extension(tags(ua))]`
895pub trait UaExtension: Extension {}
896
897/// Proxy related extensions.
898///
899/// Derive with `#[extension(tags(proxy))]`
900pub trait ProxyExtension: Extension {}
901
902/// WebSocket related extensions.
903///
904/// Derive with `#[extension(tags(ws))]`
905pub trait WsExtension: Extension {}
906
907/// DNS resolution related extensions.
908///
909/// Derive with `#[extension(tags(dns))]`
910pub trait DnsExtension: Extension {}
911
912/// gRPC related extensions.
913///
914/// Derive with `#[extension(tags(grpc))]`
915pub trait GrpcExtension: Extension {}
916
917pub trait ExtensionsRef {
918    /// Get reference to the underlying [`Extensions`] store
919    fn extensions(&self) -> &Extensions;
920}
921
922impl ExtensionsRef for Extensions {
923    fn extensions(&self) -> &Extensions {
924        self
925    }
926}
927
928impl<T> ExtensionsRef for &T
929where
930    T: ExtensionsRef,
931{
932    #[inline(always)]
933    fn extensions(&self) -> &Extensions {
934        (**self).extensions()
935    }
936}
937
938impl<T> ExtensionsRef for &mut T
939where
940    T: ExtensionsRef,
941{
942    #[inline(always)]
943    fn extensions(&self) -> &Extensions {
944        (**self).extensions()
945    }
946}
947
948impl<T> ExtensionsRef for Box<T>
949where
950    T: ExtensionsRef,
951{
952    fn extensions(&self) -> &Extensions {
953        (**self).extensions()
954    }
955}
956
957impl<T> ExtensionsRef for Pin<Box<T>>
958where
959    T: ExtensionsRef,
960{
961    fn extensions(&self) -> &Extensions {
962        (**self).extensions()
963    }
964}
965
966impl<T> ExtensionsRef for Arc<T>
967where
968    T: ExtensionsRef,
969{
970    fn extensions(&self) -> &Extensions {
971        (**self).extensions()
972    }
973}
974
975macro_rules! impl_extensions_either {
976    ($id:ident, $($param:ident),+ $(,)?) => {
977        impl<$($param),+,> ExtensionsRef for crate::combinators::$id<$($param),+>
978        where
979            $($param: ExtensionsRef,)+
980        {
981            fn extensions(&self) -> &Extensions {
982                match self {
983                    $(crate::combinators::$id::$param(s) => s.extensions(),)+
984                }
985            }
986        }
987    };
988}
989
990crate::combinators::impl_either!(impl_extensions_either);
991
992mod seal {
993    pub struct Seal;
994}
995
996#[cfg(test)]
997mod tests {
998    use super::*;
999    use core::any::TypeId;
1000    use core::pin::Pin;
1001    use core::sync::atomic::{AtomicUsize, Ordering};
1002
1003    #[derive(Debug, Clone, PartialEq, Eq, Extension)]
1004    struct TraceNote(String);
1005
1006    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Extension)]
1007    struct RetryBudget(u32);
1008
1009    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Extension)]
1010    struct ConnectionTimeoutMs(u64);
1011
1012    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Extension)]
1013    struct WorkerId(i32);
1014
1015    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Extension)]
1016    struct HealthSignal(u8);
1017
1018    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Extension)]
1019    struct FeatureToggle(bool);
1020
1021    #[test]
1022    fn get_ref_returns_last_inserted() {
1023        let ext = Extensions::new();
1024        ext.insert(TraceNote("first".to_owned()));
1025        ext.insert(TraceNote("second".to_owned()));
1026        ext.insert(TraceNote("third".to_owned()));
1027
1028        assert_eq!(
1029            ext.get_ref::<TraceNote>(),
1030            Some(&TraceNote("third".to_owned()))
1031        );
1032    }
1033
1034    #[test]
1035    fn clone_shares_backing_store() {
1036        let ext = Extensions::new();
1037        ext.insert(TraceNote("first".to_owned()));
1038
1039        let clone = ext.clone();
1040        clone.insert(TraceNote("second".to_owned()));
1041
1042        assert_eq!(
1043            ext.get_ref::<TraceNote>(),
1044            Some(&TraceNote("second".to_owned()))
1045        );
1046        assert_eq!(
1047            clone.get_ref::<TraceNote>(),
1048            Some(&TraceNote("second".to_owned()))
1049        );
1050    }
1051
1052    #[test]
1053    fn get_ref_none_when_absent() {
1054        let ext = Extensions::new();
1055        assert_eq!(ext.get_ref::<TraceNote>(), None);
1056    }
1057
1058    #[test]
1059    fn get_arc_none_when_absent() {
1060        let ext = Extensions::new();
1061        assert!(ext.get_arc::<TraceNote>().is_none());
1062    }
1063
1064    #[test]
1065    fn first_ref_none_when_absent() {
1066        let ext = Extensions::new();
1067        assert_eq!(ext.self_first_ref::<TraceNote>(), None);
1068    }
1069
1070    #[test]
1071    fn first_arc_none_when_absent() {
1072        let ext = Extensions::new();
1073        assert!(ext.self_first_arc::<TraceNote>().is_none());
1074    }
1075
1076    #[test]
1077    fn first_ref_returns_first_inserted() {
1078        let ext = Extensions::new();
1079        ext.insert(TraceNote("first".to_owned()));
1080        ext.insert(TraceNote("second".to_owned()));
1081
1082        assert_eq!(
1083            ext.self_first_ref::<TraceNote>(),
1084            Some(&TraceNote("first".to_owned()))
1085        );
1086    }
1087
1088    #[test]
1089    fn extend_appends_other_extensions() {
1090        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Extension)]
1091        struct DerivedMetric(i32);
1092
1093        let source = Extensions::new();
1094        source.insert(WorkerId(5));
1095        source.insert(DerivedMetric(10));
1096
1097        let target = Extensions::new();
1098        target.extend(&source);
1099
1100        assert_eq!(target.get_ref::<WorkerId>(), Some(&WorkerId(5)));
1101        assert_eq!(target.get_ref::<DerivedMetric>(), Some(&DerivedMetric(10)));
1102    }
1103
1104    #[test]
1105    fn insert_arc_can_be_retrieved_via_get_arc() {
1106        let ext = Extensions::new();
1107        let inserted = ext.insert_arc(Arc::new(TraceNote(String::from("hello"))));
1108        let retrieved = ext.get_arc::<TraceNote>();
1109
1110        assert_eq!(inserted.0.as_str(), "hello");
1111        assert_eq!(retrieved.as_deref().map(|it| it.0.as_str()), Some("hello"));
1112    }
1113
1114    #[test]
1115    fn insert_arc_can_be_retrieved_via_get_ref() {
1116        let ext = Extensions::new();
1117        ext.insert_arc(Arc::new(WorkerId(99)));
1118        assert_eq!(ext.get_ref::<WorkerId>(), Some(&WorkerId(99)));
1119    }
1120
1121    #[test]
1122    fn contains_reports_presence_and_absence() {
1123        let ext = Extensions::new();
1124        assert!(!ext.contains::<RetryBudget>());
1125
1126        ext.insert(RetryBudget(1));
1127        assert!(ext.contains::<RetryBudget>());
1128        assert!(!ext.contains::<ConnectionTimeoutMs>());
1129    }
1130
1131    #[test]
1132    fn get_arc_and_first_arc_report_latest_and_oldest() {
1133        let ext = Extensions::new();
1134        ext.insert_arc(Arc::new(TraceNote(String::from("first"))));
1135        ext.insert_arc(Arc::new(TraceNote(String::from("second"))));
1136
1137        assert_eq!(
1138            ext.self_first_arc::<TraceNote>()
1139                .as_deref()
1140                .map(|it| it.0.as_str()),
1141            Some("first")
1142        );
1143        assert_eq!(
1144            ext.get_arc::<TraceNote>()
1145                .as_deref()
1146                .map(|it| it.0.as_str()),
1147            Some("second")
1148        );
1149    }
1150
1151    #[test]
1152    fn get_ref_or_insert_uses_existing_or_inserts_once() {
1153        let ext = Extensions::new();
1154        ext.insert(RetryBudget(5));
1155
1156        let calls = AtomicUsize::new(0);
1157        let existing = ext.self_get_ref_or_insert(|| {
1158            calls.fetch_add(1, Ordering::SeqCst);
1159            RetryBudget(6)
1160        });
1161        assert_eq!(existing.0, 5u32);
1162        assert_eq!(calls.load(Ordering::SeqCst), 0);
1163
1164        let missing = ext.self_get_ref_or_insert(|| {
1165            calls.fetch_add(1, Ordering::SeqCst);
1166            ConnectionTimeoutMs(7)
1167        });
1168        assert_eq!(missing.0, 7u64);
1169        assert_eq!(calls.load(Ordering::SeqCst), 1);
1170    }
1171
1172    #[test]
1173    fn get_arc_or_insert_uses_existing_or_inserts_once() {
1174        let ext = Extensions::new();
1175        ext.insert_arc(Arc::new(TraceNote(String::from("stored"))));
1176
1177        let calls = AtomicUsize::new(0);
1178        let existing = ext.self_get_arc_or_insert(|| {
1179            calls.fetch_add(1, Ordering::SeqCst);
1180            Arc::new(TraceNote(String::from("new")))
1181        });
1182        assert_eq!(existing.0.as_str(), "stored");
1183        assert_eq!(calls.load(Ordering::SeqCst), 0);
1184
1185        let missing = ext.self_get_arc_or_insert(|| {
1186            calls.fetch_add(1, Ordering::SeqCst);
1187            Arc::new(RetryBudget(11))
1188        });
1189        assert_eq!(missing.0, 11u32);
1190        assert_eq!(calls.load(Ordering::SeqCst), 1);
1191    }
1192
1193    #[test]
1194    fn iter_all_exposes_all_items_in_insert_order() {
1195        let ext = Extensions::new();
1196        ext.insert(HealthSignal(1));
1197        ext.insert(FeatureToggle(true));
1198        ext.insert(HealthSignal(2));
1199
1200        let type_ids: Vec<TypeId> = ext
1201            .self_iter_all()
1202            .map(TypeErasedExtension::type_id)
1203            .collect();
1204        assert_eq!(
1205            type_ids,
1206            vec![
1207                TypeId::of::<HealthSignal>(),
1208                TypeId::of::<FeatureToggle>(),
1209                TypeId::of::<HealthSignal>()
1210            ]
1211        );
1212    }
1213
1214    #[test]
1215    fn iter_for_missing_type_is_empty() {
1216        let ext = Extensions::new();
1217        ext.insert(HealthSignal(1));
1218
1219        assert_eq!(ext.self_iter_ref::<TraceNote>().count(), 0);
1220        assert_eq!(ext.self_iter_arc::<TraceNote>().count(), 0);
1221    }
1222
1223    #[test]
1224    fn iter_ref_returns_items_for_present_type_in_newest_to_oldest_order() {
1225        let ext = Extensions::new();
1226        ext.insert(TraceNote(String::from("first")));
1227        ext.insert(HealthSignal(9));
1228        ext.insert(TraceNote(String::from("second")));
1229
1230        let output: Vec<&str> = ext
1231            .self_iter_ref::<TraceNote>()
1232            .map(|it| it.0.as_str())
1233            .collect();
1234        assert_eq!(output, vec!["second", "first"]);
1235    }
1236
1237    #[test]
1238    fn iter_arc_returns_items_for_present_type_in_newest_to_oldest_order() {
1239        let ext = Extensions::new();
1240        ext.insert(TraceNote(String::from("first")));
1241        ext.insert(HealthSignal(9));
1242        ext.insert(TraceNote(String::from("second")));
1243
1244        let output: Vec<String> = ext
1245            .self_iter_arc::<TraceNote>()
1246            .map(|arc| arc.0.clone())
1247            .collect();
1248        assert_eq!(output, vec!["second".to_owned(), "first".to_owned()]);
1249    }
1250
1251    #[test]
1252    fn type_erased_new_supports_downcast_ref_and_cloned_downcast() {
1253        let ext = TypeErasedExtension::new(TraceNote(String::from("hello")));
1254
1255        assert_eq!(ext.type_id(), TypeId::of::<TraceNote>());
1256        assert_eq!(
1257            ext.downcast_ref::<TraceNote>().map(|it| it.0.as_str()),
1258            Some("hello")
1259        );
1260        assert_eq!(
1261            ext.cloned_downcast::<TraceNote>()
1262                .as_deref()
1263                .map(|it| it.0.as_str()),
1264            Some("hello")
1265        );
1266        assert!(ext.downcast_ref::<RetryBudget>().is_none());
1267        assert!(ext.cloned_downcast::<RetryBudget>().is_none());
1268    }
1269
1270    #[test]
1271    fn type_erased_new_arc_supports_all_downcasts() {
1272        let ext = TypeErasedExtension::new_arc(Arc::new(TraceNote(String::from("hello"))));
1273
1274        assert_eq!(ext.type_id(), TypeId::of::<TraceNote>());
1275        assert_eq!(
1276            ext.downcast_ref::<TraceNote>().map(|it| it.0.as_str()),
1277            Some("hello")
1278        );
1279        assert_eq!(
1280            ext.cloned_downcast::<TraceNote>()
1281                .as_deref()
1282                .map(|it| it.0.as_str()),
1283            Some("hello")
1284        );
1285        assert!(ext.downcast_ref::<RetryBudget>().is_none());
1286        assert!(ext.cloned_downcast::<RetryBudget>().is_none());
1287    }
1288
1289    #[test]
1290    fn extensions_ref_blanket_impls_forward_to_underlying_extensions() {
1291        let base = Extensions::new();
1292        base.insert(RetryBudget(7));
1293
1294        let by_ref: &Extensions = &base;
1295        assert_eq!(
1296            by_ref.extensions().get_ref::<RetryBudget>(),
1297            Some(&RetryBudget(7))
1298        );
1299
1300        let mut base_for_mut = base.clone();
1301        let by_mut_ref: &mut Extensions = &mut base_for_mut;
1302        assert_eq!(
1303            by_mut_ref.extensions().get_ref::<RetryBudget>(),
1304            Some(&RetryBudget(7))
1305        );
1306
1307        let boxed = Box::new(base.clone());
1308        assert_eq!(
1309            boxed.extensions().get_ref::<RetryBudget>(),
1310            Some(&RetryBudget(7))
1311        );
1312
1313        let pinned = Pin::new(Box::new(base.clone()));
1314        assert_eq!(
1315            pinned.extensions().get_ref::<RetryBudget>(),
1316            Some(&RetryBudget(7))
1317        );
1318
1319        let arced = Arc::new(base);
1320        assert_eq!(
1321            arced.extensions().get_ref::<RetryBudget>(),
1322            Some(&RetryBudget(7))
1323        );
1324    }
1325
1326    #[derive(Debug, Clone, PartialEq, Eq, Extension)]
1327    struct ConnSocketInfo(&'static str);
1328
1329    #[derive(Debug, Clone, PartialEq, Eq, Extension)]
1330    struct RequestId(u64);
1331
1332    #[test]
1333    fn get_finds_local() {
1334        let req = Extensions::new();
1335        req.insert(RequestId(42));
1336        assert_eq!(req.get_ref::<RequestId>(), Some(&RequestId(42)));
1337    }
1338
1339    #[test]
1340    fn get_walks_parent_chain() {
1341        let req = Extensions::new();
1342        req.insert(RequestId(7));
1343
1344        let resp = req.fork();
1345        assert_eq!(resp.get_ref::<RequestId>(), Some(&RequestId(7)));
1346    }
1347
1348    #[test]
1349    fn local_shadows_parent() {
1350        let req = Extensions::new();
1351        req.insert(RequestId(7));
1352
1353        let attempt = req.fork();
1354        attempt.insert(RequestId(99));
1355
1356        assert_eq!(attempt.get_ref::<RequestId>(), Some(&RequestId(99)));
1357    }
1358
1359    #[test]
1360    fn fork_isolates_writes() {
1361        let req = Extensions::new();
1362        req.insert(RequestId(1));
1363
1364        let attempt = req.fork();
1365        attempt.insert(RequestId(2));
1366
1367        assert_eq!(req.get_ref::<RequestId>(), Some(&RequestId(1)));
1368    }
1369
1370    #[test]
1371    fn ingress_view_walks_parent() {
1372        let conn_ext = Extensions::new();
1373        conn_ext.insert(ConnSocketInfo("client-in"));
1374
1375        let req = Extensions::new();
1376        req.insert(Ingress(conn_ext));
1377
1378        assert_eq!(
1379            req.ingress().and_then(|i| i.get_ref::<ConnSocketInfo>()),
1380            Some(&ConnSocketInfo("client-in"))
1381        );
1382    }
1383
1384    #[test]
1385    fn egress_view_walks_parent() {
1386        let conn_ext = Extensions::new();
1387        conn_ext.insert(ConnSocketInfo("egress-side"));
1388
1389        let req = Extensions::new();
1390        req.insert(Egress(conn_ext));
1391
1392        assert_eq!(
1393            req.egress().and_then(|e| e.get_ref::<ConnSocketInfo>()),
1394            Some(&ConnSocketInfo("egress-side"))
1395        );
1396    }
1397
1398    #[test]
1399    fn ingress_egress_disambiguate_in_mitm() {
1400        let in_conn = Extensions::new();
1401        in_conn.insert(ConnSocketInfo("in"));
1402        let out_conn = Extensions::new();
1403        out_conn.insert(ConnSocketInfo("out"));
1404
1405        let req = Extensions::new();
1406        req.insert(Ingress(in_conn));
1407        req.insert(Egress(out_conn));
1408
1409        assert_eq!(
1410            req.ingress().and_then(|i| i.get_ref::<ConnSocketInfo>()),
1411            Some(&ConnSocketInfo("in"))
1412        );
1413        assert_eq!(
1414            req.egress().and_then(|e| e.get_ref::<ConnSocketInfo>()),
1415            Some(&ConnSocketInfo("out"))
1416        );
1417    }
1418
1419    #[test]
1420    fn egress_view_walks_through_parent_to_find_wrapper() {
1421        let conn_ext = Extensions::new();
1422        conn_ext.insert(ConnSocketInfo("inside-parent"));
1423        let req = Extensions::new();
1424        req.insert(Egress(conn_ext));
1425
1426        let resp = req.fork();
1427        assert_eq!(
1428            resp.egress().and_then(|e| e.get_ref::<ConnSocketInfo>()),
1429            Some(&ConnSocketInfo("inside-parent"))
1430        );
1431    }
1432
1433    #[test]
1434    fn ingress_egress_return_none_when_absent() {
1435        let req = Extensions::new();
1436        assert!(req.ingress().is_none());
1437        assert!(req.egress().is_none());
1438    }
1439
1440    #[test]
1441    fn iter_ref_yields_local_then_parent_newest_to_oldest() {
1442        let parent = Extensions::new();
1443        parent.insert(RequestId(1));
1444        parent.insert(RequestId(2));
1445
1446        let child = parent.fork();
1447        child.insert(RequestId(3));
1448        child.insert(RequestId(4));
1449
1450        let ids: Vec<_> = child.iter_ref::<RequestId>().map(|r| r.0).collect();
1451        assert_eq!(ids, vec![4, 3, 2, 1]);
1452    }
1453
1454    #[test]
1455    fn iter_ref_walks_egress_and_ingress_wrappers_inline() {
1456        let conn_in = Extensions::new();
1457        conn_in.insert(RequestId(10));
1458        conn_in.insert(RequestId(11));
1459        let conn_out = Extensions::new();
1460        conn_out.insert(RequestId(20));
1461
1462        let req = Extensions::new();
1463        req.insert(RequestId(1));
1464        req.insert(Ingress(conn_in));
1465        req.insert(Egress(conn_out));
1466
1467        let ids: Vec<_> = req.iter_ref::<RequestId>().map(|r| r.0).collect();
1468        assert_eq!(ids, vec![20, 11, 10, 1]);
1469    }
1470
1471    #[test]
1472    fn local_direct_after_wrapper_shadows_wrapper() {
1473        let conn = Extensions::new();
1474        conn.insert(RequestId(99));
1475        let req = Extensions::new();
1476        req.insert(Ingress(conn));
1477        req.insert(RequestId(1));
1478
1479        assert_eq!(req.get_ref::<RequestId>(), Some(&RequestId(1)));
1480    }
1481
1482    #[test]
1483    fn wrapper_after_local_direct_shadows_direct() {
1484        let conn = Extensions::new();
1485        conn.insert(RequestId(99));
1486        let req = Extensions::new();
1487        req.insert(RequestId(1));
1488        req.insert(Ingress(conn));
1489
1490        assert_eq!(req.get_ref::<RequestId>(), Some(&RequestId(99)));
1491    }
1492
1493    #[test]
1494    fn iter_ref_first_matches_get_ref() {
1495        let parent = Extensions::new();
1496        parent.insert(RequestId(1));
1497        let child = parent.fork();
1498        child.insert(RequestId(2));
1499
1500        assert_eq!(
1501            child.iter_ref::<RequestId>().next(),
1502            child.get_ref::<RequestId>()
1503        );
1504    }
1505
1506    #[test]
1507    fn get_many_present_and_absent() {
1508        let ext = Extensions::new();
1509        ext.insert(RequestId(7));
1510        ext.insert(ConnSocketInfo("a"));
1511
1512        let (id, sock, toggle) = ext.get_many_ref::<(RequestId, ConnSocketInfo, FeatureToggle)>();
1513        assert_eq!(id, Some(&RequestId(7)));
1514        assert_eq!(sock, Some(&ConnSocketInfo("a")));
1515        assert_eq!(toggle, None);
1516    }
1517
1518    #[test]
1519    fn get_many_each_slot_matches_get_ref() {
1520        let ext = Extensions::new();
1521        ext.insert(RequestId(1));
1522        ext.insert(RequestId(2)); // newest wins
1523        ext.insert(ConnSocketInfo("x"));
1524
1525        let (id, sock) = ext.get_many_ref::<(RequestId, ConnSocketInfo)>();
1526        assert_eq!(id, ext.get_ref::<RequestId>());
1527        assert_eq!(sock, ext.get_ref::<ConnSocketInfo>());
1528        assert_eq!(id, Some(&RequestId(2)));
1529    }
1530
1531    #[test]
1532    fn get_many_walks_parent_chain() {
1533        let parent = Extensions::new();
1534        parent.insert(RequestId(1));
1535        let child = parent.fork();
1536        child.insert(ConnSocketInfo("x"));
1537
1538        let (id, sock) = child.get_many_ref::<(RequestId, ConnSocketInfo)>();
1539        assert_eq!(id, Some(&RequestId(1)));
1540        assert_eq!(sock, Some(&ConnSocketInfo("x")));
1541    }
1542
1543    #[test]
1544    fn with_base_and_chain() {
1545        let base = Extensions::new();
1546        base.insert(RequestId(1));
1547        base.insert(ConnSocketInfo("base"));
1548
1549        let req = Extensions::new();
1550        req.insert(RequestId(2));
1551        req.insert(TraceNote("mid".to_owned()));
1552        let req_retry = req.fork();
1553        req_retry.insert(RequestId(3));
1554
1555        let combined = req_retry.with_base(&base);
1556
1557        assert_eq!(combined.get_ref::<RequestId>(), Some(&RequestId(3)));
1558        assert_eq!(
1559            combined.get_ref::<TraceNote>(),
1560            Some(&TraceNote("mid".to_owned())),
1561        );
1562
1563        assert_eq!(
1564            combined.get_ref::<ConnSocketInfo>(),
1565            Some(&ConnSocketInfo("base")),
1566        );
1567
1568        assert_eq!(base.get_ref::<RequestId>(), Some(&RequestId(1)));
1569        assert!(req_retry.get_ref::<ConnSocketInfo>().is_none());
1570    }
1571
1572    #[test]
1573    fn get_many_walks_wrappers() {
1574        let conn = Extensions::new();
1575        conn.insert(ConnSocketInfo("in"));
1576        let req = Extensions::new();
1577        req.insert(RequestId(7));
1578        req.insert(Ingress(conn));
1579
1580        let (id, sock) = req.get_many_ref::<(RequestId, ConnSocketInfo)>();
1581        assert_eq!(id, Some(&RequestId(7)));
1582        assert_eq!(sock, Some(&ConnSocketInfo("in")));
1583    }
1584
1585    #[derive(FromExtensions)]
1586    struct GatherView<'a> {
1587        id: Option<&'a RequestId>,
1588        sock: Option<&'a ConnSocketInfo>,
1589        toggle: Option<&'a FeatureToggle>,
1590    }
1591
1592    #[test]
1593    fn derive_from_extensions_gathers_pieces() {
1594        let ext = Extensions::new();
1595        ext.insert(RequestId(7));
1596        ext.insert(ConnSocketInfo("a"));
1597
1598        let view = GatherView::from_extensions(&ext);
1599        assert_eq!(view.id, Some(&RequestId(7)));
1600        assert_eq!(view.sock, Some(&ConnSocketInfo("a")));
1601        assert_eq!(view.toggle, None);
1602        assert_eq!(view.id, ext.get_ref::<RequestId>());
1603    }
1604
1605    #[test]
1606    fn derive_from_extensions_walks_parent() {
1607        let parent = Extensions::new();
1608        parent.insert(RequestId(1));
1609        let child = parent.fork();
1610        child.insert(ConnSocketInfo("x"));
1611
1612        let view = GatherView::from_extensions(&child);
1613        assert_eq!(view.id, Some(&RequestId(1)));
1614        assert_eq!(view.sock, Some(&ConnSocketInfo("x")));
1615        assert_eq!(view.toggle, None);
1616    }
1617
1618    #[test]
1619    fn get_many_arc_returns_owned_arcs() {
1620        let ext = Extensions::new();
1621        ext.insert(RequestId(7));
1622
1623        let (id, sock) = ext.get_many_arc::<(RequestId, ConnSocketInfo)>();
1624        assert_eq!(id.as_deref(), Some(&RequestId(7)));
1625        assert_eq!(sock, None);
1626    }
1627
1628    #[derive(FromExtensions)]
1629    struct MixedView<'a> {
1630        id_ref: Option<&'a RequestId>,
1631        sock_arc: Option<Arc<ConnSocketInfo>>,
1632    }
1633
1634    #[test]
1635    fn derive_from_extensions_mixed_ref_and_arc() {
1636        let ext = Extensions::new();
1637        ext.insert(RequestId(7));
1638        ext.insert(ConnSocketInfo("a"));
1639
1640        let view = MixedView::from_extensions(&ext);
1641        assert_eq!(view.id_ref, Some(&RequestId(7)));
1642        assert_eq!(view.sock_arc.as_deref(), Some(&ConnSocketInfo("a")));
1643    }
1644
1645    #[derive(FromExtensions)]
1646    struct AllArc {
1647        id_ref: Option<Arc<RequestId>>,
1648        sock_arc: Option<Arc<ConnSocketInfo>>,
1649    }
1650
1651    #[test]
1652    fn derive_from_extensions_all_arc() {
1653        let ext = Extensions::new();
1654        ext.insert(RequestId(7));
1655        ext.insert(ConnSocketInfo("a"));
1656
1657        let view = AllArc::from_extensions(&ext);
1658        assert_eq!(view.id_ref.as_deref(), Some(&RequestId(7)));
1659        assert_eq!(view.sock_arc.as_deref(), Some(&ConnSocketInfo("a")));
1660    }
1661
1662    #[derive(FromExtensions)]
1663    struct RankedView<'a> {
1664        id: Option<(&'a RequestId, usize)>,
1665        sock: Option<(&'a ConnSocketInfo, usize)>,
1666        toggle: Option<(&'a FeatureToggle, usize)>,
1667    }
1668
1669    #[test]
1670    fn derive_from_extensions_captures_rank() {
1671        let ext = Extensions::new();
1672        ext.insert(RequestId(7)); // oldest
1673        ext.insert(ConnSocketInfo("a")); // newest
1674
1675        let view = RankedView::from_extensions(&ext);
1676
1677        assert_eq!(view.sock, Some((&ConnSocketInfo("a"), 0)));
1678        assert_eq!(view.id, Some((&RequestId(7), 1)));
1679        assert_eq!(view.toggle, None);
1680
1681        assert!(view.sock.unwrap().1 < view.id.unwrap().1);
1682    }
1683
1684    #[test]
1685    fn derive_from_extensions_rank_arc_variant() {
1686        #[derive(FromExtensions)]
1687        struct RankedArc {
1688            id: Option<(Arc<RequestId>, usize)>,
1689        }
1690
1691        let ext = Extensions::new();
1692        ext.insert(ConnSocketInfo("a"));
1693        ext.insert(RequestId(7));
1694
1695        let view = RankedArc::from_extensions(&ext);
1696        let (id, rank) = view.id.expect("present");
1697        assert_eq!(&*id, &RequestId(7));
1698        assert_eq!(rank, 0);
1699    }
1700
1701    #[derive(Debug, PartialEq, Eq, FromExtensions)]
1702    enum AnyOf<'a> {
1703        Req(&'a RequestId),
1704        Sock(&'a ConnSocketInfo),
1705    }
1706
1707    #[test]
1708    fn derive_from_extensions_enum_newest_wins() {
1709        let ext = Extensions::new();
1710        ext.insert(ConnSocketInfo("a"));
1711        ext.insert(RequestId(7));
1712        assert_eq!(
1713            AnyOf::from_extensions(&ext),
1714            Some(AnyOf::Req(&RequestId(7)))
1715        );
1716
1717        let ext = Extensions::new();
1718        ext.insert(RequestId(7));
1719        ext.insert(ConnSocketInfo("a"));
1720        assert_eq!(
1721            AnyOf::from_extensions(&ext),
1722            Some(AnyOf::Sock(&ConnSocketInfo("a")))
1723        );
1724
1725        let ext = Extensions::new();
1726        ext.insert(RequestId(7));
1727        assert_eq!(
1728            AnyOf::from_extensions(&ext),
1729            Some(AnyOf::Req(&RequestId(7)))
1730        );
1731
1732        let ext = Extensions::new();
1733        ext.insert(FeatureToggle(true));
1734        assert_eq!(AnyOf::from_extensions(&ext), None);
1735    }
1736
1737    #[derive(FromExtensions)]
1738    struct ConfigView<'a> {
1739        toggle: Option<&'a FeatureToggle>,
1740        either: Option<AnyOf<'a>>,
1741    }
1742
1743    #[derive(Debug, PartialEq, Eq, FromExtensions)]
1744    enum SameType<'a> {
1745        First(&'a RequestId),
1746        Second(&'a RequestId),
1747    }
1748
1749    #[test]
1750    fn derive_from_extensions_enum_same_type_ties_to_earlier_variant() {
1751        // Both variants name `RequestId`, so they resolve to the same entry
1752        // (equal rank), the tie breaks deterministically toward the earlier
1753        // variant, `First`.
1754        let ext = Extensions::new();
1755        ext.insert(RequestId(7));
1756        assert_eq!(
1757            SameType::from_extensions(&ext),
1758            Some(SameType::First(&RequestId(7)))
1759        );
1760    }
1761
1762    #[test]
1763    fn derive_from_extensions_nested_group_field() {
1764        let ext = Extensions::new();
1765        ext.insert(FeatureToggle(true));
1766        ext.insert(RequestId(7));
1767        ext.insert(ConnSocketInfo("a"));
1768
1769        let view = ConfigView::from_extensions(&ext);
1770        assert_eq!(view.toggle, Some(&FeatureToggle(true)));
1771        assert_eq!(view.either, Some(AnyOf::Sock(&ConnSocketInfo("a"))));
1772
1773        let ext = Extensions::new();
1774        ext.insert(FeatureToggle(false));
1775        let view = ConfigView::from_extensions(&ext);
1776        assert_eq!(view.toggle, Some(&FeatureToggle(false)));
1777        assert_eq!(view.either, None);
1778    }
1779
1780    #[test]
1781    fn derive_from_extensions_enum_newest_wins_across_parent() {
1782        let parent = Extensions::new();
1783        parent.insert(RequestId(7));
1784        let child = parent.fork();
1785        child.insert(ConnSocketInfo("a"));
1786
1787        assert_eq!(
1788            AnyOf::from_extensions(&child),
1789            Some(AnyOf::Sock(&ConnSocketInfo("a")))
1790        );
1791
1792        let parent = Extensions::new();
1793        parent.insert(ConnSocketInfo("a"));
1794        let child = parent.fork();
1795        child.insert(RequestId(7));
1796
1797        assert_eq!(
1798            AnyOf::from_extensions(&child),
1799            Some(AnyOf::Req(&RequestId(7)))
1800        );
1801    }
1802}