Context as type map (#504)

* Context as type map

* Added header injector policy

* Rylev's suggestions and clippy

* removed iterator, pub internal HeaderMap

Author: MindFlavor

sdk/core/src/context.rs

@@ -1,11 +1,175 @@
+use std::any::{Any, TypeId};
+use std::collections::HashMap;
+use std::sync::Arc;
+
 /// Pipeline execution context.
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct Context {
-    _priv: (),
+    type_map: HashMap<TypeId, Arc<dyn Any + Send + Sync>>,
+}
+
+impl Default for Context {
+    fn default() -> Self {
+        Self::new()
+    }
 }
 
 impl Context {
+    /// Creates a new, empty Context.
     pub fn new() -> Self {
-        Self::default()
+        Self {
+            type_map: HashMap::new(),
+        }
+    }
+
+    /// Inserts or replaces an entity in the type map. If an entity with the same type was displaced
+    /// by the insert, it will be returned to the caller.
+    pub fn insert_or_replace<E>(&mut self, entity: E) -> Option<Arc<E>>
+    where
+        E: Send + Sync + 'static,
+    {
+        // we make sure that for every TypeId of E as key we ALWAYS retrieve an Option<Arc<E>>. That's why
+        // the `unwrap` below is safe.
+        self.type_map
+            .insert(TypeId::of::<E>(), Arc::new(entity))
+            .map(|displaced| displaced.downcast().unwrap())
+    }
+
+    /// Inserts an entity in the type map. If the an entity with the same type signature is
+    /// already present it will be silently dropped. This function returns a mutable reference to
+    /// the same Context so it can be chained to itself.
+    pub fn insert<E>(&mut self, entity: E) -> &mut Self
+    where
+        E: Send + Sync + 'static,
+    {
+        self.type_map.insert(TypeId::of::<E>(), Arc::new(entity));
+
+        self
+    }
+
+    /// Removes an entity from the type map. If present, the entity will be returned.
+    pub fn remove<E>(&mut self) -> Option<Arc<E>>
+    where
+        E: Send + Sync + 'static,
+    {
+        self.type_map
+            .remove(&TypeId::of::<E>())
+            .map(|removed| removed.downcast().unwrap())
+    }
+
+    /// Returns a reference of the entity of the specified type signature, if it exists.
+    ///
+    /// If there is no entity with the specific type signature, `None` is returned instead.
+    pub fn get<E>(&self) -> Option<&E>
+    where
+        E: Send + Sync + 'static,
+    {
+        self.type_map
+            .get(&TypeId::of::<E>())
+            .map(|item| item.downcast_ref())
+            .flatten()
+    }
+
+    /// Returns the number of entities in the type map.
+    pub fn len(&self) -> usize {
+        self.type_map.len()
+    }
+
+    /// Returns `true` if the type map is empty, `false` otherwise.
+    pub fn is_empty(&self) -> bool {
+        self.type_map.is_empty()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Mutex;
+
+    #[test]
+    fn insert_get_string() {
+        let mut context = Context::new();
+        context.insert_or_replace("pollo".to_string());
+        assert_eq!(Some(&"pollo".to_string()), context.get());
+    }
+
+    #[test]
+    fn insert_get_custom_structs() {
+        #[derive(Debug, PartialEq, Eq)]
+        struct S1 {}
+        #[derive(Debug, PartialEq, Eq)]
+        struct S2 {}
+
+        let mut context = Context::new();
+        context.insert_or_replace(S1 {});
+        context.insert_or_replace(S2 {});
+
+        assert_eq!(Some(Arc::new(S1 {})), context.insert_or_replace(S1 {}));
+        assert_eq!(Some(Arc::new(S2 {})), context.insert_or_replace(S2 {}));
+
+        assert_eq!(Some(&S1 {}), context.get());
+        assert_eq!(Some(&S2 {}), context.get());
+    }
+
+    #[test]
+    fn insert_fluent_syntax() {
+        #[derive(Debug, PartialEq, Eq, Default)]
+        struct S1 {}
+        #[derive(Debug, PartialEq, Eq, Default)]
+        struct S2 {}
+
+        let mut context = Context::new();
+
+        context
+            .insert("static str")
+            .insert("a String".to_string())
+            .insert(S1::default())
+            .insert(S1::default()) // notice we are REPLACING S1. This call will *not* increment the counter
+            .insert(S2::default());
+
+        assert_eq!(4, context.len());
+        assert_eq!(Some(&"static str"), context.get());
+    }
+
+    fn require_send_sync<T: Send + Sync>(_: &T) {}
+
+    #[test]
+    fn test_require_send_sync() {
+        // this won't compile if Context as a whole is not Send + Sync
+        require_send_sync(&Context::new())
+    }
+
+    #[test]
+    fn mutability() {
+        #[derive(Debug, PartialEq, Eq, Default)]
+        struct S1 {
+            num: u8,
+        }
+        let mut context = Context::new();
+        context.insert_or_replace(Mutex::new(S1::default()));
+
+        // the stored value is 0.
+        assert_eq!(0, context.get::<Mutex<S1>>().unwrap().lock().unwrap().num);
+
+        // we change the number to 42 in a thread safe manner.
+        context.get::<Mutex<S1>>().unwrap().lock().unwrap().num = 42;
+
+        // now the number is 42.
+        assert_eq!(42, context.get::<Mutex<S1>>().unwrap().lock().unwrap().num);
+
+        // we replace the struct with a new one.
+        let displaced = context
+            .insert_or_replace(Mutex::new(S1::default()))
+            .unwrap();
+
+        // the displaced struct still holds 42 as number
+        assert_eq!(42, displaced.lock().unwrap().num);
+
+        // the new struct has 0 has number.
+        assert_eq!(0, context.get::<Mutex<S1>>().unwrap().lock().unwrap().num);
+
+        context.insert_or_replace(Mutex::new(33u32));
+        *context.get::<Mutex<u32>>().unwrap().lock().unwrap() = 42;
+        assert_eq!(42, *context.get::<Mutex<u32>>().unwrap().lock().unwrap());
     }
 }

sdk/core/src/pipeline.rs

@@ -1,6 +1,6 @@
 #[cfg(not(target_arch = "wasm32"))]
 use crate::policies::TransportPolicy;
-use crate::policies::{Policy, TelemetryPolicy};
+use crate::policies::{CustomHeadersInjectorPolicy, Policy, TelemetryPolicy};
 use crate::{ClientOptions, Error, HttpClient, PipelineContext, Request, Response};
 use std::sync::Arc;
 
@@ -70,6 +70,8 @@ where
         let telemetry_policy = TelemetryPolicy::new(crate_name, crate_version, &options.telemetry);
         pipeline.push(Arc::new(telemetry_policy));
 
+        pipeline.push(Arc::new(CustomHeadersInjectorPolicy::default()));
+
         let retry_policy = options.retry.to_policy();
         pipeline.push(retry_policy);
 

sdk/core/src/policies/custom_headers_injector_policy.rs

@@ -0,0 +1,47 @@
+use crate::policies::{Policy, PolicyResult};
+use crate::{PipelineContext, Request, Response};
+use http::header::HeaderMap;
+use std::sync::Arc;
+
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct CustomHeaders(pub HeaderMap);
+
+impl From<HeaderMap> for CustomHeaders {
+    fn from(header_map: HeaderMap) -> Self {
+        Self(header_map)
+    }
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct CustomHeadersInjectorPolicy {}
+
+#[async_trait::async_trait]
+impl<C> Policy<C> for CustomHeadersInjectorPolicy
+where
+    C: Send + Sync,
+{
+    async fn send(
+        &self,
+        ctx: &mut PipelineContext<C>,
+        request: &mut Request,
+        next: &[Arc<dyn Policy<C>>],
+    ) -> PolicyResult<Response> {
+        if let Some(CustomHeaders(custom_headers)) = ctx.get_inner_context().get::<CustomHeaders>()
+        {
+            custom_headers
+                .iter()
+                .for_each(|(header_name, header_value)| {
+                    log::trace!(
+                        "injecting custom context header {:?} with value {:?}",
+                        header_name,
+                        header_value
+                    );
+                    request
+                        .headers_mut()
+                        .insert(header_name, header_value.to_owned());
+                });
+        }
+
+        next[0].send(ctx, request, &next[1..]).await
+    }
+}

sdk/core/src/policies/mod.rs

@@ -1,3 +1,4 @@
+mod custom_headers_injector_policy;
 #[cfg(feature = "mock_transport_framework")]
 mod mock_transport_player_policy;
 #[cfg(feature = "mock_transport_framework")]
@@ -7,6 +8,7 @@ mod telemetry_policy;
 mod transport;
 
 use crate::{PipelineContext, Request, Response};
+pub use custom_headers_injector_policy::{CustomHeaders, CustomHeadersInjectorPolicy};
 #[cfg(feature = "mock_transport_framework")]
 pub use mock_transport_player_policy::MockTransportPlayerPolicy;
 #[cfg(feature = "mock_transport_framework")]

sdk/core/src/prelude.rs

@@ -1,4 +1,5 @@
 pub use crate::etag::Etag;
+pub use crate::policies::CustomHeaders;
 pub use crate::request_options::*;
 pub use crate::{
     new_http_client, AddAsHeader, AppendToUrlQuery, Context, HttpClient, RequestId, SessionToken,

sdk/cosmos/examples/get_database.rs

@@ -1,5 +1,6 @@
 use azure_core::prelude::*;
 use azure_cosmos::prelude::*;
+use http::{HeaderMap, HeaderValue};
 use std::error::Error;
 
 #[tokio::main]
@@ -24,8 +25,19 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
 
     let database_client = client.into_database_client(database_name.clone());
 
+    let mut context = Context::new();
+
+    // Next we create a CustomHeaders type and insert it into the context allowing us to insert custom headers.
+    let custom_headers: CustomHeaders = {
+        let mut custom_headers = HeaderMap::new();
+        custom_headers.insert("MyCoolHeader", HeaderValue::from_static("CORS maybe?"));
+        custom_headers.into()
+    };
+
+    context.insert(custom_headers);
+
     let response = database_client
-        .get_database(Context::new(), GetDatabaseOptions::new())
+        .get_database(context, GetDatabaseOptions::new())
         .await?;
     println!("response == {:?}", response);