Make in-process loglet readers obey readable tails

Author: AhmedSoliman

crates/bifrost/src/providers/local_loglet/mod.rs

@@ -138,8 +138,7 @@ impl Loglet for LocalLoglet {
         from: LogletOffset,
     ) -> Result<SendableLogletReadStream, OperationError> {
         let readable_tail = OffsetWatch::default();
-        let read_stream =
-            LocalLogletReadStream::create(self, filter, from, readable_tail.clone()).await?;
+        let read_stream = LocalLogletReadStream::create(self, filter, from, readable_tail.clone())?;
         Ok(SendableLogletReadStream::new(read_stream, readable_tail))
     }
 
@@ -389,17 +388,19 @@ mod tests {
             ("record-1", Keys::Single(1)).into(),
             ("record-2", Keys::Single(2)).into(),
             ("record-3", Keys::Single(1)).into(),
+            ("record-4", Keys::Single(2)).into(),
         ]
         .into();
         let offset = loglet.enqueue_batch(batch).await?.await?;
 
         let key_filter = KeyFilter::Include(1);
-        let read_stream = loglet
+        let mut read_stream = loglet
             .create_read_stream(key_filter, LogletOffset::OLDEST)
             .await?;
         read_stream.notify_readable_tail(offset.next());
 
         let records: Vec<_> = read_stream
+            .by_ref()
             .take(2)
             .try_collect::<Vec<_>>()
             .await?
@@ -419,6 +420,9 @@ mod tests {
                 eq((LogletOffset::from(3), "record-3".to_owned()))
             ]
         );
+        let filtered = read_stream.next().await.unwrap()?;
+        assert_that!(filtered.kind(), eq(crate::RecordKind::Filtered));
+        assert_that!(read_stream.read_pointer(), eq(offset.next()));
 
         Ok(())
     }

crates/bifrost/src/providers/local_loglet/read_stream.rs

@@ -20,9 +20,9 @@ use tracing::{debug, error, warn};
 
 use restate_core::ShutdownError;
 use restate_rocksdb::RocksDbReadPerfGuard;
-use restate_types::logs::{KeyFilter, LogletOffset, OffsetWatch, SequenceNumber, TailState};
+use restate_types::logs::{KeyFilter, LogletOffset, OffsetWatch, SequenceNumber};
 
-use crate::loglet::{Loglet, LogletReadStream, OperationError};
+use crate::loglet::{LogletReadStream, OperationError};
 use crate::providers::local_loglet::LogStoreError;
 use crate::providers::local_loglet::record_format::decode_and_filter_record;
 use crate::{LogEntry, Result};
