batch insertion of gauges and counts

sozu-proxy · Jul 13, 2022 · f2d07eb · f2d07eb
1 parent 3418b6f
commit f2d07eb
Showing 1 changed file with 26 additions and 24 deletions.
diff --git a/lib/src/metrics/local_drain.rs b/lib/src/metrics/local_drain.rs
@@ -563,11 +563,8 @@ impl LocalDrain {
       let complete_key = format!("{}\t{}", key, timestamp);
 
       trace!("store gauge at {} -> {}", complete_key, i);
-      if is_backend {
-          self.backend_tree.insert(complete_key.as_bytes(), &i.to_le_bytes())?;
-      } else {
-          self.cluster_tree.insert(complete_key.as_bytes(), &i.to_le_bytes())?;
-      }
+      let mut batch = sled::Batch::default();
+      batch.insert(complete_key.as_bytes(), &i.to_le_bytes());
 
       // aggregate at the last hour
       let second = now.second();
@@ -578,7 +575,7 @@ impl LocalDrain {
           let complete_key = format!("{}\t{}", key, timestamp);
 
           if !self.tree(is_backend).contains_key(complete_key.as_bytes())? {
-              self.tree(is_backend).insert(complete_key.as_bytes(), &i.to_le_bytes())?;
+              batch.insert(complete_key.as_bytes(), &i.to_le_bytes());
           }
 
           let minute = previous_minute.minute();
@@ -588,11 +585,13 @@ impl LocalDrain {
 
               let complete_key = format!("{}\t{}", key, timestamp);
               if !self.tree(is_backend).contains_key(complete_key.as_bytes())? {
-                  self.tree(is_backend).insert(complete_key.as_bytes(), &i.to_le_bytes())?;
+                  batch.insert(complete_key.as_bytes(), &i.to_le_bytes());
               }
           }
       }
 
+      self.tree(is_backend).apply_batch(batch)?;
+
       Ok(())
   }
 
@@ -602,6 +601,8 @@ impl LocalDrain {
       let complete_key = format!("{}\t{}", key, timestamp);
 
       trace!("add gauge at {} -> {}", complete_key, i);
+      let mut batch = sled::Batch::default();
+
       let value = match self.tree(is_backend).get(complete_key.as_bytes())? {
           Some(v) => i64::from_le_bytes((*v).try_into().unwrap()),
           // start from the last known value, or zero
@@ -614,7 +615,7 @@ impl LocalDrain {
       };
 
       let new_value = value + i;
-      self.tree(is_backend).insert(complete_key.as_bytes(), &new_value.to_le_bytes())?;
+      batch.insert(complete_key.as_bytes(), &new_value.to_le_bytes());
 
       // aggregate at the last hour
       let second = now.second();
@@ -636,7 +637,7 @@ impl LocalDrain {
           };
 
           let new_value = value + i;
-          self.tree(is_backend).insert(complete_key.as_bytes(), &new_value.to_le_bytes())?;
+          batch.insert(complete_key.as_bytes(), &new_value.to_le_bytes());
 
           let minute = previous_minute.minute();
           if minute != 0 {
@@ -657,10 +658,12 @@ impl LocalDrain {
               };
 
               let new_value = value + i;
-              self.tree(is_backend).insert(complete_key.as_bytes(), &new_value.to_le_bytes())?;
+              batch.insert(complete_key.as_bytes(), &new_value.to_le_bytes());
           }
       }
 
+      self.tree(is_backend).apply_batch(batch)?;
+
       Ok(())
   }
 
@@ -669,19 +672,16 @@ impl LocalDrain {
       let timestamp = now.unix_timestamp();
       let complete_key = format!("{}\t{}", key, timestamp);
 
-      let tree = if is_backend {
-          &mut self.backend_tree
-      } else {
-          &mut self.cluster_tree
-      };
       trace!("store count at {} -> {}", complete_key, i);
-      match tree.get(complete_key.as_bytes())? {
+      let mut batch = sled::Batch::default();
+
+      match self.tree(is_backend).get(complete_key.as_bytes())? {
           None => {
-              tree.insert(complete_key.as_bytes(), &i.to_le_bytes())?;
+              batch.insert(complete_key.as_bytes(), &i.to_le_bytes());
           },
           Some(v) => {
               let i2 = i64::from_le_bytes((*v).try_into().unwrap());
-              tree.insert(complete_key.as_bytes(), &(i+i2).to_le_bytes())?;
+              batch.insert(complete_key.as_bytes(), &(i+i2).to_le_bytes());
           }
       };
 
@@ -692,13 +692,13 @@ impl LocalDrain {
           let timestamp = previous_minute.unix_timestamp();
 
           let complete_key = format!("{}\t{}", key, timestamp);
-          match tree.get(complete_key.as_bytes())? {
+          match self.tree(is_backend).get(complete_key.as_bytes())? {
               None => {
-                  tree.insert(complete_key.as_bytes(), &i.to_le_bytes())?;
+                  batch.insert(complete_key.as_bytes(), &i.to_le_bytes());
               },
               Some(v) => {
                   let i2 = i64::from_le_bytes((*v).try_into().unwrap());
-                  tree.insert(complete_key.as_bytes(), &(i+i2).to_le_bytes())?;
+                  batch.insert(complete_key.as_bytes(), &(i+i2).to_le_bytes());
               }
           };
 
@@ -708,18 +708,20 @@ impl LocalDrain {
               let timestamp = previous_hour.unix_timestamp();
 
               let complete_key = format!("{}\t{}", key, timestamp);
-              match tree.get(complete_key.as_bytes())? {
+              match self.tree(is_backend).get(complete_key.as_bytes())? {
                   None => {
-                      tree.insert(complete_key.as_bytes(), &i.to_le_bytes())?;
+                      batch.insert(complete_key.as_bytes(), &i.to_le_bytes());
                   },
                   Some(v) => {
                       let i2 = i64::from_le_bytes((*v).try_into().unwrap());
-                      tree.insert(complete_key.as_bytes(), &(i+i2).to_le_bytes())?;
+                      batch.insert(complete_key.as_bytes(), &(i+i2).to_le_bytes());
                   }
               };
           }
       }
 
+      self.tree(is_backend).apply_batch(batch)?;
+
       Ok(())
   }