Attempt to replace GraphMap’s Map with a CLruCache

Fixes issue trustification#1218

Author: chirino

Cargo.lock

@@ -54,6 +54,7 @@ bytes = "1.5"
 bytesize = "1.3"
 chrono = { version = "0.4.35", default-features = false }
 clap = "4"
+clru = "0.6.2"
 concat-idents = "1"
 cpe = "0.1.5"
 criterion = "0.5.1"

Cargo.toml

@@ -15,6 +15,7 @@ actix-web = { workspace = true }
 anyhow = { workspace = true }
 cpe = { workspace = true }
 log = { workspace = true }
+clru = { workspace = true }
 parking_lot = { workspace = true }
 petgraph = { workspace = true }
 sea-orm = { workspace = true }

modules/analysis/Cargo.toml

@@ -1,10 +1,13 @@
 use petgraph::Graph;
 use serde::Serialize;
 use std::{
-    collections::HashMap,
     fmt,
     ops::{Deref, DerefMut},
 };
+use std::hash::RandomState;
+use std::num::NonZeroUsize;
+use clru::{CLruCache, CLruCacheConfig, WeightScale};
+use petgraph::graph::Edge;
 use trustify_common::{cpe::Cpe, purl::Purl};
 use trustify_entity::relationship::Relationship;
 use utoipa::ToSchema;
@@ -35,7 +38,39 @@ pub struct PackageNode {
     pub document_id: String,
     pub product_name: String,
     pub product_version: String,
+    pub approximate_memory_size: usize,
 }
+
+impl PackageNode {
+    pub(crate) fn set_approximate_memory_size(&self) -> PackageNode {
+
+        // Is there a better way to do this?  I think this will under-estimate the memory size
+        let approximate_memory_size =
+            size_of::<PackageNode>() +
+            self.sbom_id.len() +
+            self.node_id.len() +
+            self.purl.iter().fold(0, |acc, purl|
+                // use the json string length as an approximation of the memory size
+                acc + serde_json::to_string(purl).unwrap_or_else(|_| "".to_string()).len()
+            ) +
+            self.cpe.iter().fold(0, |acc, cpe|
+                // use the json string length as an approximation of the memory size
+                acc + serde_json::to_string(cpe).unwrap_or_else(|_| "".to_string()).len()
+            ) +
+            self.name.len() +
+            self.version.len() +
+            self.published.len() +
+            self.document_id.len() +
+            self.product_name.len() +
+            self.product_version.len();
+
+        PackageNode {
+            approximate_memory_size,
+            ..self.clone()
+        }
+    }
+}
+
 impl fmt::Display for PackageNode {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{:?}", self.purl)
@@ -151,22 +186,38 @@ impl DerefMut for DepSummary {
     }
 }
 
-#[derive(Debug)]
 pub struct GraphMap {
-    map: HashMap<String, Graph<PackageNode, Relationship, petgraph::Directed>>,
+    map: CLruCache<String, Graph<PackageNode, Relationship, petgraph::Directed>, RandomState, GraphMapWeightScale>,
 }
 
+struct GraphMapWeightScale;
+
+impl WeightScale<String, Graph<PackageNode, Relationship, petgraph::Directed>> for GraphMapWeightScale {
+    fn weight(&self, key: &String, value: &Graph<PackageNode, Relationship, petgraph::Directed>) -> usize {
+        let mut result = key.len();
+        for n in value.raw_nodes() {
+            result + n.weight.approximate_memory_size;
+        }
+        result = value.raw_edges().len() * size_of::<Edge<Relationship>>();
+        result
+    }
+}
+
+
 impl GraphMap {
     // Create a new instance of GraphMap
-    pub fn new() -> Self {
+    pub fn new(cap: NonZeroUsize) -> Self {
+        let x = CLruCache::with_config(
+            CLruCacheConfig::new(cap).with_scale(GraphMapWeightScale {}),
+        );
         GraphMap {
-            map: HashMap::new(),
+            map: x,
         }
     }
 
     // Check if the map contains a key
     pub fn contains_key(&self, key: &str) -> bool {
-        self.map.contains_key(key)
+        self.map.contains(key)
     }
 
     // Get the number of graphs in the map
@@ -185,27 +236,22 @@ impl GraphMap {
         key: String,
         graph: Graph<PackageNode, Relationship, petgraph::Directed>,
     ) {
-        self.map.insert(key, graph);
+        self.map.put_with_weight(key, graph);
     }
 
     // Retrieve a reference to a graph by its key (read access)
-    pub fn get(&self, key: &str) -> Option<&Graph<PackageNode, Relationship, petgraph::Directed>> {
+    pub fn get(&mut self, key: &str) -> Option<&Graph<PackageNode, Relationship, petgraph::Directed>> {
         self.map.get(key)
     }
 
-    // Retrieve all sbom ids(read access)
-    pub fn sbom_ids(&self) -> Vec<String> {
-        self.map.keys().cloned().collect()
-    }
-
     // Clear all graphs from the map
     pub fn clear(&mut self) {
         self.map.clear();
     }
 }
 
-impl Default for GraphMap {
-    fn default() -> Self {
-        Self::new()
-    }
-}
+// impl Default for GraphMap {
+//     fn default() -> Self {
+//         Self::new()
+//     }
+// }

modules/analysis/src/model.rs

@@ -261,7 +261,8 @@ impl AnalysisService {
                         document_id: package.document_id.clone().unwrap_or_default(),
                         product_name: package.product_name.clone().unwrap_or_default(),
                         product_version: package.product_version.clone().unwrap_or_default(),
-                    });
+                        approximate_memory_size: 0,
+                    }.set_approximate_memory_size());
 
                     log::debug!("Inserting - id: {}, index: {index:?}", entry.key());
 

modules/analysis/src/service/load.rs

@@ -27,6 +27,7 @@ use std::{
     fmt::Debug,
     sync::Arc,
 };
+use std::num::NonZeroUsize;
 use tracing::instrument;
 use trustify_common::{
     db::query::Value,
@@ -154,9 +155,9 @@ impl AnalysisService {
     /// Also, we do not implement default because of this. As a new instance has the implication
     /// of having its own cache. So creating a new instance should be a deliberate choice.
     #[allow(clippy::new_without_default)]
-    pub fn new() -> Self {
+    pub fn new(cap: NonZeroUsize) -> Self {
         Self {
-            graph: Default::default(),
+            graph: Arc::new(RwLock::new(GraphMap::new(cap))),
         }
     }