From 6738b50f19e14d85cc6a9707f1346380d0e6f41c Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Fri, 29 Aug 2025 00:06:59 +0000
Subject: [PATCH] storage: optimize lease entry message build\n\n- Use
 SmallVec<[&[u8]; 32]> for borrowed ids\n- Skip sort for n<=1 to avoid
 overhead\n- Add unit tests for single and multiple ids\n- Add smallvec
 dependency

---
 Cargo.lock     |  1 +
 Cargo.toml     |  1 +
 src/storage.rs | 62 ++++++++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e4c906a..c688409 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1468,6 +1468,7 @@ dependencies = [
  "rand 0.8.5",
  "rocksdb",
  "scopeguard",
+ "smallvec",
  "socket2 0.5.10",
  "tempfile",
  "thiserror",
diff --git a/Cargo.toml b/Cargo.toml
index bde4685..7384789 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -24,6 +24,7 @@ console-subscriber = "0.4"
 uuid = { version = "1.18.0", features = ["std", "v4", "v7", "zerocopy"] }
 rand = "0.8"
 socket2 = "0.5"
+smallvec = { version = "1.13" }
 
 [build-dependencies]
 capnpc = "0.21.2"
diff --git a/src/storage.rs b/src/storage.rs
index 61cb6af..a6260d5 100644
--- a/src/storage.rs
+++ b/src/storage.rs
@@ -731,14 +731,16 @@ fn build_lease_entry_message(
     let mut lease_entry_builder = lease_entry.init_root::<protocol::lease_entry::Builder>();
     lease_entry_builder.set_expiry_ts_secs(expiry_ts_secs);
     lease_entry_builder.set_expiry_ts_index_key(expiry_index_key_bytes);
-    // Write ids unsorted first.
+    // Collect borrowed id slices. For small batches, avoid heap alloc by using SmallVec.
     let n = polled_items.len();
-    // Stable sort without copying id bytes: collect borrowed slices, sort stably, then write
-    let mut borrowed: Vec<&[u8]> = Vec::with_capacity(n);
+    use smallvec::SmallVec;
+    let mut borrowed: SmallVec<[&[u8]; 32]> = SmallVec::with_capacity(n.min(32));
     for item in polled_items.iter() {
         borrowed.push(item.get()?.get_id()?);
     }
-    borrowed.sort(); // stable
+    if n > 1 {
+        borrowed.sort();
+    }
     let mut out_ids = lease_entry_builder.init_ids(n as u32);
     for (i, id) in borrowed.into_iter().enumerate() {
         out_ids.set(i as u32, id);
@@ -1150,6 +1152,58 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn build_lease_entry_message_skips_sort_for_single() -> Result<()> {
+        // Create a single polled item
+        let mut msg = Builder::new_default();
+        let mut pi = msg.init_root::<protocol::polled_item::Builder>();
+        pi.set_id(b"z");
+        pi.set_contents(b"c");
+        let typed = msg.into_typed();
+        let owned: PolledItemOwnedReader = typed.into_reader();
+
+        // Build lease entry and read back via serialization
+        let lease_msg = build_lease_entry_message(123, b"k", &[owned])?;
+        let mut buf = Vec::new();
+        serialize::write_message(&mut buf, &lease_msg)?;
+        let msg =
+            serialize::read_message_from_flat_slice(&mut &buf[..], message::ReaderOptions::new())?;
+        let reader = msg.get_root::<protocol::lease_entry::Reader>()?;
+        let ids = reader.get_ids()?;
+        assert_eq!(ids.len(), 1);
+        assert_eq!(ids.get(0)?, b"z");
+        Ok(())
+    }
+
+    #[test]
+    fn build_lease_entry_message_sorts_multiple_ids() -> Result<()> {
+        // Two items out of order by id
+        let mut m1 = Builder::new_default();
+        let mut p1 = m1.init_root::<protocol::polled_item::Builder>();
+        p1.set_id(b"b");
+        p1.set_contents(b"x");
+        let o1: PolledItemOwnedReader = m1.into_typed().into_reader();
+
+        let mut m2 = Builder::new_default();
+        let mut p2 = m2.init_root::<protocol::polled_item::Builder>();
+        p2.set_id(b"a");
+        p2.set_contents(b"y");
+        let o2: PolledItemOwnedReader = m2.into_typed().into_reader();
+
+        // Build lease entry with unsorted inputs [b, a]
+        let lease_msg = build_lease_entry_message(1, b"k", &[o1, o2])?;
+        let mut buf = Vec::new();
+        serialize::write_message(&mut buf, &lease_msg)?;
+        let msg =
+            serialize::read_message_from_flat_slice(&mut &buf[..], message::ReaderOptions::new())?;
+        let reader = msg.get_root::<protocol::lease_entry::Reader>()?;
+        let ids = reader.get_ids()?;
+        assert_eq!(ids.len(), 2);
+        assert_eq!(ids.get(0)?, b"a");
+        assert_eq!(ids.get(1)?, b"b");
+        Ok(())
+    }
+
     #[test]
     fn poll_moves_multiple_items_and_updates_indexes()
     -> std::result::Result<(), Box<dyn std::error::Error>> {