Switch to reading dictionaries during the fastly_dictionary_open ca…

…ll (#379) Read dictionaries during the `open` host call instead of once for each `get` call. This improves viceroy performance for large dictionaries, and more closely matches the semantics of the production environment, as changes made while the dictionary is open won't be reflected in subsequent `get` calls. Additionally, remove the `DictionaryName` newtype, as it wasn't offering anything over using a `String` directly. Co-authored-by: Jamey Sharp <[email protected]>
fastly · Jun 11, 2024 · 2925fe0 · 2925fe0
1 parent bc77f49
commit 2925fe0
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 74 deletions.
diff --git a/lib/src/component/config_store.rs b/lib/src/component/config_store.rs
@@ -16,10 +16,7 @@ impl config_store::Host for Session {
         name: String,
         max_len: u64,
     ) -> Result<Option<String>, types::Error> {
-        let dict = self
-            .dictionary(store.into())?
-            .contents()
-            .map_err(|err| error::Error::Other(err.into()))?;
+        let dict = &self.dictionary(store.into())?.contents;
 
         let item = if let Some(item) = dict.get(&name) {
             item

diff --git a/lib/src/component/dictionary.rs b/lib/src/component/dictionary.rs
@@ -16,10 +16,7 @@ impl dictionary::Host for Session {
         key: String,
         max_len: u64,
     ) -> Result<Option<String>, types::Error> {
-        let dict = self
-            .dictionary(h.into())?
-            .contents()
-            .map_err(|err| error::Error::Other(err.into()))?;
+        let dict = &self.dictionary(h.into())?.contents;
 
         let item = if let Some(item) = dict.get(&key) {
             item

diff --git a/lib/src/config.rs b/lib/src/config.rs
@@ -21,10 +21,9 @@ mod limits;
 /// Types and deserializers for dictionaries configuration settings.
 mod dictionaries;
 
-pub use self::dictionaries::Dictionary;
-pub use self::dictionaries::DictionaryName;
+pub use self::dictionaries::{Dictionary, LoadedDictionary};
 
-pub type Dictionaries = HashMap<DictionaryName, Dictionary>;
+pub type Dictionaries = HashMap<String, Dictionary>;
 
 /// Types and deserializers for backend configuration settings.
 mod backends;

diff --git a/lib/src/config/dictionaries.rs b/lib/src/config/dictionaries.rs
@@ -2,44 +2,26 @@ use {
     crate::error::DictionaryConfigError,
     std::{
         collections::HashMap,
-        fmt, fs,
+        fs,
         path::{Path, PathBuf},
+        sync::Arc,
     },
 };
 
-#[derive(Clone, Debug, Eq, PartialEq, Hash)]
-pub struct DictionaryName(String);
-
-impl fmt::Display for DictionaryName {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}", self.0)
-    }
-}
-
-impl DictionaryName {
-    pub fn new(name: String) -> Self {
-        Self(name)
-    }
-}
-
 /// A single Dictionary definition.
 ///
 /// A Dictionary consists of a file and format, but more fields may be added in the future.
 #[derive(Clone, Debug)]
 pub enum Dictionary {
-    InlineToml { contents: HashMap<String, String> },
-    Json { file: PathBuf },
+    InlineToml {
+        contents: Arc<HashMap<String, String>>,
+    },
+    Json {
+        file: PathBuf,
+    },
 }
 
 impl Dictionary {
-    /// Returns a reference to the dictionary's contents.
-    pub fn contents(&self) -> Result<HashMap<String, String>, DictionaryConfigError> {
-        match self {
-            Self::InlineToml { contents } => Ok(contents.clone()),
-            Self::Json { file } => Self::read_json_contents(file),
-        }
-    }
-
     /// Returns `true` if this is dictionary uses an external JSON file.
     pub fn is_json(&self) -> bool {
         matches!(self, Self::Json { .. })
@@ -86,11 +68,28 @@ impl Dictionary {
 
         Ok(contents)
     }
+
+    pub fn load(&self) -> Result<LoadedDictionary, DictionaryConfigError> {
+        let contents = match self {
+            Dictionary::InlineToml { contents } => Arc::clone(contents),
+            Dictionary::Json { file } => {
+                let contents = Self::read_json_contents(file)?;
+                Arc::new(contents)
+            }
+        };
+
+        Ok(LoadedDictionary { contents })
+    }
+}
+
+#[derive(Clone)]
+pub struct LoadedDictionary {
+    pub contents: Arc<HashMap<String, String>>,
 }
 
 /// A map of [`Dictionary`] definitions, keyed by their name.
 #[derive(Clone, Debug, Default)]
-pub struct DictionariesConfig(pub HashMap<DictionaryName, Dictionary>);
+pub struct DictionariesConfig(pub HashMap<String, Dictionary>);
 
 /// This module contains [`TryFrom`] implementations used when deserializing a `fastly.toml`.
 ///
@@ -99,12 +98,12 @@ pub struct DictionariesConfig(pub HashMap<DictionaryName, Dictionary>);
 /// not valid, a [`FastlyConfigError`] will be returned.
 mod deserialization {
     use {
-        super::{DictionariesConfig, Dictionary, DictionaryName},
+        super::{DictionariesConfig, Dictionary},
         crate::{
             config::limits::{DICTIONARY_ITEM_KEY_MAX_LEN, DICTIONARY_ITEM_VALUE_MAX_LEN},
             error::{DictionaryConfigError, FastlyConfigError},
         },
-        std::{collections::HashMap, path::PathBuf, str::FromStr},
+        std::{collections::HashMap, path::PathBuf, sync::Arc},
         toml::value::{Table, Value},
         tracing::info,
     };
@@ -129,7 +128,7 @@ mod deserialization {
             fn process_entry(
                 name: &str,
                 entry: Value,
-            ) -> Result<(DictionaryName, Dictionary), DictionaryConfigError> {
+            ) -> Result<(String, Dictionary), DictionaryConfigError> {
                 let mut toml = match entry {
                     Value::Table(table) => table,
                     _ => return Err(DictionaryConfigError::InvalidEntryType),
@@ -154,10 +153,9 @@ mod deserialization {
                     }
                 };
 
-                let name = name.parse()?;
                 check_for_unrecognized_keys(&toml)?;
 
-                Ok((name, dictionary))
+                Ok((name.to_string(), dictionary))
             }
 
             toml.into_iter()
@@ -195,7 +193,9 @@ mod deserialization {
         // Validate that the dictionary adheres to Fastly's API.
         validate_dictionary_contents(&contents)?;
 
-        Ok(Dictionary::InlineToml { contents })
+        Ok(Dictionary::InlineToml {
+            contents: Arc::new(contents),
+        })
     }
 
     fn process_json_dictionary(toml: &mut Table) -> Result<Dictionary, DictionaryConfigError> {
@@ -241,12 +241,4 @@ mod deserialization {
 
         Ok(())
     }
-
-    impl FromStr for DictionaryName {
-        type Err = DictionaryConfigError;
-        fn from_str(name: &str) -> Result<Self, Self::Err> {
-            // We do not do any validation on the name as Config Stores and Dictionaries have different validation rules
-            Ok(Self(name.to_owned()))
-        }
-    }
 }
diff --git a/lib/src/config/unit_tests.rs b/lib/src/config/unit_tests.rs
@@ -1,6 +1,6 @@
 use {
     super::{FastlyConfig, LocalServerConfig, RawLocalServerConfig},
-    crate::{config::DictionaryName, error::FastlyConfigError},
+    crate::error::FastlyConfigError,
     std::{convert::TryInto, fs::File, io::Write},
     tempfile::tempdir,
 };
@@ -132,7 +132,7 @@ fn fastly_toml_files_with_simple_dictionary_configurations_can_be_read() {
 
     let dictionary = config
         .dictionaries()
-        .get(&DictionaryName::new("a".to_string()))
+        .get("a")
         .expect("dictionary configurations can be accessed");
     assert_eq!(dictionary.file_path().unwrap(), file_path);
     assert!(dictionary.is_json());
@@ -171,7 +171,7 @@ fn fastly_toml_files_with_simple_config_store_configurations_can_be_read() {
 
     let dictionary = config
         .dictionaries()
-        .get(&DictionaryName::new("a".to_string()))
+        .get("a")
         .expect("dictionary configurations can be accessed");
     assert_eq!(dictionary.file_path().unwrap(), file_path);
     assert!(dictionary.is_json());

diff --git a/lib/src/session.rs b/lib/src/session.rs
@@ -11,10 +11,7 @@ use {
     self::downstream::DownstreamResponse,
     crate::{
         body::Body,
-        config::{
-            Backend, Backends, DeviceDetection, Dictionaries, Dictionary, DictionaryName,
-            Geolocation,
-        },
+        config::{Backend, Backends, DeviceDetection, Dictionaries, Geolocation, LoadedDictionary},
         error::{Error, HandleError},
         logging::LogEndpoint,
         object_store::{ObjectKey, ObjectStoreError, ObjectStoreKey, ObjectStores},
@@ -96,10 +93,8 @@ pub struct Session {
     ///
     /// Populated prior to guest execution, and never modified.
     dictionaries: Arc<Dictionaries>,
-    /// The dictionaries configured for this execution.
-    ///
-    /// Populated prior to guest execution, and never modified.
-    dictionaries_by_name: PrimaryMap<DictionaryHandle, DictionaryName>,
+    /// The dictionaries that have been opened by the guest.
+    loaded_dictionaries: PrimaryMap<DictionaryHandle, LoadedDictionary>,
     /// The ObjectStore configured for this execution.
     ///
     /// Populated prior to guest execution and can be modified during requests.
@@ -171,7 +166,7 @@ impl Session {
             dynamic_backends: Backends::default(),
             tls_config,
             dictionaries,
-            dictionaries_by_name: PrimaryMap::new(),
+            loaded_dictionaries: PrimaryMap::new(),
             object_store,
             object_store_by_name: PrimaryMap::new(),
             secret_stores,
@@ -603,9 +598,9 @@ impl Session {
 
     /// Look up a dictionary-handle by name.
     pub fn dictionary_handle(&mut self, name: &str) -> Result<DictionaryHandle, Error> {
-        let dict = DictionaryName::new(name.to_string());
-        if self.dictionaries.contains_key(&dict) {
-            Ok(self.dictionaries_by_name.push(dict))
+        if let Some(dict) = self.dictionaries.get(name) {
+            let loaded = dict.load().map_err(|err| Error::Other(err.into()))?;
+            Ok(self.loaded_dictionaries.push(loaded))
         } else {
             Err(Error::DictionaryError(
                 crate::wiggle_abi::DictionaryError::UnknownDictionary(name.to_owned()),
@@ -614,10 +609,9 @@ impl Session {
     }
 
     /// Look up a dictionary by dictionary-handle.
-    pub fn dictionary(&self, handle: DictionaryHandle) -> Result<&Dictionary, HandleError> {
-        self.dictionaries_by_name
+    pub fn dictionary(&self, handle: DictionaryHandle) -> Result<&LoadedDictionary, HandleError> {
+        self.loaded_dictionaries
             .get(handle)
-            .and_then(|name| self.dictionaries.get(name))
             .ok_or(HandleError::InvalidDictionaryHandle(handle))
     }
 

diff --git a/lib/src/wiggle_abi/dictionary_impl.rs b/lib/src/wiggle_abi/dictionary_impl.rs
@@ -45,10 +45,7 @@ impl FastlyDictionary for Session {
         buf: &GuestPtr<u8>,
         buf_len: u32,
     ) -> Result<u32, Error> {
-        let dict = self
-            .dictionary(dictionary)?
-            .contents()
-            .map_err(|err| Error::Other(err.into()))?;
+        let dict = &self.dictionary(dictionary)?.contents;
 
         let item_bytes = {
             let key: &str = &key.as_str()?.ok_or(Error::SharedMemory)?;