apalis_core/task/
metadata.rs

1//! Task metadata extension trait and implementations
2//!
3//! ## Overview
4//! - `Metadata`: A trait for extracting and injecting metadata.
5use crate::task::Task;
6use crate::task::from_request::FromRequest;
7use std::convert::Infallible;
8#[cfg(feature = "tracing")]
9use std::fmt;
10use std::ops::{Deref, Index};
11#[cfg(feature = "tracing")]
12use std::str::FromStr;
13use std::{collections::HashMap, ops::IndexMut};
14
15/// Metadata wrapper for task contexts.
16#[derive(Debug, Clone)]
17pub struct Meta<T>(pub T);
18
19impl<T> Deref for Meta<T> {
20    type Target = T;
21    fn deref(&self) -> &Self::Target {
22        &self.0
23    }
24}
25
26/// A lightweight key-value metadata store backed by a `HashMap`.
27///
28/// `MetadataStore` is designed for storing arbitrary string metadata such as
29/// task attributes, labels, annotations, headers, or contextual information.
30///
31/// Keys are unique within the store. Attempting to insert a duplicate key
32/// returns a [`MetadataError::DuplicateKey`] error.
33///
34/// # Examples
35///
36/// ```rust
37/// # use std::collections::HashMap;
38/// # use apalis_core::task::metadata::Metadata;
39/// # use apalis_core::task::metadata::MetadataStore;
40/// # use apalis_core::task::metadata::MetadataError;
41/// #
42/// let mut metadata = MetadataStore::new();
43///
44/// metadata.insert("request_id", "abc-123")?;
45/// metadata.insert("environment", "production")?;
46///
47/// assert_eq!(
48///     metadata.get("request_id"),
49///     Some("abc-123")
50/// );
51///
52/// assert!(metadata.contains_key("environment"));
53///
54/// # Ok::<(), MetadataError>(())
55/// ```
56#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
57#[derive(Debug, Clone, Default, PartialEq, Eq)]
58pub struct MetadataStore(HashMap<String, String>);
59
60/// Errors returned by [`MetadataStore`].
61#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
62#[non_exhaustive]
63pub enum MetadataError {
64    /// Returned when attempting to insert a key that already exists.
65    #[error("The key already exists in the store")]
66    DuplicateKey(String),
67}
68
69impl MetadataStore {
70    /// Creates an empty [`MetadataStore`].
71    ///
72    /// # Examples
73    ///
74    /// ```rust
75    /// # use apalis_core::task::metadata::Metadata;
76    /// # use apalis_core::task::metadata::MetadataStore;
77    /// let metadata = MetadataStore::new();
78    ///
79    /// assert_eq!(metadata.iter().count(), 0);
80    /// ```
81    #[must_use]
82    pub fn new() -> Self {
83        Self(HashMap::new())
84    }
85
86    /// Inserts a key-value pair into the store.
87    ///
88    /// Returns an error if the key already exists.
89    ///
90    /// # Errors
91    ///
92    /// Returns [`MetadataError::DuplicateKey`] if the provided key is already
93    /// present in the store.
94    ///
95    /// # Examples
96    ///
97    /// ```rust
98    /// # use apalis_core::task::metadata::Metadata;
99    /// # use apalis_core::task::metadata::MetadataStore;
100    /// # use apalis_core::task::metadata::MetadataError;
101    /// let mut metadata = MetadataStore::new();
102    ///
103    /// metadata.insert("region", "us-east-1")?;
104    ///
105    /// assert_eq!(
106    ///     metadata.get("region"),
107    ///     Some("us-east-1")
108    /// );
109    ///
110    /// # Ok::<(), MetadataError>(())
111    /// ```
112    ///
113    /// Duplicate keys are rejected:
114    ///
115    /// ```rust
116    /// # use apalis_core::task::metadata::Metadata;
117    /// # use apalis_core::task::metadata::MetadataStore;
118    /// # use apalis_core::task::metadata::MetadataError;
119    /// let mut metadata = MetadataStore::new();
120    ///
121    /// metadata.insert("service", "api")?;
122    ///
123    /// let err = metadata.insert("service", "worker").unwrap_err();
124    ///
125    /// assert_eq!(
126    ///     err,
127    ///     MetadataError::DuplicateKey("service".to_string())
128    /// );
129    ///
130    /// # Ok::<(), MetadataError>(())
131    /// ```
132    pub fn insert<K, V>(&mut self, key: K, value: V) -> Result<(), MetadataError>
133    where
134        K: Into<String>,
135        V: Into<String>,
136    {
137        let key = key.into();
138
139        if self.0.contains_key(&key) {
140            return Err(MetadataError::DuplicateKey(key));
141        }
142
143        self.0.insert(key, value.into());
144
145        Ok(())
146    }
147
148    /// Returns a reference to the value corresponding to the given key.
149    ///
150    /// Returns `None` if the key does not exist.
151    ///
152    /// # Examples
153    ///
154    /// ```rust
155    /// # use apalis_core::task::metadata::Metadata;
156    /// # use apalis_core::task::metadata::MetadataStore;
157    /// # use apalis_core::task::metadata::MetadataError;
158    /// let mut metadata = MetadataStore::new();
159    ///
160    /// metadata.insert("version", "1.0")?;
161    ///
162    /// assert_eq!(
163    ///     metadata.get("version"),
164    ///     Some("1.0")
165    /// );
166    ///
167    /// assert_eq!(metadata.get("missing"), None);
168    ///
169    /// # Ok::<(), MetadataError>(())
170    /// ```
171    #[must_use]
172    pub fn get(&self, key: &str) -> Option<&str> {
173        self.0.get(key).map(String::as_str)
174    }
175
176    /// Removes a key from the store, returning the stored value if it existed.
177    ///
178    /// # Examples
179    ///
180    /// ```rust
181    /// # use apalis_core::task::metadata::Metadata;
182    /// # use apalis_core::task::metadata::MetadataStore;
183    /// # use apalis_core::task::metadata::MetadataError;
184    /// let mut metadata = MetadataStore::new();
185    ///
186    /// metadata.insert("token", "secret")?;
187    ///
188    /// assert_eq!(
189    ///     metadata.remove("token"),
190    ///     Some("secret".to_string())
191    /// );
192    ///
193    /// assert!(!metadata.contains_key("token"));
194    ///
195    /// # Ok::<(), MetadataError>(())
196    /// ```
197    pub fn remove(&mut self, key: &str) -> Option<String> {
198        self.0.remove(key)
199    }
200
201    /// Returns `true` if the store contains the specified key.
202    ///
203    /// # Examples
204    ///
205    /// ```rust
206    /// # use apalis_core::task::metadata::Metadata;
207    /// # use apalis_core::task::metadata::MetadataStore;
208    /// # use apalis_core::task::metadata::MetadataError;
209    /// let mut metadata = MetadataStore::new();
210    ///
211    /// metadata.insert("owner", "alice")?;
212    ///
213    /// assert!(metadata.contains_key("owner"));
214    /// assert!(!metadata.contains_key("missing"));
215    ///
216    /// # Ok::<(), MetadataError>(())
217    /// ```
218    #[must_use]
219    pub fn contains_key(&self, key: &str) -> bool {
220        self.0.contains_key(key)
221    }
222
223    /// Returns an iterator over all key-value pairs in the store.
224    ///
225    /// The iterator yields `(&String, &String)` pairs.
226    ///
227    /// # Examples
228    ///
229    /// ```rust
230    /// # use apalis_core::task::metadata::Metadata;
231    /// # use apalis_core::task::metadata::MetadataStore;
232    /// # use apalis_core::task::metadata::MetadataError;
233    /// let mut metadata = MetadataStore::new();
234    ///
235    /// metadata.insert("a", "1")?;
236    /// metadata.insert("b", "2")?;
237    ///
238    /// let items: Vec<_> = metadata.iter().collect();
239    ///
240    /// assert_eq!(items.len(), 2);
241    ///
242    /// # Ok::<(), MetadataError>(())
243    /// ```
244    pub fn iter(&self) -> impl Iterator<Item = (&String, &String)> {
245        self.0.iter()
246    }
247
248    /// Consumes the store and returns the underlying `HashMap`.
249    ///
250    /// # Examples
251    ///
252    /// ```rust
253    /// # use apalis_core::task::metadata::Metadata;
254    /// # use apalis_core::task::metadata::MetadataStore;
255    /// # use apalis_core::task::metadata::MetadataError;
256    /// let mut metadata = MetadataStore::new();
257    ///
258    /// metadata.insert("key", "value")?;
259    ///
260    /// let inner = metadata.into_inner();
261    ///
262    /// assert_eq!(
263    ///     inner.get("key"),
264    ///     Some(&"value".to_string())
265    /// );
266    ///
267    /// # Ok::<(), MetadataError>(())
268    /// ```
269    #[must_use]
270    pub fn into_inner(self) -> HashMap<String, String> {
271        self.0
272    }
273
274    /// Get a typed metadata entry.
275    #[must_use = "Extracted metadata should be used or handled"]
276    pub fn extract_as<M: Metadata>(&self) -> Result<M, M::Error> {
277        M::extract(self)
278    }
279
280    /// Create a `MetadataStore` from a `HashMap<String, String>`.
281    #[must_use]
282    pub fn from_map(map: HashMap<String, String>) -> Self {
283        Self(map)
284    }
285
286    /// Returns the number of entries in the store.
287    #[must_use]
288    pub fn len(&self) -> usize {
289        self.0.len()
290    }
291
292    /// Returns `true` if the store contains no entries.
293    #[must_use]
294    pub fn is_empty(&self) -> bool {
295        self.0.is_empty()
296    }
297
298    /// Removes all entries from the store.
299    pub fn clear(&mut self) {
300        self.0.clear();
301    }
302
303    /// Returns a mutable reference to the value corresponding to the key.
304    ///
305    /// Returns `None` if the key does not exist.
306    #[must_use]
307    pub fn get_mut(&mut self, key: &str) -> Option<&mut String> {
308        self.0.get_mut(key)
309    }
310
311    /// Returns a mutable iterator over all key-value pairs in the store.
312    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&String, &mut String)> {
313        self.0.iter_mut()
314    }
315
316    /// Returns an iterator over all keys in the store.
317    pub fn keys(&self) -> impl Iterator<Item = &String> {
318        self.0.keys()
319    }
320
321    /// Returns an iterator over all values in the store.
322    pub fn values(&self) -> impl Iterator<Item = &String> {
323        self.0.values()
324    }
325
326    /// Returns a mutable iterator over all values in the store.
327    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut String> {
328        self.0.values_mut()
329    }
330
331    /// Extends the store with key-value pairs.
332    ///
333    /// Existing keys are overwritten.
334    ///
335    /// Use [`MetadataStore::insert`] when duplicate keys should be rejected.
336    pub fn extend<I, K, V>(&mut self, iter: I)
337    where
338        I: IntoIterator<Item = (K, V)>,
339        K: Into<String>,
340        V: Into<String>,
341    {
342        self.0.extend(
343            iter.into_iter()
344                .map(|(key, value)| (key.into(), value.into())),
345        );
346    }
347
348    /// Returns a reference to the underlying map.
349    #[must_use]
350    pub fn as_map(&self) -> &HashMap<String, String> {
351        &self.0
352    }
353
354    /// Updates an existing value.
355    ///
356    /// Returns `None` if the key does not exist.
357    pub fn replace<K, V>(&mut self, key: K, value: V) -> Option<String>
358    where
359        K: AsRef<str>,
360        V: Into<String>,
361    {
362        self.0.insert(key.as_ref().to_owned(), value.into())
363    }
364
365    /// Inserts a value only if the key does not already exist.
366    pub fn insert_if_absent<K, V>(&mut self, key: K, value: V) -> bool
367    where
368        K: Into<String>,
369        V: Into<String>,
370    {
371        let key = key.into();
372
373        if let std::collections::hash_map::Entry::Vacant(e) = self.0.entry(key) {
374            e.insert(value.into());
375            true
376        } else {
377            false
378        }
379    }
380}
381
382impl<K, V> FromIterator<(K, V)> for MetadataStore
383where
384    K: Into<String>,
385    V: Into<String>,
386{
387    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
388        Self(
389            iter.into_iter()
390                .map(|(key, value)| (key.into(), value.into()))
391                .collect(),
392        )
393    }
394}
395
396impl IntoIterator for MetadataStore {
397    type Item = (String, String);
398    type IntoIter = std::collections::hash_map::IntoIter<String, String>;
399
400    fn into_iter(self) -> Self::IntoIter {
401        self.0.into_iter()
402    }
403}
404
405impl<'a> IntoIterator for &'a MetadataStore {
406    type Item = (&'a String, &'a String);
407    type IntoIter = std::collections::hash_map::Iter<'a, String, String>;
408
409    fn into_iter(self) -> Self::IntoIter {
410        self.0.iter()
411    }
412}
413
414impl<'a> IntoIterator for &'a mut MetadataStore {
415    type Item = (&'a String, &'a mut String);
416    type IntoIter = std::collections::hash_map::IterMut<'a, String, String>;
417
418    fn into_iter(self) -> Self::IntoIter {
419        self.0.iter_mut()
420    }
421}
422
423impl AsRef<HashMap<String, String>> for MetadataStore {
424    fn as_ref(&self) -> &HashMap<String, String> {
425        &self.0
426    }
427}
428
429impl AsMut<HashMap<String, String>> for MetadataStore {
430    fn as_mut(&mut self) -> &mut HashMap<String, String> {
431        &mut self.0
432    }
433}
434
435impl Index<&str> for MetadataStore {
436    type Output = String;
437
438    fn index(&self, key: &str) -> &Self::Output {
439        &self.0[key]
440    }
441}
442
443impl IndexMut<&str> for MetadataStore {
444    fn index_mut(&mut self, key: &str) -> &mut Self::Output {
445        self.0
446            .get_mut(key)
447            .unwrap_or_else(|| panic!("metadata key `{key}` does not exist"))
448    }
449}
450
451/// Implemented by types that can be stored as metadata.
452/// Provides a stable key and string-based serialization.
453pub trait Metadata: Sized {
454    /// The error produced when extracting the Metadata
455    type Error;
456
457    /// Extract `Metadata` from the store
458    fn extract(store: &MetadataStore) -> Result<Self, Self::Error>;
459
460    /// Inject [`Self`] into the store
461    fn inject(&self, map: &mut MetadataStore) -> Result<(), Self::Error>;
462}
463
464impl<T: Metadata, Args: Send + Sync> FromRequest<Task<Args>> for Meta<T> {
465    type Error = T::Error;
466
467    async fn from_request(task: &Task<Args>) -> Result<Self, Self::Error> {
468        let metadata = task.metadata();
469        let value = T::extract(metadata)?;
470        Ok(Self(value))
471    }
472}
473
474impl<Args: Send + Sync> FromRequest<Task<Args>> for MetadataStore {
475    type Error = Infallible;
476
477    async fn from_request(task: &Task<Args>) -> Result<Self, Self::Error> {
478        Ok(task.metadata().clone())
479    }
480}
481
482/// Metadata used specifically for storing tracing context.
483#[cfg(feature = "tracing")]
484#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
485#[derive(Debug, Default, Clone)]
486pub struct TracingContext {
487    trace_id: Option<String>,
488    span_id: Option<String>,
489    trace_flags: Option<u8>,
490    trace_state: Option<String>,
491}
492
493#[cfg(feature = "tracing")]
494impl TracingContext {
495    /// Create a new empty `TracingContext`.
496    #[must_use]
497    pub fn new() -> Self {
498        Self::default()
499    }
500
501    /// Set the trace ID.
502    #[must_use]
503    pub fn with_trace_id(mut self, trace_id: impl Into<String>) -> Self {
504        self.trace_id = Some(trace_id.into());
505        self
506    }
507
508    /// Set the span ID.
509    #[must_use]
510    pub fn with_span_id(mut self, span_id: impl Into<String>) -> Self {
511        self.span_id = Some(span_id.into());
512        self
513    }
514
515    /// Set the trace flags.
516    #[must_use]
517    pub fn with_trace_flags(mut self, trace_flags: u8) -> Self {
518        self.trace_flags = Some(trace_flags);
519        self
520    }
521
522    /// Set the trace state.
523    #[must_use]
524    pub fn with_trace_state(mut self, trace_state: impl Into<String>) -> Self {
525        self.trace_state = Some(trace_state.into());
526        self
527    }
528
529    /// Get the trace ID.
530    #[must_use]
531    pub fn trace_id(&self) -> &Option<String> {
532        &self.trace_id
533    }
534
535    /// Get the span ID.
536    #[must_use]
537    pub fn span_id(&self) -> &Option<String> {
538        &self.span_id
539    }
540
541    /// Get the trace flags.
542    #[must_use]
543    pub fn trace_flags(&self) -> &Option<u8> {
544        &self.trace_flags
545    }
546
547    /// Get the trace state.
548    #[must_use]
549    pub fn trace_state(&self) -> &Option<String> {
550        &self.trace_state
551    }
552}
553
554#[cfg(feature = "tracing")]
555/// Error provided by parsing TracingContext
556#[derive(Debug, thiserror::Error)]
557#[non_exhaustive]
558pub enum TracingContextParseError {
559    /// Missing Field
560    #[error("Missing Field: {0}")]
561    MissingField(&'static str),
562    /// Invalid flags
563    #[error("Invalid flags: {0}")]
564    InvalidTraceFlags(std::num::ParseIntError),
565    /// Invalid Format
566    #[error("Invalid Format")]
567    InvalidFormat,
568    /// Key {apalis_core.tracing.context} not found in Metadata
569    #[error("Key {{apalis_core.tracing.context}} not found in Metadata")]
570    MissingKey,
571    /// Duplicate entry
572    #[error("Duplicate entry: {0}")]
573    DuplicateEntry(#[from] MetadataError),
574}
575
576// Serialization format: a single W3C traceparent-style string.
577//
578//   <trace_id>;<span_id>;<trace_flags>;<trace_state>
579//
580// Each field is either its value or `-` if None.
581//
582// Example:
583//   "4bf92f3577b34da6a3ce929d0e0e4736;00f067aa0ba902b7;01;congo=t61rcWkgMzE"
584//   "4bf92f3577b34da6a3ce929d0e0e4736;00f067aa0ba902b7;-;-"
585#[cfg(feature = "tracing")]
586impl fmt::Display for TracingContext {
587    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
588        write!(
589            f,
590            "{};{};{};{}",
591            self.trace_id.as_deref().unwrap_or("-"),
592            self.span_id.as_deref().unwrap_or("-"),
593            self.trace_flags
594                .map(|v| v.to_string())
595                .as_deref()
596                .unwrap_or("-"),
597            self.trace_state.as_deref().unwrap_or("-"),
598        )
599    }
600}
601
602#[cfg(feature = "tracing")]
603impl FromStr for TracingContext {
604    type Err = TracingContextParseError;
605
606    fn from_str(s: &str) -> Result<Self, Self::Err> {
607        let mut parts = s.splitn(4, ';');
608
609        let mut next = |field| {
610            parts
611                .next()
612                .ok_or(TracingContextParseError::MissingField(field))
613        };
614
615        let trace_id = match next("trace_id")? {
616            "-" => None,
617            v => Some(v.to_owned()),
618        };
619
620        let span_id = match next("span_id")? {
621            "-" => None,
622            v => Some(v.to_owned()),
623        };
624
625        let trace_flags = match next("trace_flags")? {
626            "-" => None,
627            v => Some(
628                v.parse::<u8>()
629                    .map_err(TracingContextParseError::InvalidTraceFlags)?,
630            ),
631        };
632
633        let trace_state = match next("trace_state")? {
634            "-" => None,
635            v => Some(v.to_owned()),
636        };
637
638        Ok(Self {
639            trace_id,
640            span_id,
641            trace_flags,
642            trace_state,
643        })
644    }
645}
646
647#[cfg(feature = "tracing")]
648const TRACING_CONTENT_KEY: &str = "apalis_core.tracing.context";
649
650#[cfg(feature = "tracing")]
651impl Metadata for TracingContext {
652    type Error = TracingContextParseError;
653    fn extract(store: &MetadataStore) -> Result<Self, Self::Error> {
654        store
655            .get(TRACING_CONTENT_KEY)
656            .ok_or(TracingContextParseError::InvalidFormat)?
657            .parse()
658    }
659
660    fn inject(&self, map: &mut MetadataStore) -> Result<(), Self::Error> {
661        Ok(map.insert(TRACING_CONTENT_KEY, self.to_string())?)
662    }
663}
664
665#[cfg(test)]
666#[allow(unused)]
667mod tests {
668    use std::{convert::Infallible, fmt::Debug, num::ParseIntError, task::Poll, time::Duration};
669
670    use crate::{
671        error::BoxDynError,
672        task::from_request::FromRequest,
673        task::{
674            Task,
675            metadata::{Meta, Metadata, MetadataStore},
676        },
677    };
678    use futures_core::future::BoxFuture;
679    use tower::Service;
680
681    #[derive(Debug, Clone)]
682    struct ExampleService<S> {
683        service: S,
684    }
685    #[derive(Debug, Clone, Default)]
686    struct ExampleConfig {
687        timeout: Duration,
688    }
689
690    const EXAMPLE_CONFIG: &str = "apalis_core.example.config";
691
692    impl Metadata for ExampleConfig {
693        type Error = ParseIntError;
694        fn extract(store: &MetadataStore) -> Result<Self, Self::Error> {
695            let timeout = store
696                .get(EXAMPLE_CONFIG)
697                .unwrap()
698                .parse::<u64>()
699                .map(Duration::from_secs)?;
700            Ok(ExampleConfig { timeout })
701        }
702
703        fn inject(&self, map: &mut MetadataStore) -> Result<(), ParseIntError> {
704            let value = self.timeout.as_secs().to_string();
705            map.insert(EXAMPLE_CONFIG, value).unwrap();
706            Ok(())
707        }
708    }
709
710    impl<S, Args: Send + Sync + 'static> Service<Task<Args>> for ExampleService<S>
711    where
712        S: Service<Task<Args>> + Clone + Send + 'static,
713        S::Future: Send + 'static,
714    {
715        type Response = S::Response;
716        type Error = S::Error;
717        type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
718
719        fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
720            self.service.poll_ready(cx)
721        }
722
723        fn call(&mut self, request: Task<Args>) -> Self::Future {
724            let mut svc = self.service.clone();
725
726            // Do something with config
727            Box::pin(async move {
728                let _config: Meta<ExampleConfig> = request.extract().await.unwrap();
729                svc.call(request).await
730            })
731        }
732    }
733
734    #[cfg(feature = "tracing")]
735    #[test]
736    fn tracing_context_keeps_explicit_fields() {
737        let context = crate::task::metadata::TracingContext::new()
738            .with_trace_id("4bf92f3577b34da6a3ce929d0e0e4736")
739            .with_span_id("00f067aa0ba902b7")
740            .with_trace_flags(1)
741            .with_trace_state("vendor=acme");
742
743        assert_eq!(
744            context.trace_id(),
745            &Some("4bf92f3577b34da6a3ce929d0e0e4736".to_string())
746        );
747        assert_eq!(context.span_id(), &Some("00f067aa0ba902b7".to_string()));
748        assert_eq!(context.trace_flags(), &Some(1));
749        assert_eq!(context.trace_state(), &Some("vendor=acme".to_string()));
750    }
751}
apalis_core/task/metadata.rs

apalis_core/task/
metadata.rs
