Document queue APIs; ensure user key/value sizes are checked; fix bug…

… in MAX_VALUE_SIZE

src/lib.rs

@@ -139,7 +139,7 @@ pub(crate) const MAX_TOTAL_VALUE_SIZE: usize = 0xffff; // 16 bits
 pub(crate) const NAMESPACING_RESERVED_SIZE: usize = 0xff;
 pub(crate) const VALUE_RESERVED_SIZE: usize = 0xff;
 pub const MAX_KEY_SIZE: usize = MAX_TOTAL_KEY_SIZE - NAMESPACING_RESERVED_SIZE;
-pub const MAX_VALUE_SIZE: usize = MAX_TOTAL_KEY_SIZE - VALUE_RESERVED_SIZE;
+pub const MAX_VALUE_SIZE: usize = MAX_TOTAL_VALUE_SIZE - VALUE_RESERVED_SIZE;
 
 const _: () = assert!(MAX_KEY_SIZE <= u16::MAX as usize);
 const _: () = assert!(MAX_VALUE_SIZE <= u16::MAX as usize);

src/queues.rs

@@ -16,6 +16,15 @@ struct Queue {
     num_items: u64,
 }
 
+impl Queue {
+    fn span(&self) -> u64 {
+        self.tail_idx - self.head_idx
+    }
+    fn is_empty(&self) -> bool {
+        self.head_idx == self.tail_idx
+    }
+}
+
 enum QueuePos {
     Head,
     Tail,
@@ -141,13 +150,16 @@ impl CandyStore {
         Ok(item_idx as usize)
     }
 
+    /// Pushed a new element at the front (head) of the queue, returning the element's index in the queue
     pub fn push_to_queue_head<B1: AsRef<[u8]> + ?Sized, B2: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B1,
         val: &B2,
     ) -> Result<usize> {
         self._push_to_queue(queue_key.as_ref(), val.as_ref(), QueuePos::Head)
     }
+
+    /// Pushed a new element at the end (tail) of the queue, returning the element's index in the queue
     pub fn push_to_queue_tail<B1: AsRef<[u8]> + ?Sized, B2: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B1,
@@ -191,7 +203,7 @@ impl CandyStore {
             }
         }
 
-        if queue.head_idx == queue.tail_idx {
+        if queue.is_empty() {
             self.remove_raw(&full_queue_key)?;
         } else {
             self.set_raw(&full_queue_key, &queue_bytes)?;
@@ -200,19 +212,28 @@ impl CandyStore {
         Ok(val)
     }
 
+    /// Removes and returns the head element of the queue, or None if the queue is empty
     pub fn pop_queue_head<B: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B,
     ) -> Result<Option<Vec<u8>>> {
         self._pop_queue(queue_key.as_ref(), QueuePos::Head)
     }
+
+    /// Removes and returns the tail element of the queue, or None if the queue is empty
     pub fn pop_queue_tail<B: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B,
     ) -> Result<Option<Vec<u8>>> {
         self._pop_queue(queue_key.as_ref(), QueuePos::Tail)
     }
 
+    /// Removes an element by index from the queue, returning the value it had or None if it did not exist (as well
+    /// as if the queue itself does not exist).
+    ///
+    /// This will leave a "hole" in the queue, which means we will skip over it in future iterations, but this could
+    /// lead to inefficienies as if you keep only the head and tail elements of a long queue, while removing elements
+    /// from the middle.
     pub fn remove_from_queue<B: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B,
@@ -236,7 +257,7 @@ impl CandyStore {
                 queue.tail_idx -= 1;
             }
             queue.num_items -= 1;
