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