[VQueues] Optimize sys_vqueue_meta point lookups

crates/partition-store/src/inbox_table/mod.rs

@@ -43,10 +43,9 @@ define_table_key!(
 );
 
 fn any_inbox_entry_in_range<S: StorageAccess>(storage: &mut S, range: KeyRange) -> Result<bool> {
-    storage.get_first_blocking(
-        TableScan::FullScanPartitionKeyRange::<InboxKey>(range),
-        |kv| Ok(kv.is_some()),
-    )
+    storage.get_first_blocking(TableScan::ScanPartitionKeyRange::<InboxKey>(range), |kv| {
+        Ok(kv.is_some())
+    })
 }
 
 fn peek_inbox<S: StorageAccess>(
@@ -58,16 +57,13 @@ fn peek_inbox<S: StorageAccess>(
         .service_name(service_id.service_name.clone())
         .service_key(service_id.key.clone());
 
-    storage.get_first_blocking(
-        TableScan::SinglePartitionKeyPrefix(service_id.partition_key(), key),
-        |kv| match kv {
-            Some((k, v)) => {
-                let entry = decode_inbox_key_value(k, v)?;
-                Ok(Some(entry))
-            }
-            None => Ok(None),
-        },
-    )
+    storage.get_first_blocking(TableScan::Prefix(key), |kv| match kv {
+        Some((k, v)) => {
+            let entry = decode_inbox_key_value(k, v)?;
+            Ok(Some(entry))
+        }
+        None => Ok(None),
+    })
 }
 
 fn inbox<S: StorageAccess>(
@@ -80,7 +76,7 @@ fn inbox<S: StorageAccess>(
         .service_key(service_id.key.clone());
 
     Ok(stream::iter(storage.for_each_key_value_in_place(
-        TableScan::SinglePartitionKeyPrefix(service_id.partition_key(), key),
+        TableScan::Prefix(key),
         |k, v| {
             let inbox_entry = decode_inbox_key_value(k, v);
             TableScanIterationDecision::Emit(inbox_entry)
@@ -121,7 +117,7 @@ impl ScanInboxTable for PartitionStore {
         self.iterator_for_each(
             "df-inbox",
             Priority::Low,
-            TableScan::FullScanPartitionKeyRange::<InboxKey>(range),
+            TableScan::ScanPartitionKeyRange::<InboxKey>(range),
             move |(mut key, mut value)| {
                 let key = break_on_err(InboxKey::deserialize_from(&mut key))?;
                 let inbox_entry = break_on_err(InboxEntry::decode(&mut value))?;

crates/partition-store/src/invocation_status_table/mod.rs

@@ -8,11 +8,14 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::collections::BTreeSet;
 use std::ops::ControlFlow;
 
-use futures::Stream;
+use bytes::BytesMut;
+use futures::{FutureExt, Stream};
+use rocksdb::ReadOptions;
 
-use restate_rocksdb::{Priority, RocksDbReadPerfGuard};
+use restate_rocksdb::{Priority, RocksDbReadPerfGuard, StorageTaskKind};
 use restate_storage_api::invocation_status_table::{
     InvocationLite, InvocationStatus, InvocationStatusDiscriminants, InvokedInvocationStatusLite,
     ReadInvocationStatusTable, ScanInvocationStatusTable, ScanInvocationStatusTableRange,
@@ -25,8 +28,8 @@ use restate_types::identifiers::{InvocationId, InvocationUuid, PartitionKey, Wit
 use restate_types::sharding::KeyRange;
 use restate_util_string::format_restring;
 
-use crate::TableScan::FullScanPartitionKeyRange;
-use crate::keys::{DecodeTableKey, KeyKind, define_table_key};
+use crate::TableScan::ScanPartitionKeyRange;
+use crate::keys::{DecodeTableKey, EncodeTableKey, KeyKind, define_table_key};
 use crate::scan::TableScan;
 use crate::{PartitionStore, PartitionStoreTransaction, StorageAccess, TableKind, break_on_err};
 
@@ -39,6 +42,20 @@ define_table_key!(
     )
 );
 
+impl InvocationStatusKey {
+    pub const fn serialized_length_fixed() -> usize {
+        KeyKind::SERIALIZED_LENGTH
+            + std::mem::size_of::<PartitionKey>()
+            + InvocationUuid::RAW_BYTES_LEN
+    }
+}
+
+/// Maximum number of invocation-status keys passed to one RocksDB multi-get call.
+/// This one is set to a low value due to the possibility of loading a large value from rocksdb
+/// when the invocation holds a large output payload. This is an unfortunate side effect of storing
+/// the output payload into invocation-status!
+const INVOCATION_STATUS_MULTI_GET_BATCH_SIZE: usize = 25;
+
 #[inline]
 fn create_invocation_status_key(invocation_id: &InvocationId) -> InvocationStatusKey {
     InvocationStatusKey {
@@ -56,6 +73,133 @@ fn invocation_id_from_key_bytes<B: bytes::Buf>(bytes: &mut B) -> crate::Result<I
     ))
 }
 
+fn multi_get_invocation_status_lazy<E, F>(
+    store: &PartitionStore,
+    ids: BTreeSet<InvocationId>,
+    mut f: F,
+) -> impl Future<Output = Result<()>> + Send
+where
+    E: Into<anyhow::Error> + 'static,
+    F: for<'a> FnMut(
+            (InvocationId, &'a InvocationStatusV2Lazy<'a>),
+        ) -> ControlFlow<std::result::Result<(), E>>
+        + Send
+        + Sync
+        + 'static,
+{
+    const KEY_LEN: usize = InvocationStatusKey::serialized_length_fixed();
+
+    let rocksdb = store.partition_db().rocksdb().clone();
+    let cf_name: restate_rocksdb::CfName = store.partition_db().partition().cf_name().into();
+
+    async move {
+        rocksdb
+            .run_background_read_op(
+                "df-for-each-invocation-status",
+                StorageTaskKind::MultiGet,
+                Priority::Low,
+                move |raw_db| -> Result<()> {
+                    let Some(cf) = raw_db.cf_handle(cf_name.as_str()) else {
+                        return Err(StorageError::Generic(anyhow::anyhow!(
+                            "column family {cf_name} not found for invocation-status multi-get"
+                        )));
+                    };
+
+                    let batch_capacity = ids.len().min(INVOCATION_STATUS_MULTI_GET_BATCH_SIZE);
+                    let mut key_buf = BytesMut::with_capacity(batch_capacity * KEY_LEN);
+                    let mut batch_ids = Vec::with_capacity(batch_capacity);
+
+                    let mut readopts = ReadOptions::default();
+                    // future proofing to make use of parallel L0 reads and async-io
+                    // if/when we build rocksdb with COROUTINES=1 and IO-URING support.
+                    // by default, this will not do anything.
+                    readopts.set_async_io(true);
+                    readopts.set_optimize_multiget_for_io(true);
+
+                    let mut ids = ids.into_iter();
+                    loop {
+                        key_buf.clear();
+                        batch_ids.clear();
+
+                        for id in ids.by_ref().take(INVOCATION_STATUS_MULTI_GET_BATCH_SIZE) {
+                            EncodeTableKey::serialize_to(
+                                &create_invocation_status_key(&id),
+                                &mut key_buf,
+                            );
+                            batch_ids.push(id);
+                        }
+
+                        if batch_ids.is_empty() {
+                            break;
+                        }
+
+                        let results = raw_db.batched_multi_get_cf_opt(
+                            &cf,
+                            key_buf.chunks_exact(KEY_LEN),
+                            true,
+                            &readopts,
+                        );
+
+                        for (id, result) in batch_ids.iter().zip(results) {
+                            let Some(value) = result.map_err(|e| StorageError::Generic(e.into()))?
+                            else {
+                                continue;
+                            };
+                            let mut value = value.as_ref();
+
+                            if value.len() < std::mem::size_of::<u8>() {
+                                return Err(StorageError::Conversion(
+                                    restate_types::storage::StorageDecodeError::ReadingCodec(
+                                        format_restring!(
+                                            "remaining bytes in buf '{}' < version bytes '{}'",
+                                            value.len(),
+                                            std::mem::size_of::<u8>()
+                                        ),
+                                    )
+                                    .into(),
+                                ));
+                            }
+
+                            let codec = restate_types::storage::StorageCodecKind::try_from(
+                                bytes::Buf::get_u8(&mut value),
+                            )
+                            .map_err(|e| StorageError::Conversion(e.into()))?;
+                            let restate_types::storage::StorageCodecKind::Protobuf = codec else {
+                                return Err(StorageError::Conversion(
+                                    restate_types::storage::StorageDecodeError::UnsupportedCodecKind(
+                                        codec,
+                                    )
+                                    .into(),
+                                ));
+                            };
+
+                            let mut inv_status_v2_lazy = InvocationStatusV2Lazy::default();
+                            inv_status_v2_lazy
+                                .merge(value)
+                                .map_err(|e| StorageError::Conversion(e.into()))?;
+
+                            match f((*id, &inv_status_v2_lazy)) {
+                                ControlFlow::Continue(()) => {}
+                                ControlFlow::Break(Ok(())) => return Ok(()),
+                                ControlFlow::Break(Err(e)) => {
+                                    return Err(StorageError::Conversion(e.into()));
+                                }
+                            }
+                        }
+
+                        if batch_ids.len() < INVOCATION_STATUS_MULTI_GET_BATCH_SIZE {
+                            break;
+                        }
+                    }
+
+                    Ok(())
+                },
+            )
+            .await
+            .map_err(|_| StorageError::OperationalError)?
+    }
+}
+
 fn put_invocation_status<S: StorageAccess>(
     storage: &mut S,
     invocation_id: &InvocationId,
@@ -95,7 +239,7 @@ fn any_non_completed_invocation_in_range<S: StorageAccess>(
     storage: &S,
     range: KeyRange,
 ) -> Result<bool> {
-    let mut iterator = storage.iterator_from(TableScan::FullScanPartitionKeyRange::<
+    let mut iterator = storage.iterator_from(TableScan::ScanPartitionKeyRange::<
         InvocationStatusKey,
     >(range))?;
 
@@ -155,14 +299,14 @@ impl ScanInvocationStatusTable for PartitionStore {
         self.iterator_filter_map(
             "scan-all-invoked",
             Priority::High,
-            FullScanPartitionKeyRange::<InvocationStatusKey>(self.partition_key_range()),
+            ScanPartitionKeyRange::<InvocationStatusKey>(self.partition_key_range()),
             read_invoked_full_invocation_id,
         )
         .map_err(|_| StorageError::OperationalError)
     }
 
     fn for_each_invocation_status_lazy<
-        E: Into<anyhow::Error>,
+        E: Into<anyhow::Error> + 'static,
         F: for<'a> FnMut(
                 (InvocationId, &'a InvocationStatusV2Lazy<'a>),
             ) -> ControlFlow<std::result::Result<(), E>>
@@ -174,9 +318,13 @@ impl ScanInvocationStatusTable for PartitionStore {
         range: ScanInvocationStatusTableRange,
         mut f: F,
     ) -> Result<impl Future<Output = Result<()>> + Send> {
+        if let ScanInvocationStatusTableRange::InvocationIdSet(ids) = range {
+            return Ok(multi_get_invocation_status_lazy(self, ids, f).boxed());
+        }
+
         let scan = match range {
             ScanInvocationStatusTableRange::PartitionKey(partition_key) => {
-                TableScan::FullScanPartitionKeyRange::<InvocationStatusKeyBuilder>(partition_key)
+                TableScan::ScanPartitionKeyRange::<InvocationStatusKeyBuilder>(partition_key)
             }
             ScanInvocationStatusTableRange::InvocationId(invocation_id) => {
                 let start = InvocationStatusKey::builder()
@@ -187,8 +335,9 @@ impl ScanInvocationStatusTable for PartitionStore {
                     .partition_key(invocation_id.end().partition_key())
                     .invocation_uuid(invocation_id.end().invocation_uuid());
 
-                TableScan::KeyRangeInclusiveInSinglePartition(self.partition_id(), start, end)
+                TableScan::RangeInclusive(start, end)
             }
+            ScanInvocationStatusTableRange::InvocationIdSet(_) => unreachable!("handled above"),
         };
 
         let new_status_keys = self
@@ -234,7 +383,7 @@ impl ScanInvocationStatusTable for PartitionStore {
             )
             .map_err(|_| StorageError::OperationalError)?;
 
-        Ok(new_status_keys)
+        Ok(new_status_keys.boxed())
     }
 
     fn filter_map_invocation_status_lazy<
@@ -254,7 +403,7 @@ impl ScanInvocationStatusTable for PartitionStore {
             .iterator_filter_map(
                 "df-filter-map-invocation-status",
                 Priority::Low,
-                TableScan::FullScanPartitionKeyRange::<InvocationStatusKey>(
+                TableScan::ScanPartitionKeyRange::<InvocationStatusKey>(
                     self.partition_key_range(),
                 ),
                 {

crates/partition-store/src/journal_events/mod.rs

@@ -106,34 +106,31 @@ fn get_journal_events<S: StorageAccess>(
         .partition_key(invocation_id.partition_key())
         .invocation_uuid(invocation_id.invocation_uuid());
 
-    storage.for_each_key_value_in_place(
-        TableScan::SinglePartitionKeyPrefix(invocation_id.partition_key(), key),
-        move |mut k, mut v| {
-            TableScanIterationDecision::Emit(
-                JournalEventKey::deserialize_from(&mut k)
-                    .map(|k| k.split())
-                    .and_then(|k_elements| {
-                        EventWrapper::decode(&mut v).map(|value| (k_elements, value.0))
-                    })
-                    .map(
-                        |(
-                            (_, _, event_type, timestamp, _),
-                            PbEvent {
-                                content,
-                                after_journal_entry_index,
-                            },
-                        )| EventView {
-                            event: RawEvent::new(
-                                EventType::from_repr(event_type).unwrap_or(EventType::Unknown),
-                                content,
-                            ),
+    storage.for_each_key_value_in_place(TableScan::Prefix(key), move |mut k, mut v| {
+        TableScanIterationDecision::Emit(
+            JournalEventKey::deserialize_from(&mut k)
+                .map(|k| k.split())
+                .and_then(|k_elements| {
+                    EventWrapper::decode(&mut v).map(|value| (k_elements, value.0))
+                })
+                .map(
+                    |(
+                        (_, _, event_type, timestamp, _),
+                        PbEvent {
+                            content,
                             after_journal_entry_index,
-                            append_time: MillisSinceEpoch::new(timestamp),
                         },
-                    ),
-            )
-        },
-    )
+                    )| EventView {
+                        event: RawEvent::new(
+                            EventType::from_repr(event_type).unwrap_or(EventType::Unknown),
+                            content,
+                        ),
+                        after_journal_entry_index,
+                        append_time: MillisSinceEpoch::new(timestamp),
+                    },
+                ),
+        )
+    })
 }
 
 fn delete_journal_events<S: StorageAccess>(
@@ -146,10 +143,9 @@ fn delete_journal_events<S: StorageAccess>(
         .partition_key(invocation_id.partition_key())
         .invocation_uuid(invocation_id.invocation_uuid());
 
-    let keys = storage.for_each_key_value_in_place(
-        TableScan::SinglePartitionKeyPrefix(invocation_id.partition_key(), prefix_key),
-        |k, _| TableScanIterationDecision::Emit(Ok(Box::from(k))),
-    )?;
+    let keys = storage.for_each_key_value_in_place(TableScan::Prefix(prefix_key), |k, _| {
+        TableScanIterationDecision::Emit(Ok(Box::from(k)))
+    })?;
 
     for k in keys {
         let key = k?;
@@ -178,7 +174,7 @@ impl ScanJournalEventsTable for PartitionStore {
     ) -> Result<impl Future<Output = Result<()>> + Send> {
         let scan = match range {
             ScanJournalEventsTableRange::PartitionKey(partition_key) => {
-                TableScan::FullScanPartitionKeyRange::<JournalEventKeyBuilder>(partition_key)
+                TableScan::ScanPartitionKeyRange::<JournalEventKeyBuilder>(partition_key)
             }
             ScanJournalEventsTableRange::InvocationId(invocation_id) => {
                 let start = JournalEventKey::builder()
@@ -189,7 +185,7 @@ impl ScanJournalEventsTable for PartitionStore {
                     .partition_key(invocation_id.end().partition_key())
                     .invocation_uuid(invocation_id.end().invocation_uuid());
 
-                TableScan::KeyRangeInclusiveInSinglePartition(self.partition_id(), start, end)
+                TableScan::RangeInclusive(start, end)
             }
         };