@@ -37,12 +37,9 @@ pub(crate) struct LocalLogletReadStream {
     serde_buffer: BytesMut,
     // the next record this stream will attempt to read
     read_pointer: LogletOffset,
-    /// stop when read_pointer is at or beyond this offset
-    last_known_tail: LogletOffset,
     readable_tail_watch: BoxStream<'static, LogletOffset>,
     readable_tail: LogletOffset,
     iterator: DBRawIteratorWithThreadMode<'static, DB>,
-    tail_watch: BoxStream<'static, TailState<LogletOffset>>,
     terminated: bool,
     // IMPORTANT: Do not reorder, this should be dropped last since `iterator` holds a reference
     // into the underlying database.
@@ -64,7 +61,7 @@ unsafe fn ignore_iterator_lifetime<'a>(
 }
 
 impl LocalLogletReadStream {
-    pub(crate) async fn create(
+    pub(crate) fn create(
         loglet: Arc<LocalLoglet>,
         filter: KeyFilter,
         from_offset: LogletOffset,
@@ -95,12 +92,6 @@ impl LocalLogletReadStream {
         );
 
         let log_store = &loglet.log_store;
-        let mut tail_watch = loglet.watch_tail();
-        let last_known_tail = tail_watch
-            .next()
-            .await
-            .expect("loglet watch returns tail pointer")
-            .offset();
 
         // ## Safety:
         // the iterator is guaranteed to be dropped before the loglet is dropped, we hold to the
@@ -123,8 +114,6 @@ impl LocalLogletReadStream {
             read_pointer: from_offset,
             iterator: iter,
             terminated: false,
-            tail_watch,
-            last_known_tail,
             readable_tail_watch: Box::pin(readable_tail.to_stream()),
             readable_tail: readable_tail.get(),
         })
@@ -154,8 +143,15 @@ impl Stream for LocalLogletReadStream {
         }
 
         let perf_guard = RocksDbReadPerfGuard::new("local-loglet-next");
+        let mut filtered_from = None;
         loop {
             if self.read_pointer >= self.readable_tail {
+                if let Some(filtered_from) = filtered_from {
+                    return Poll::Ready(Some(Ok(LogEntry::new_filtered_gap(
+                        filtered_from,
+                        self.read_pointer.prev(),
+                    ))));
+                }
                 let maybe_readable_tail = match self.readable_tail_watch.poll_next_unpin(cx) {
                     Poll::Ready(tail) => tail,
                     Poll::Pending => {
@@ -174,36 +170,6 @@ impl Stream for LocalLogletReadStream {
                     }
                 }
             }
-            // Are we reading after commit offset?
-            // We are at tail. We need to wait until new records have been released.
-            if self.read_pointer >= self.last_known_tail {
-                let maybe_tail_state = match self.tail_watch.poll_next_unpin(cx) {
-                    Poll::Ready(t) => t,
-                    Poll::Pending => {
-                        perf_guard.forget();
-                        return Poll::Pending;
-                    }
-                };
-
-                match maybe_tail_state {
-                    Some(tail_state) => {
-                        // tail has been updated.
-                        self.last_known_tail = tail_state.offset();
-                        continue;
-                    }
-                    None => {
-                        // system shutdown. Or that the loglet has been unexpectedly shutdown.
-                        self.terminated = true;
-                        return Poll::Ready(Some(Err(OperationError::Shutdown(ShutdownError))));
-                    }
-                }
-            }
-            // tail has been updated.
-            let last_known_tail = self.last_known_tail;
-
-            // assert that we are behind tail
-            assert!(last_known_tail > self.read_pointer);
-
             // Trim point is the slot **before** the first readable record (if it exists)
             // trim point might have been updated since last time.
             let trim_point =
@@ -213,10 +179,11 @@ impl Stream for LocalLogletReadStream {
             assert!(self.read_pointer > LogletOffset::from(0));
 
             if self.read_pointer < head_offset {
-                let trim_gap = LogEntry::new_trim_gap(self.read_pointer, trim_point);
+                let gap_to = trim_point.min(self.readable_tail.prev());
+                let trim_gap = LogEntry::new_trim_gap(self.read_pointer, gap_to);
                 // next record should be beyond at the head
-                self.read_pointer = head_offset;
-                let key = RecordKey::new(self.loglet_id, trim_point);
+                self.read_pointer = gap_to.next();
+                let key = RecordKey::new(self.loglet_id, gap_to);
                 // park the iterator at the trim point, next iteration will seek it forward.
                 let key_bytes = key.encode_and_split(&mut self.serde_buffer);
                 self.iterator.seek(key_bytes);
@@ -251,7 +218,7 @@ impl Stream for LocalLogletReadStream {
                     loglet_id = self.loglet_id,
                     read_pointer = %self.read_pointer,
                     trim_point = %potentially_different_trim_point,
-                    last_known_tail = %self.last_known_tail,
+                    readable_tail = %self.readable_tail,
                     "poll_next() has moved to a non-existent record, that should not happen!"
                 );
                 panic!("poll_next() has moved to a non-existent record, that should not happen!");
@@ -283,7 +250,7 @@ impl Stream for LocalLogletReadStream {
             if let Some(record) = maybe_record {
                 return Poll::Ready(Some(Ok(LogEntry::new_data(key.offset, record))));
             }
-            // Didn't match, loop and read the next record if possible.
+            filtered_from.get_or_insert(key.offset);
         }
     }
 }

crates/bifrost/src/providers/memory_loglet.rs

@@ -200,34 +200,22 @@ struct MemoryReadStream {
     filter: KeyFilter,
     /// The next offset to read from
     read_pointer: LogletOffset,
-    tail_watch: BoxStream<'static, TailState<LogletOffset>>,
-    /// stop when read_pointer is at or beyond this offset
-    last_known_tail: LogletOffset,
     readable_tail_watch: BoxStream<'static, LogletOffset>,
     readable_tail: LogletOffset,
     terminated: bool,
 }
 
 impl MemoryReadStream {
-    async fn create(
+    fn create(
         loglet: Arc<MemoryLoglet>,
         filter: KeyFilter,
         from_offset: LogletOffset,
         readable_tail: OffsetWatch,
     ) -> Self {
-        let mut tail_watch = loglet.watch_tail();
-        let last_known_tail = tail_watch
-            .next()
-            .await
-            .expect("loglet watch returns tail pointer")
-            .offset();
-
         Self {
             loglet,
             filter,
             read_pointer: from_offset,
-            tail_watch,
-            last_known_tail,
             readable_tail_watch: Box::pin(readable_tail.to_stream()),
             readable_tail: readable_tail.get(),
             terminated: false,
@@ -257,10 +245,17 @@ impl Stream for MemoryReadStream {
             return Poll::Ready(None);
         }
 
+        let mut filtered_from = None;
         loop {
             let next_offset = self.read_pointer;
 
             if next_offset >= self.readable_tail {
+                if let Some(filtered_from) = filtered_from {
+                    return Poll::Ready(Some(Ok(LogEntry::new_filtered_gap(
+                        filtered_from,
+                        next_offset.prev(),
+                    ))));
+                }
                 match ready!(self.readable_tail_watch.poll_next_unpin(cx)) {
                     Some(readable_tail) => {
                         self.readable_tail = readable_tail;
@@ -273,28 +268,6 @@ impl Stream for MemoryReadStream {
                 }
             }
 
-            // Are we reading after commit offset?
-            // We are at tail. We need to wait until new records have been released.
-            if next_offset >= self.last_known_tail {
-                match ready!(self.tail_watch.poll_next_unpin(cx)) {
-                    Some(tail_state) => {
-                        self.last_known_tail = tail_state.offset();
-                        continue;
-                    }
-                    None => {
-                        // system shutdown. Or that the loglet has been unexpectedly shutdown.
-                        self.terminated = true;
-                        return Poll::Ready(Some(Err(OperationError::Shutdown(ShutdownError))));
-                    }
-                }
-            }
-
-            // tail has been updated.
-            let last_known_tail = self.last_known_tail;
-
-            // assert that we are behind tail
-            assert!(last_known_tail > next_offset);
-
             // Trim point is the the slot **before** the first readable record (if it exists)
             // trim point might have been updated since last time.
             let trim_point =
@@ -304,9 +277,10 @@ impl Stream for MemoryReadStream {
             // Are we reading behind the loglet head? -> TrimGap
             assert!(next_offset > LogletOffset::from(0));
             if next_offset < head_offset {
-                let trim_gap = LogEntry::new_trim_gap(next_offset, trim_point);
+                let gap_to = trim_point.min(self.readable_tail.prev());
+                let trim_gap = LogEntry::new_trim_gap(next_offset, gap_to);
                 // next record should be beyond at the head
-                self.read_pointer = head_offset;
+                self.read_pointer = gap_to.next();
                 return Poll::Ready(Some(Ok(trim_gap)));
             }
 
@@ -324,8 +298,7 @@ impl Stream for MemoryReadStream {
             if let Some(data_record) = next_record.as_record()
                 && !data_record.matches_key_query(&self.filter)
             {
-                // read_pointer is already advanced, just don't return the
-                // record and fast-forward.
+                filtered_from.get_or_insert(next_offset);
                 continue;
             }
 
@@ -350,7 +323,7 @@ impl Loglet for MemoryLoglet {
         from: LogletOffset,
     ) -> Result<SendableLogletReadStream, OperationError> {
         let readable_tail = OffsetWatch::default();
-        let read_stream = MemoryReadStream::create(self, filter, from, readable_tail.clone()).await;
+        let read_stream = MemoryReadStream::create(self, filter, from, readable_tail.clone());
         Ok(SendableLogletReadStream::new(read_stream, readable_tail))
     }
 

crates/bifrost/src/read_stream.rs

@@ -212,13 +212,7 @@ impl LogReadStream {
     /// The read pointer points to the next LSN will be attempted on the next
     /// `poll_next()`.
     fn calculate_read_pointer(record: &LogEntry) -> Lsn {
-        // On trim gaps, we fast-forward the read pointer beyond the end of the gap. We do
-        // this after delivering a TrimGap record. This means that the next read operation
-        // skips over the boundary of the gap.
-        record
-            .trim_gap_to_sequence_number()
-            .unwrap_or_else(|| record.sequence_number())
-            .next()
+        record.next_sequence_number()
     }
 
     pub fn safe_known_tail(&self) -> Option<Lsn> {