fix: improve performance for EnvIO (#1574)

Co-authored-by: Tushar Mathur <[email protected]>

aws-lambda/src/runtime.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::sync::Arc;
 
 use anyhow::anyhow;
@@ -12,10 +13,10 @@ use crate::http::init_http;
 pub struct LambdaEnv;
 
 impl EnvIO for LambdaEnv {
-    fn get(&self, key: &str) -> Option<String> {
+    fn get(&self, key: &str) -> Option<Cow<'_, str>> {
         // AWS Lambda sets environment variables
         // as real env vars in the runtime.
-        std::env::var(key).ok()
+        std::env::var(key).ok().map(Cow::from)
     }
 }
 

benches/data_loader_bench.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::collections::BTreeSet;
 use std::num::NonZeroU64;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -30,7 +31,7 @@ impl HttpIO for MockHttpClient {
 struct Env {}
 #[async_trait]
 impl EnvIO for Env {
-    fn get(&self, _: &str) -> Option<String> {
+    fn get(&self, _: &str) -> Option<Cow<'_, str>> {
         unimplemented!("Not needed for this bench")
     }
 }

benches/impl_path_string_for_evaluation_context.rs

@@ -81,7 +81,7 @@ impl HttpIO for Http {
 struct Env {}
 
 impl EnvIO for Env {
-    fn get(&self, _: &str) -> Option<String> {
+    fn get(&self, _: &str) -> Option<Cow<'_, str>> {
         unimplemented!("Not needed for this bench")
     }
 }

cloudflare/src/env.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::rc::Rc;
 
 use tailcall::EnvIO;
@@ -11,8 +12,8 @@ unsafe impl Send for CloudflareEnv {}
 unsafe impl Sync for CloudflareEnv {}
 
 impl EnvIO for CloudflareEnv {
-    fn get(&self, key: &str) -> Option<String> {
-        self.env.var(key).ok().map(|s| s.to_string())
+    fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+        self.env.var(key).ok().map(|v| Cow::from(v.to_string()))
     }
 }
 

cloudflare/src/file.rs

@@ -13,10 +13,8 @@ pub struct CloudflareFileIO {
 }
 
 impl CloudflareFileIO {
-    pub fn init(env: Rc<Env>, bucket_id: String) -> anyhow::Result<Self> {
-        let bucket = env
-            .bucket(bucket_id.as_str())
-            .map_err(|e| anyhow!(e.to_string()))?;
+    pub fn init(env: Rc<Env>, bucket_id: &str) -> anyhow::Result<Self> {
+        let bucket = env.bucket(bucket_id).map_err(|e| anyhow!(e.to_string()))?;
         let bucket = Rc::new(bucket);
         Ok(CloudflareFileIO { bucket })
     }

cloudflare/src/runtime.rs

@@ -12,7 +12,7 @@ fn init_env(env: Rc<worker::Env>) -> Arc<dyn EnvIO> {
     Arc::new(env::CloudflareEnv::init(env))
 }
 
-fn init_file(env: Rc<worker::Env>, bucket_id: String) -> anyhow::Result<Arc<dyn FileIO>> {
+fn init_file(env: Rc<worker::Env>, bucket_id: &str) -> anyhow::Result<Arc<dyn FileIO>> {
     Ok(Arc::new(file::CloudflareFileIO::init(env, bucket_id)?))
 }
 
@@ -35,7 +35,7 @@ pub fn init(env: Rc<worker::Env>) -> anyhow::Result<TargetRuntime> {
         http: http.clone(),
         http2_only: http.clone(),
         env: init_env(env.clone()),
-        file: init_file(env.clone(), bucket_id)?,
+        file: init_file(env.clone(), &bucket_id)?,
         cache: init_cache(env),
         extensions: Arc::new(vec![]),
     })

src/cli/runtime/env.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::collections::HashMap;
 
 use crate::EnvIO;
@@ -8,8 +9,8 @@ pub struct EnvNative {
 }
 
 impl EnvIO for EnvNative {
-    fn get(&self, key: &str) -> Option<String> {
-        self.vars.get(key).cloned()
+    fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+        self.vars.get(key).map(Cow::from)
     }
 }
 
@@ -35,7 +36,7 @@ mod tests {
         vars.insert("EXISTING_VAR".to_string(), "value".to_string());
         let test_env = EnvNative { vars };
         let result = test_env.get("EXISTING_VAR");
-        assert_eq!(result, Some("value".to_string()));
+        assert_eq!(result, Some("value".into()));
     }
 
     #[test]

src/config/reader_context.rs

@@ -19,7 +19,7 @@ impl<'a> PathString for ConfigReaderContext<'a> {
         path.split_first()
             .and_then(|(head, tail)| match head.as_ref() {
                 "vars" => self.vars.get(tail[0].as_ref()).map(|v| v.into()),
-                "env" => self.env.get(tail[0].as_ref()).map(|v| v.into()),
+                "env" => self.env.get(tail[0].as_ref()),
                 _ => None,
             })
     }

src/lambda/evaluation_context.rs

@@ -86,7 +86,7 @@ impl<'a, Ctx: ResolverContextLike<'a>> EvaluationContext<'a, Ctx> {
         value.to_str().ok()
     }
 
-    pub fn env_var(&self, key: &str) -> Option<String> {
+    pub fn env_var(&self, key: &str) -> Option<Cow<'_, str>> {
         self.request_ctx.runtime.env.get(key)
     }
 

src/lib.rs

@@ -33,14 +33,15 @@ pub mod tracing;
 pub mod try_fold;
 pub mod valid;
 
+use std::borrow::Cow;
 use std::hash::Hash;
 use std::num::NonZeroU64;
 
 use async_graphql_value::ConstValue;
 use http::Response;
 
 pub trait EnvIO: Send + Sync + 'static {
-    fn get(&self, key: &str) -> Option<String>;
+    fn get(&self, key: &str) -> Option<Cow<'_, str>>;
 }
 
 #[async_trait::async_trait]
@@ -94,8 +95,8 @@ pub mod tests {
     pub struct TestEnvIO(HashMap<String, String>);
 
     impl EnvIO for TestEnvIO {
-        fn get(&self, key: &str) -> Option<String> {
-            self.0.get(key).cloned()
+        fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+            self.0.get(key).map(Cow::from)
         }
     }
 

src/path.rs

@@ -72,7 +72,7 @@ impl<'a, Ctx: ResolverContextLike<'a>> PathString for EvaluationContext<'a, Ctx>
                 "args" => convert_value(ctx.path_arg(tail)?),
                 "headers" => ctx.header(tail[0].as_ref()).map(|v| v.into()),
                 "vars" => ctx.var(tail[0].as_ref()).map(|v| v.into()),
-                "env" => ctx.env_var(tail[0].as_ref()).map(|v| v.into()),
+                "env" => ctx.env_var(tail[0].as_ref()),
                 _ => None,
             })
     }
@@ -123,8 +123,8 @@ mod tests {
         }
 
         impl EnvIO for Env {
-            fn get(&self, key: &str) -> Option<String> {
-                self.env.get(key).cloned()
+            fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+                self.env.get(key).map(Cow::from)
             }
         }
 

src/runtime.rs

@@ -36,6 +36,7 @@ impl TargetRuntime {
 
 #[cfg(test)]
 pub mod test {
+    use std::borrow::Cow;
     use std::collections::HashMap;
     use std::sync::Arc;
     use std::time::Duration;
@@ -152,8 +153,8 @@ pub mod test {
     }
 
     impl EnvIO for TestEnvIO {
-        fn get(&self, key: &str) -> Option<String> {
-            self.vars.get(key).cloned()
+        fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+            self.vars.get(key).map(Cow::from)
         }
     }
 

tests/execution_spec.rs

@@ -2,6 +2,7 @@ extern crate core;
 
 mod telemetry;
 
+use std::borrow::Cow;
 use std::collections::{BTreeMap, HashMap};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
@@ -38,6 +39,7 @@ use url::Url;
 
 #[cfg(test)]
 pub mod test {
+    use std::borrow::Cow;
     use std::collections::HashMap;
     use std::sync::Arc;
     use std::time::Duration;
@@ -154,8 +156,8 @@ pub mod test {
     }
 
     impl EnvIO for TestEnvIO {
-        fn get(&self, key: &str) -> Option<String> {
-            self.vars.get(key).cloned()
+        fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+            self.vars.get(key).map(Cow::from)
         }
     }
 
@@ -238,8 +240,8 @@ pub struct Env {
 }
 
 impl EnvIO for Env {
-    fn get(&self, key: &str) -> Option<String> {
-        self.env.get(key).cloned()
+    fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+        self.env.get(key).map(Cow::from)
     }
 }
 

tests/server_spec.rs

@@ -2,6 +2,7 @@ use tailcall::{blueprint, EnvIO, FileIO, HttpIO};
 
 #[cfg(test)]
 pub mod test {
+    use std::borrow::Cow;
     use std::collections::HashMap;
     use std::sync::Arc;
     use std::time::Duration;
@@ -118,8 +119,8 @@ pub mod test {
     }
 
     impl EnvIO for TestEnvIO {
-        fn get(&self, key: &str) -> Option<String> {
-            self.vars.get(key).cloned()
+        fn get(&self, key: &str) -> Option<Cow<'_, str>> {
+            self.vars.get(key).map(Cow::from)
         }
     }