From 23372fc53574ef06e852c7c5a745714f8fc685a9 Mon Sep 17 00:00:00 2001
From: glendc <contact@glendc.com>
Date: Thu, 22 Aug 2024 14:21:40 +0200
Subject: [PATCH] support maybe_insert for Extensions/Context

---
 examples/http_connect_proxy.rs    |  6 ++----
 src/proxy/proxydb/layer.rs        | 20 +++++---------------
 src/service/context/extensions.rs | 10 ++++++++++
 src/service/context/mod.rs        | 10 ++++++++++
 4 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/examples/http_connect_proxy.rs b/examples/http_connect_proxy.rs
index 9c256f20..39771652 100644
--- a/examples/http_connect_proxy.rs
+++ b/examples/http_connect_proxy.rs
@@ -257,10 +257,8 @@ impl UsernameLabelParser for PriorityUsernameLabelParser {
         UsernameLabelState::Used
     }
 
-    fn build(mut self, ext: &mut Extensions) -> Result<(), Self::Error> {
-        if let Some(priority) = self.priority.take() {
-            ext.insert(priority);
-        }
+    fn build(self, ext: &mut Extensions) -> Result<(), Self::Error> {
+        ext.maybe_insert(self.priority);
         Ok(())
     }
 }
diff --git a/src/proxy/proxydb/layer.rs b/src/proxy/proxydb/layer.rs
index a8690c47..6c302074 100644
--- a/src/proxy/proxydb/layer.rs
+++ b/src/proxy/proxydb/layer.rs
@@ -1046,9 +1046,7 @@ mod tests {
             ),
         ] {
             let mut ctx = Context::default();
-            if let Some(filter) = filter {
-                ctx.insert(filter);
-            }
+            ctx.maybe_insert(filter);
 
             let maybe_proxy_address = service.serve(ctx, req).await.unwrap();
 
@@ -1085,9 +1083,7 @@ mod tests {
             let mut seen_addresses = Vec::new();
             for _ in 0..5000 {
                 let mut ctx = Context::default();
-                if let Some(filter) = filter.clone() {
-                    ctx.insert(filter);
-                }
+                ctx.maybe_insert(filter.clone());
 
                 let req = Request::builder()
                     .version(req_info.0)
@@ -1142,9 +1138,7 @@ mod tests {
             let mut seen_addresses = Vec::new();
             for _ in 0..5000 {
                 let mut ctx = Context::default();
-                if let Some(filter) = filter.clone() {
-                    ctx.insert(filter);
-                }
+                ctx.maybe_insert(filter.clone());
 
                 let req = Request::builder()
                     .version(req_info.0)
@@ -1231,9 +1225,7 @@ mod tests {
             ),
         ] {
             let mut ctx = Context::default();
-            if let Some(filter) = filter {
-                ctx.insert(filter);
-            }
+            ctx.maybe_insert(filter.clone());
 
             let proxy_address_result = service.serve(ctx, req).await;
             match expected_address {
@@ -1332,9 +1324,7 @@ mod tests {
             ),
         ] {
             let mut ctx = Context::default();
-            if let Some(filter) = filter {
-                ctx.insert(filter);
-            }
+            ctx.maybe_insert(filter);
 
             let proxy_result = service.serve(ctx, req).await;
             match expected {
diff --git a/src/service/context/extensions.rs b/src/service/context/extensions.rs
index 86f3c1c1..4909359b 100644
--- a/src/service/context/extensions.rs
+++ b/src/service/context/extensions.rs
@@ -56,6 +56,16 @@ impl Extensions {
             .and_then(|boxed| boxed.into_any().downcast().ok().map(|boxed| *boxed))
     }
 
+    /// Insert a type only into this `Extensions`, if the value is `Some(T)`.
+    ///
+    /// See [`Self::insert`] for more information.
+    pub fn maybe_insert<T: Clone + Send + Sync + 'static>(
+        &mut self,
+        mut val: Option<T>,
+    ) -> Option<T> {
+        val.take().and_then(|val| self.insert(val))
+    }
+
     /// Extend these extensions with another Extensions.
     pub fn extend(&mut self, other: Extensions) {
         if let Some(other_map) = other.map {
diff --git a/src/service/context/mod.rs b/src/service/context/mod.rs
index 4e8a30d2..52eaad34 100644
--- a/src/service/context/mod.rs
+++ b/src/service/context/mod.rs
@@ -437,6 +437,16 @@ impl<S> Context<S> {
         self.extensions.insert(extension)
     }
 
+    /// Insert a type only into this [`Context`], if the extension is `Some(T)`.
+    ///
+    /// See [`Self::insert`] for more information.
+    pub fn maybe_insert<T: Clone + Send + Sync + 'static>(
+        &mut self,
+        extension: Option<T>,
+    ) -> Option<T> {
+        self.extensions.maybe_insert(extension)
+    }
+
     /// Return the entire dynamic state of the [`Context`] by reference.
     ///
     /// Useful only in case you have a function which works with [`Extensions`] rather