-            if queue.head_idx == queue.tail_idx {
+            if queue.is_empty() {
                 self.remove_raw(&full_queue_key)?;
             } else {
                 self.set_raw(&full_queue_key, &queue_bytes)?;
@@ -246,6 +267,7 @@ impl CandyStore {
         Ok(Some(val))
     }
 
+    /// Discards the queue (dropping all elements in contains). Returns true if it had existed before, false otherwise
     pub fn discard_queue<B: AsRef<[u8]> + ?Sized>(&self, queue_key: &B) -> Result<bool> {
         let queue_key = queue_key.as_ref();
         let (queue_ph, full_queue_key) = self.make_queue_key(queue_key);
@@ -275,6 +297,14 @@ impl CandyStore {
         }
     }
 
+    /// Extends the queue with elements from the given iterator. The queue will be created if it did not exist before,
+    /// and elements are pushed at the tail-end of the queue. This is more efficient than calling
+    /// [Self::push_to_queue_tail] in a loop
+    ///
+    /// Note: this is not an atomic (crash-safe) operation: if your program crashes while extending the queue, it
+    /// is possible that only some of the elements will have been appended.
+    ///
+    /// Returns the indices of the elements added (a range)
     pub fn extend_queue<'a, B: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B,
@@ -314,6 +344,7 @@ impl CandyStore {
         Ok(indices)
     }
 
+    /// Returns (without removing) the head element of the queue, or None if the queue is empty
     pub fn peek_queue_head<B: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B,
@@ -325,6 +356,7 @@ impl CandyStore {
         Ok(Some(v))
     }
 
+    /// Returns (without removing) the tail element of the queue, or None if the queue is empty
     pub fn peek_queue_tail<B: AsRef<[u8]> + ?Sized>(
         &self,
         queue_key: &B,
@@ -336,6 +368,8 @@ impl CandyStore {
         Ok(Some(v))
     }
 
+    /// Returns a forward iterator (head to tail) over the elements of the queue. If the queue does not exist,
+    /// this is an empty iterator.
     pub fn iter_queue<'a, B: AsRef<[u8]> + ?Sized>(&'a self, queue_key: &B) -> QueueIterator<'a> {
         QueueIterator {
             store: &self,
@@ -346,6 +380,8 @@ impl CandyStore {
         }
     }
 
+    /// Returns a backward iterator (tail to head) over the elements of the queue. If the queue does not exist,
+    /// this is an empty iterator.
     pub fn iter_queue_backwards<'a, B: AsRef<[u8]> + ?Sized>(
         &'a self,
         queue_key: &B,
@@ -359,6 +395,7 @@ impl CandyStore {
         }
     }
 
+    /// Returns a the length of the given queue (number of elements in the queue) or 0 if the queue does not exist
     pub fn queue_len<B: AsRef<[u8]> + ?Sized>(&self, queue_key: &B) -> Result<usize> {
         let Some(queue) = self.fetch_queue(queue_key.as_ref())? else {
             return Ok(0);

src/store.rs

@@ -1,4 +1,4 @@
-use anyhow::{anyhow, bail};
+use anyhow::{anyhow, bail, ensure};
 use fslock::LockFile;
 use parking_lot::Mutex;
 use std::{
@@ -10,7 +10,7 @@ use crate::{
     hashing::{HashSeed, PartedHash},
     router::ShardRouter,
     shard::{CompactionThreadPool, InsertMode, InsertStatus, KVPair},
-    Stats, MAX_TOTAL_VALUE_SIZE,
+    Stats, MAX_KEY_SIZE, MAX_TOTAL_VALUE_SIZE,
 };
 use crate::{
     shard::{NUM_ROWS, ROW_WIDTH},
@@ -266,6 +266,16 @@ impl CandyStore {
         Ok(())
     }
 
+    pub(crate) fn ensure_sizes(key: &[u8], val: &[u8]) -> Result<()> {
+        ensure!(key.len() <= MAX_KEY_SIZE, CandyError::KeyTooLong(key.len()));
+        ensure!(
+            val.len() <= MAX_VALUE_SIZE,
+            CandyError::ValueTooLong(val.len())
+        );
+
+        Ok(())
+    }
+
     pub(crate) fn make_user_key(&self, mut key: Vec<u8>) -> Vec<u8> {
         key.extend_from_slice(USER_NAMESPACE);
         key
@@ -329,12 +339,15 @@ impl CandyStore {
     ) -> Result<InsertStatus> {
         let ph = PartedHash::new(&self.config.hash_seed, full_key);
 
-        if full_key.len() > MAX_TOTAL_KEY_SIZE as usize {
-            return Err(anyhow!(CandyError::KeyTooLong(full_key.len())));
-        }
-        if val.len() > MAX_TOTAL_VALUE_SIZE as usize {
-            return Err(anyhow!(CandyError::ValueTooLong(val.len())));
-        }
+        ensure!(
+            full_key.len() <= MAX_TOTAL_KEY_SIZE,
+            CandyError::KeyTooLong(full_key.len())
+        );
+        ensure!(
+            val.len() <= MAX_TOTAL_VALUE_SIZE,
+            CandyError::ValueTooLong(val.len())
+        );
+
         if full_key.len() + val.len() > self.config.max_shard_size as usize {
             return Err(anyhow!(CandyError::EntryCannotFitInShard(
                 full_key.len() + val.len(),
@@ -373,6 +386,7 @@ impl CandyStore {
 
     /// Same as [Self::set], but the key passed owned to this function
     pub fn owned_set(&self, key: Vec<u8>, val: &[u8]) -> Result<SetStatus> {
+        Self::ensure_sizes(&key, &val)?;
         self.set_raw(&self.make_user_key(key), val)
     }
 
@@ -416,7 +430,8 @@ impl CandyStore {
         val: &[u8],
         expected_val: Option<&[u8]>,
     ) -> Result<ReplaceStatus> {
-        self.replace_raw(&self.make_user_key(key), val.as_ref(), expected_val)
+        Self::ensure_sizes(&key, &val)?;
+        self.replace_raw(&self.make_user_key(key), val, expected_val)
     }
 
     pub(crate) fn get_or_create_raw(
@@ -453,6 +468,7 @@ impl CandyStore {
         key: Vec<u8>,
         default_val: Vec<u8>,
     ) -> Result<GetOrCreateStatus> {
+        Self::ensure_sizes(&key, &default_val)?;
         self.get_or_create_raw(&self.make_user_key(key), default_val)
     }
 

tests/test_logic.rs

@@ -2,7 +2,7 @@ mod common;
 
 use std::collections::HashSet;
 
-use candystore::{CandyStore, Config, Result};
+use candystore::{CandyStore, Config, Result, MAX_VALUE_SIZE};
 
 use crate::common::{run_in_tempdir, LONG_VAL};
 
@@ -127,7 +127,7 @@ fn test_histogram() -> Result<()> {
         db.set("k4", &vec![b'b'; 5000])?;
         db.set("k4", &vec![b'b'; 4500])?;
         db.set("k5", &vec![b'b'; 50000])?;
-        db.set("kkkkkkkkkkkkkkk", &vec![b'b'; 0xffff])?;
+        db.set("kkkkkkkkkkkkkkk", &vec![b'b'; MAX_VALUE_SIZE])?;
 
         let stats = db.stats();
         assert_eq!(stats.entries_under_128, 2);