fix(server): handle byteview with empty bytes at the end by letting the amount of elements be specified in the blocks

server/src/db/block.rs

···
-use ordered_varint::Variable;
+use std::{
+    io::{self, Read, Write},
+    marker::PhantomData,
+    usize,
+};
+
 use rkyv::{
     Archive, Serialize, api::high::HighSerializer, rancor, ser::allocator::ArenaHandle,
     util::AlignedVec,
 };
-use std::{
-    io::{self, Read, Write},
-    marker::PhantomData,
-};
+
+use crate::utils::{ReadVariableExt, WriteVariableExt};
 
 pub struct Item<T> {
     pub timestamp: u64,
···
     writer: W,
     prev_timestamp: u64,
     prev_delta: i64,
+    item_count: usize,
     _item: PhantomData<T>,
 }
 
 impl<W: Write, T> ItemEncoder<W, T> {
-    pub fn new(writer: W) -> Self {
+    pub fn new(writer: W, item_count: usize) -> Self {
+        assert!(item_count > 0);
         ItemEncoder {
             writer,
             prev_timestamp: 0,
             prev_delta: 0,
+            item_count,
             _item: PhantomData,
         }
     }
 
+    /// NOTE: this is a best effort estimate of the encoded length of the block.
+    /// if T contains variable-length data, the encoded length may be larger than this estimate.
+    pub fn encoded_len(item_count: usize) -> usize {
+        // items length + item count * delta length + data length
+        size_of::<usize>() + item_count * size_of::<(i64, T)>()
+    }
+
     pub fn encode(&mut self, item: &Item<T>) -> io::Result<()> {
         if self.prev_timestamp == 0 {
+            self.writer.write_varint(self.item_count)?;
             // self.writer.write_varint(item.timestamp)?;
             self.prev_timestamp = item.timestamp;
             self.write_data(&item.data)?;
···
     reader: R,
     current_timestamp: u64,
     current_delta: i64,
-    first_item: bool,
+    items_read: usize,
+    expected: usize,
     _item: PhantomData<T>,
 }
 
 impl<R: Read, T: Archive> ItemDecoder<R, T> {
-    pub fn new(reader: R, start_timestamp: u64) -> io::Result<Self> {
+    pub fn new(mut reader: R, start_timestamp: u64) -> io::Result<Self> {
+        let expected = match reader.read_varint() {
+            Ok(expected) => expected,
+            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => 0,
+            Err(e) => return Err(e.into()),
+        };
+
         Ok(ItemDecoder {
             reader,
             current_timestamp: start_timestamp,
             current_delta: 0,
-            first_item: true,
+            items_read: 0,
+            expected,
             _item: PhantomData,
         })
     }
 
     pub fn decode(&mut self) -> io::Result<Option<Item<T>>> {
-        if self.first_item {
+        if self.items_read == 0 {
             // read the first timestamp
             // let timestamp = match self.reader.read_varint::<u64>() {
             //     Ok(timestamp) => timestamp,
···
             let Some(data_raw) = self.read_item()? else {
                 return Ok(None);
             };
-            self.first_item = false;
+
+            self.items_read += 1;
             return Ok(Some(Item {
                 timestamp: self.current_timestamp,
                 data: data_raw,
···
             }));
         }
 
+        if self.items_read >= self.expected {
+            return Ok(None);
+        }
+
         let Some(_delta) = self.read_timestamp()? else {
             return Ok(None);
         };
···
             }
         };
 
+        self.items_read += 1;
         Ok(Some(Item {
             timestamp: self.current_timestamp,
             data: data_raw,
···
     fn next(&mut self) -> Option<Self::Item> {
         self.decode().transpose()
     }
-}
 
-pub trait WriteVariableExt: Write {
-    fn write_varint(&mut self, value: impl Variable) -> io::Result<usize> {
-        value.encode_variable(self)
-    }
-}
-impl<W: Write> WriteVariableExt for W {}
-
-pub trait ReadVariableExt: Read {
-    fn read_varint<T: Variable>(&mut self) -> io::Result<T> {
-        T::decode_variable(self)
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.expected, Some(self.expected))
     }
 }
-impl<R: Read> ReadVariableExt for R {}
 
 #[cfg(test)]
 mod test {
···
 
         // encode
         let mut buffer = Vec::new();
-        let mut encoder = ItemEncoder::new(&mut buffer);
+        let mut encoder = ItemEncoder::new(&mut buffer, 1);
         encoder.encode(&item).unwrap();
         encoder.finish().unwrap();
 
···
 
         // encode
         let mut buffer = Vec::new();
-        let mut encoder = ItemEncoder::new(&mut buffer);
+        let mut encoder = ItemEncoder::new(&mut buffer, items.len());
 
         for item in &items {
             encoder.encode(item).unwrap();
···
 
         // encode
         let mut buffer = Vec::new();
-        let mut encoder = ItemEncoder::new(&mut buffer);
+        let mut encoder = ItemEncoder::new(&mut buffer, items.len());
 
         for item in &items {
             encoder.encode(item).unwrap();
···
         ];
 
         let mut buffer = Vec::new();
-        let mut encoder = ItemEncoder::new(&mut buffer);
+        let mut encoder = ItemEncoder::new(&mut buffer, items.len());
 
         for item in &items {
             encoder.encode(item).unwrap();
···
         ];
 
         let mut buffer = Vec::new();
-        let mut encoder = ItemEncoder::new(&mut buffer);
+        let mut encoder = ItemEncoder::new(&mut buffer, items.len());
 
         for item in &items {
             encoder.encode(item).unwrap();
···
         let items = vec![Item::new(1000, &small_data), Item::new(1001, &large_data)];
 
         let mut buffer = Vec::new();
-        let mut encoder = ItemEncoder::new(&mut buffer);
+        let mut encoder = ItemEncoder::new(&mut buffer, items.len());
 
         for item in &items {
             encoder.encode(item).unwrap();

server/src/db/mod.rs

···
 use std::{
-    io::{Cursor, Write},
+    io::Cursor,
     ops::{Bound, Deref, RangeBounds},
     path::Path,
     sync::{
···
 use byteview::ByteView;
 use fjall::{Config, Keyspace, Partition, PartitionCreateOptions, Slice};
 use itertools::{Either, Itertools};
-use ordered_varint::Variable;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 use rkyv::{Archive, Deserialize, Serialize, rancor::Error};
 use smol_str::SmolStr;
···
 use tokio_util::sync::CancellationToken;
 
 use crate::{
-    db::block::{ReadVariableExt, WriteVariableExt},
     error::{AppError, AppResult},
     jetstream::JetstreamEvent,
-    utils::{CLOCK, DefaultRateTracker, RateTracker},
+    utils::{
+        CLOCK, DefaultRateTracker, RateTracker, ReadVariableExt, WritableByteView,
+        varints_unsigned_encoded,
+    },
 };
 
 mod block;
···
 type ItemEncoder = block::ItemEncoder<WritableByteView, NsidHit>;
 type Item = block::Item<NsidHit>;
 
-struct WritableByteView {
-    view: ByteView,
-    written: usize,
-}
-
-impl WritableByteView {
-    // returns None if the view already has a reference to it
-    fn with_size(capacity: usize) -> Self {
-        Self {
-            view: ByteView::with_size(capacity),
-            written: 0,
-        }
-    }
-
-    #[inline(always)]
-    fn into_inner(self) -> ByteView {
-        self.view
-    }
-}
-
-impl Write for WritableByteView {
-    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-        let len = buf.len();
-        if len > self.view.len() - self.written {
-            return Err(std::io::Error::new(
-                std::io::ErrorKind::StorageFull,
-                "buffer full",
-            ));
-        }
-        // SAFETY: this is safe because we have checked that the buffer is not full
-        // SAFETY: we own the mutator so no other references to the view exist
-        unsafe {
-            std::ptr::copy_nonoverlapping(
-                buf.as_ptr(),
-                self.view
-                    .get_mut()
-                    .unwrap_unchecked()
-                    .as_mut_ptr()
-                    .add(self.written),
-                len,
-            );
-            self.written += len;
-        }
-        Ok(len)
-    }
-
-    #[inline(always)]
-    fn flush(&mut self) -> std::io::Result<()> {
-        Ok(())
-    }
-}
-
 struct Block {
     written: usize,
     key: ByteView,
···
         self.eps.observe();
     }
 
-    fn encode_block(&self, max_block_size: usize) -> AppResult<Option<Block>> {
-        let buf_size =
-            size_of::<u64>() + self.item_count().min(max_block_size) * size_of::<(u64, NsidHit)>();
-        let mut writer = ItemEncoder::new(WritableByteView::with_size(buf_size));
+    fn encode_block(&self, item_count: usize) -> AppResult<Option<Block>> {
+        let mut writer = ItemEncoder::new(
+            WritableByteView::with_size(ItemEncoder::encoded_len(item_count)),
+            item_count,
+        );
         let mut start_timestamp = None;
         let mut end_timestamp = None;
         let mut written = 0_usize;
···
                 start_timestamp = Some(event.timestamp);
             }
             end_timestamp = Some(event.timestamp);
-            if written >= max_block_size {
+            if written >= item_count {
                 break;
             }
             written += 1;
         }
+        if written != item_count {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidData,
+                "unexpected number of items, invalid data?",
+            )
+            .into());
+        }
         if let (Some(start_timestamp), Some(end_timestamp)) = (start_timestamp, end_timestamp) {
             self.buf_len.store(0, AtomicOrdering::SeqCst);
             let value = writer.finish()?;
-            let mut key = WritableByteView::with_size(size_of::<u64>() * 2);
-            key.write_varint(start_timestamp)?;
-            key.write_varint(end_timestamp)?;
+            let key = varints_unsigned_encoded([start_timestamp, end_timestamp]);
             return Ok(Some(Block {
                 written,
-                key: key.into_inner(),
+                key,
                 data: value.into_inner(),
             }));
         }
···
         nsid: &str,
         range: impl RangeBounds<u64> + std::fmt::Debug,
     ) -> impl Iterator<Item = AppResult<Item>> {
-        let start = range
-            .start_bound()
-            .cloned()
-            .map(|t| unsafe { t.to_variable_vec().unwrap_unchecked() });
-        let end = range
-            .end_bound()
-            .cloned()
-            .map(|t| unsafe { t.to_variable_vec().unwrap_unchecked() });
-        let limit = match range.end_bound().cloned() {
+        let start_limit = match range.start_bound().cloned() {
+            Bound::Included(start) => start,
+            Bound::Excluded(start) => start.saturating_add(1),
+            Bound::Unbounded => 0,
+        };
+        let end_limit = match range.end_bound().cloned() {
             Bound::Included(end) => end,
             Bound::Excluded(end) => end.saturating_sub(1),
             Bound::Unbounded => u64::MAX,
         };
+        let end_key = varints_unsigned_encoded([end_limit]);
 
         self.maybe_run_in_nsid_tree(nsid, move |handle| {
             let map_block = move |(key, val)| {
                 let mut key_reader = Cursor::new(key);
                 let start_timestamp = key_reader.read_varint::<u64>()?;
+                if start_timestamp < start_limit {
+                    return Ok(None);
+                }
                 let items = ItemDecoder::new(Cursor::new(val), start_timestamp)?
                     .take_while(move |item| {
-                        item.as_ref().map_or(true, |item| item.timestamp <= limit)
+                        item.as_ref().map_or(true, |item| {
+                            item.timestamp <= end_limit && item.timestamp >= start_limit
+                        })
                     })
                     .map(|res| res.map_err(AppError::from));
-                Ok(items)
+                Ok(Some(items))
             };
 
             Either::Left(
                 handle
                     .tree
-                    .range(TimestampRange { start, end })
-                    .map(move |res| res.map_err(AppError::from).and_then(map_block))
+                    .range(..end_key)
+                    .rev()
+                    .map_while(move |res| {
+                        res.map_err(AppError::from).and_then(map_block).transpose()
+                    })
+                    .collect::<Vec<_>>()
+                    .into_iter()
+                    .rev()
                     .flatten()
                     .flatten(),
             )
···
         .unwrap_or(Ok(0))
     }
 }
-
-type TimestampRepr = Vec<u8>;
-
-struct TimestampRange {
-    start: Bound<TimestampRepr>,
-    end: Bound<TimestampRepr>,
-}
-
-impl RangeBounds<TimestampRepr> for TimestampRange {
-    #[inline(always)]
-    fn start_bound(&self) -> Bound<&TimestampRepr> {
-        self.start.as_ref()
-    }
-
-    #[inline(always)]
-    fn end_bound(&self) -> Bound<&TimestampRepr> {
-        self.end.as_ref()
-    }
-}

server/src/utils.rs

···
+use std::io::{self, Read, Write};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
+
+use byteview::ByteView;
+use ordered_varint::Variable;
+
+pub trait WriteVariableExt: Write {
+    fn write_varint(&mut self, value: impl Variable) -> io::Result<usize> {
+        value.encode_variable(self)
+    }
+}
+impl<W: Write> WriteVariableExt for W {}
+
+pub trait ReadVariableExt: Read {
+    fn read_varint<T: Variable>(&mut self) -> io::Result<T> {
+        T::decode_variable(self)
+    }
+}
+impl<R: Read> ReadVariableExt for R {}
+
+pub struct WritableByteView {
+    view: ByteView,
+    written: usize,
+}
+
+impl WritableByteView {
+    // returns None if the view already has a reference to it
+    pub fn with_size(capacity: usize) -> Self {
+        Self {
+            view: ByteView::with_size(capacity),
+            written: 0,
+        }
+    }
+
+    #[inline(always)]
+    pub fn into_inner(self) -> ByteView {
+        self.view
+    }
+}
+
+impl Write for WritableByteView {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let len = buf.len();
+        if len > self.view.len() - self.written {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::StorageFull,
+                "buffer full",
+            ));
+        }
+        // SAFETY: this is safe because we have checked that the buffer is not full
+        // SAFETY: we own the mutator so no other references to the view exist
+        unsafe {
+            std::ptr::copy_nonoverlapping(
+                buf.as_ptr(),
+                self.view
+                    .get_mut()
+                    .unwrap_unchecked()
+                    .as_mut_ptr()
+                    .add(self.written),
+                len,
+            );
+            self.written += len;
+        }
+        Ok(len)
+    }
+
+    #[inline(always)]
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
+
+pub fn varints_unsigned_encoded<const N: usize>(values: [u64; N]) -> ByteView {
+    let mut buf =
+        WritableByteView::with_size(values.into_iter().map(varint_unsigned_encoded_len).sum());
+    for value in values {
+        // cant fail
+        let _ = buf.write_varint(value);
+    }
+    buf.into_inner()
+}
+
+// gets the encoded length of a varint-encoded unsigned integer
+// see ordered_varint
+pub fn varint_unsigned_encoded_len(value: u64) -> usize {
+    let value = value.to_be_bytes();
+    value
+        .iter()
+        .enumerate()
+        .find_map(|(index, &byte)| {
+            (byte > 0).then(|| {
+                let extra_bytes = 7 - index;
+                (byte < 16)
+                    .then(|| extra_bytes + 1)
+                    .unwrap_or_else(|| extra_bytes + 2)
+            })
+        })
+        .unwrap_or(0)
+        .max(1)
+}
 
 pub static CLOCK: std::sync::LazyLock<quanta::Clock> =
     std::sync::LazyLock::new(|| quanta::Clock::new());