[F42] propagate the configured datastore key query cap from RPC/gRPC …

massa-api/src/public.rs

@@ -1636,7 +1636,7 @@ fn get_address_datastore_keys_to_state_query_item(
         (Some(k), false) => std::ops::Bound::Excluded(k),
     };
 
-    let (prefix, start_key, end_key) = cleanup_datastore_key_range_query(
+    let (prefix, start_key, end_key, count) = cleanup_datastore_key_range_query(
         &value.prefix,
         start_key,
         end_key,
@@ -1657,15 +1657,15 @@ fn get_address_datastore_keys_to_state_query_item(
             prefix,
             start_key,
             end_key,
-            count: value.count,
+            count,
         })
     } else {
         Ok(ExecutionQueryRequestItem::AddressDatastoreKeysCandidate {
             address: value.address,
             prefix,
             start_key,
             end_key,
-            count: value.count,
+            count,
         })
     }
 }

massa-execution-exports/src/mapping_grpc.rs

@@ -71,7 +71,7 @@ pub fn to_querystate_filter(
                     (Some(k), false) => std::ops::Bound::Excluded(k),
                 };
 
-                let (prefix, start_key, end_key) = cleanup_datastore_key_range_query(
+                let (prefix, start_key, end_key, count) = cleanup_datastore_key_range_query(
                     &value.prefix,
                     start_key,
                     end_key,
@@ -85,7 +85,7 @@ pub fn to_querystate_filter(
                     prefix,
                     start_key,
                     end_key,
-                    count: value.limit,
+                    count,
                 })
             }
             exec::RequestItem::AddressDatastoreKeysFinal(value) => {
@@ -102,7 +102,7 @@ pub fn to_querystate_filter(
                     (Some(k), false) => std::ops::Bound::Excluded(k),
                 };
 
-                let (prefix, start_key, end_key) = cleanup_datastore_key_range_query(
+                let (prefix, start_key, end_key, count) = cleanup_datastore_key_range_query(
                     &value.prefix,
                     start_key,
                     end_key,
@@ -116,7 +116,7 @@ pub fn to_querystate_filter(
                     prefix,
                     start_key,
                     end_key,
-                    count: value.limit,
+                    count,
                 })
             }
             exec::RequestItem::AddressDatastoreValueCandidate(value) => {

massa-execution-worker/src/context.rs

@@ -660,7 +660,7 @@ impl ExecutionContext {
         let max_datastore_query = None;
 
         // cleanup bounds
-        let (prefix, start_key, end_key) = cleanup_datastore_key_range_query(
+        let (prefix, start_key, end_key, count) = cleanup_datastore_key_range_query(
             prefix,
             start_key,
             end_key,

massa-execution-worker/src/datastore_scan.rs

@@ -8,7 +8,7 @@ use std::{
 };
 
 use massa_final_state::FinalStateController;
-use massa_ledger_exports::LedgerChanges;
+use massa_ledger_exports::{LedgerChanges, LedgerEntry};
 use massa_models::{
     address::Address,
     datastore::{get_prefix_bounds, range_intersection},
@@ -68,6 +68,7 @@ pub fn scan_datastore(
     let mut key_updates = BTreeMap::new();
     {
         let mut update_indices = VecDeque::new();
+        let mut reset_entry: Option<&LedgerEntry> = None;
         let history_lock = active_history.read();
 
         let it = history_lock
@@ -84,13 +85,9 @@ pub fn scan_datastore(
 
                 // address ledger entry being reset to an absolute new list of keys
                 Some(SetUpdateOrDelete::Set(v)) => {
-                    if let Some(k_range) = key_range.as_ref() {
-                        key_updates = v
-                            .datastore
-                            .range(k_range.clone())
-                            .map(|(k, _v)| (k.clone(), true))
-                            .collect();
-                    }
+                    // the entry is only scanned once the newer updates are known, so that
+                    // the scan can be bounded instead of materializing the whole range
+                    reset_entry = Some(v);
                     speculative_reset = SpeculativeResetType::Set;
                     break;
                 }
@@ -146,6 +143,27 @@ pub fn scan_datastore(
                 panic!("unexpected state change");
             }
         }
+
+        // Scan the keys of a reset entry, now that the newer updates are known.
+        // Only the first `count + key_updates.len()` keys of the entry can matter: the
+        // updates already gathered are the only thing that can delete a key from that
+        // range, so any key past that bound is guaranteed not to reach the final result.
+        // Updates are newer than the reset, so they take precedence over the entry.
+        if let (Some(entry), Some(k_range)) = (reset_entry, key_range.as_ref()) {
+            let base_it = entry.datastore.range(k_range.clone()).map(|(k, _v)| k);
+            match count.map(|cnt| (cnt as usize).saturating_add(key_updates.len())) {
+                Some(limit) => {
+                    for k in base_it.take(limit) {
+                        key_updates.entry(k.clone()).or_insert(true);
+                    }
+                }
+                None => {
+                    for k in base_it {
+                        key_updates.entry(k.clone()).or_insert(true);
+                    }
+                }
+            }
+        }
     }
 
     // process reset-related edge cases
@@ -264,6 +282,12 @@ pub fn scan_datastore(
         if final_keys_queue.is_empty() {
             if let Some(last_k) = last_final_batch_key.take() {
                 // the last final item was consumed: replenish the queue by querying more
+                // only the keys still missing to reach `count` are worth fetching
+                let remaining = count.map(|cnt| {
+                    (cnt as u64)
+                        .saturating_sub(speculative_keys.len() as u64)
+                        .max(1) as u32
+                });
                 final_keys_queue = final_state
                     .read()
                     .get_ledger()
@@ -272,7 +296,7 @@ pub fn scan_datastore(
                         prefix,
                         std::ops::Bound::Excluded(last_k),
                         end_key.clone(),
-                        count,
+                        remaining,
                     )
                     .expect("address expected to exist in final state")
                     .iter()

massa-execution-worker/src/tests/tests_scan_datastore.rs

@@ -433,3 +433,141 @@ fn display_bound_human_readable(bound: Bound<Vec<u8>>) {
         Bound::Unbounded => dbg!("bound key Unbounded".to_string()),
     };
 }
+
+/// Bounding the speculative scan must not change what is returned: for any `count`,
+/// the result has to be exactly the first `count` keys of the unbounded result.
+///
+/// The reset (`Set`) path is the interesting one, because deletions applied on top of
+/// the reset entry mean the scan has to look past the first `count` entry keys to
+/// still produce `count` keys.
+#[test]
+fn test_scan_datastore_count_is_a_prefix_of_unbounded() {
+    let mut rng = thread_rng();
+    for _ in 0..20 {
+        scan_datastore_count_prefix_case(rng.gen_range(20..60));
+    }
+}
+
+fn scan_datastore_count_prefix_case(nb_keys: usize) {
+    let keypair = KeyPair::generate(0).unwrap();
+    let addr = Address::from_public_key(&keypair.get_public_key());
+
+    let mut foreign_controllers = ExecutionForeignControllers::new_with_mocks();
+    foreign_controllers
+        .ledger_controller
+        .set_expectations(|ledger_controller| {
+            ledger_controller
+                .expect_get_datastore_keys()
+                .returning(move |_, _, _, _, _| None);
+        });
+    foreign_controllers
+        .final_state
+        .write()
+        .expect_get_ledger()
+        .return_const(Box::new(foreign_controllers.ledger_controller.clone()));
+
+    let mut rng = thread_rng();
+
+    // the reset entry
+    let mut data = BTreeMap::new();
+    for i in 0..nb_keys {
+        // fixed-width keys so that byte order matches the generation order
+        let key = format!("key{:04}", i).into_bytes();
+        let value: Vec<u8> = (0..rng.gen_range(1..10))
+            .map(|_| rng.sample(Alphanumeric) as u8)
+            .collect();
+        data.insert(key, value);
+    }
+    let existing_keys: Vec<_> = data.keys().cloned().collect();
+
+    let mut changes = PreHashMap::default();
+    changes.insert(
+        addr,
+        massa_models::types::SetUpdateOrDelete::Set(LedgerEntry {
+            datastore: data.clone(),
+            ..Default::default()
+        }),
+    );
+
+    // updates newer than the reset: deletions concentrated on the lowest keys, so a
+    // naive `take(count)` on the entry would come up short, plus a few added keys
+    let mut datastore_updates = BTreeMap::new();
+    for key in existing_keys.iter().take(nb_keys / 3) {
+        datastore_updates.insert(key.clone(), massa_models::types::SetOrDelete::Delete);
+    }
+    for i in 0..5 {
+        // "add" sorts before "key", so these land at the front of the range
+        datastore_updates.insert(
+            format!("add{:04}", i).into_bytes(),
+            massa_models::types::SetOrDelete::Set(vec![1, 2, 3]),
+        );
+    }
+
+    let mut update_changes = PreHashMap::default();
+    update_changes.insert(
+        addr,
+        massa_models::types::SetUpdateOrDelete::Update(LedgerEntryUpdate {
+            datastore: datastore_updates.clone(),
+            ..Default::default()
+        }),
+    );
+
+    let mk_output = |slot, ledger_changes| ExecutionOutput {
+        slot,
+        block_info: None,
+        state_changes: StateChanges {
+            ledger_changes,
+            async_pool_changes: Default::default(),
+            deferred_call_changes: Default::default(),
+            pos_changes: Default::default(),
+            executed_ops_changes: Default::default(),
+            executed_denunciations_changes: Default::default(),
+            execution_trail_hash_change: Default::default(),
+        },
+        events: Default::default(),
+        #[cfg(feature = "execution-trace")]
+        slot_trace: Default::default(),
+        #[cfg(feature = "dump-block")]
+        storage: None,
+        deferred_credits_execution: Default::default(),
+        cancel_async_message_execution: Default::default(),
+        auto_sell_execution: Default::default(),
+        transfers_history: Default::default(),
+        execution_info: None,
+    };
+
+    let active_history = Arc::new(RwLock::new(ActiveHistory(VecDeque::from([
+        mk_output(Slot::new(1, 0), LedgerChanges(changes)),
+        mk_output(Slot::new(2, 0), LedgerChanges(update_changes)),
+    ]))));
+
+    let scan = |count| {
+        scan_datastore(
+            &addr,
+            &[],
+            Bound::Unbounded,
+            Bound::Unbounded,
+            count,
+            foreign_controllers.final_state.clone(),
+            active_history.clone(),
+            None,
+        )
+        .1
+        .expect("expected candidate keys")
+    };
+
+    let unbounded: Vec<_> = scan(None).into_iter().collect();
+
+    // sanity: the deletions really do force the scan past the first keys of the entry
+    assert!(unbounded.len() > nb_keys / 2);
+
+    for count in 0..=(unbounded.len() + 2) {
+        let bounded: Vec<_> = scan(Some(count as u32)).into_iter().collect();
+        let expected: Vec<_> = unbounded.iter().take(count).cloned().collect();
+        assert_eq!(
+            bounded, expected,
+            "count={} produced a different result than the unbounded scan",
+            count
+        );
+    }
+}

massa-models/src/datastore.rs

@@ -236,7 +236,10 @@ pub fn range_intersection<T: Ord>(
 }
 
 /// Checks and cleans up a datastore key range query
-/// Returns: (prefix, start_bound, end_bound) or error
+/// Returns: (prefix, start_bound, end_bound, count) or error
+/// The returned count is the effective item count limit: it falls back to the
+/// configured maximum when the caller did not provide one, so that callers can
+/// forward it to the datastore scan instead of running an unbounded enumeration.
 /// Note: only useful to cleanup user-supplied requests (API/ABI)
 #[allow(clippy::type_complexity)]
 pub fn cleanup_datastore_key_range_query(
@@ -246,7 +249,7 @@ pub fn cleanup_datastore_key_range_query(
     count: Option<u32>,
     max_datastore_key_length: u8,
     max_datastore_query_config: Option<u32>,
-) -> Result<(Vec<u8>, Bound<Vec<u8>>, Bound<Vec<u8>>), ModelsError> {
+) -> Result<(Vec<u8>, Bound<Vec<u8>>, Bound<Vec<u8>>, Option<u32>), ModelsError> {
     // check item count
     let count = count.or(max_datastore_query_config);
     if let (Some(cnt), Some(max_cnt)) = (count.as_ref(), max_datastore_query_config.as_ref()) {
@@ -265,6 +268,7 @@ pub fn cleanup_datastore_key_range_query(
             Vec::new(),
             std::ops::Bound::Excluded(Vec::new()),
             std::ops::Bound::Excluded(Vec::new()),
+            count,
         ));
     } else {
         prefix.to_vec()
@@ -308,7 +312,7 @@ pub fn cleanup_datastore_key_range_query(
         }
     };
 
-    Ok((prefix, start_bound, end_bound))
+    Ok((prefix, start_bound, end_bound, count))
 }
 
 #[cfg(test)]
@@ -497,10 +501,11 @@ mod tests {
             max_query_config,
         );
         assert!(result.is_ok());
-        let (res_prefix, res_start, res_end) = result.unwrap();
+        let (res_prefix, res_start, res_end, res_count) = result.unwrap();
         assert_eq!(res_prefix, prefix);
         assert_eq!(res_start, start_bound);
         assert_eq!(res_end, end_bound);
+        assert_eq!(res_count, count);
 
         // Case 2: Prefix length exceeds max length
         let long_prefix = vec![b'a'; 30];
@@ -513,8 +518,9 @@ mod tests {
             None,
         );
         assert!(result.is_ok());
-        let (res_prefix, res_start, res_end) = result.unwrap();
+        let (res_prefix, res_start, res_end, res_count) = result.unwrap();
         assert!(res_prefix.is_empty());
+        assert_eq!(res_count, None);
         assert_eq!(res_start, Bound::Excluded(Vec::new()));
         assert_eq!(res_end, Bound::Excluded(Vec::new()));
 
@@ -531,8 +537,9 @@ mod tests {
             None,
         );
         assert!(result.is_ok());
-        let (res_prefix, res_start, res_end) = result.unwrap();
+        let (res_prefix, res_start, res_end, res_count) = result.unwrap();
         assert_eq!(res_prefix, prefix);
+        assert_eq!(res_count, None);
         assert_eq!(
             res_start,
             Bound::Excluded(long_key[0..10].to_vec()) // Start key truncated
@@ -570,9 +577,25 @@ mod tests {
             None,
         );
         assert!(result.is_ok());
-        let (res_prefix, res_start, res_end) = result.unwrap();
+        let (res_prefix, res_start, res_end, res_count) = result.unwrap();
         assert_eq!(res_prefix, prefix);
         assert_eq!(res_start, Bound::Unbounded);
         assert_eq!(res_end, Bound::Unbounded);
+        assert_eq!(res_count, None);
+
+        // Case 6: No count provided but a max query config is set:
+        // the configured maximum is used as the effective count so that the
+        // datastore scan stays bounded.
+        let result = cleanup_datastore_key_range_query(
+            &prefix,
+            Bound::Unbounded,
+            Bound::Unbounded,
+            None,
+            10,
+            Some(50),
+        );
+        assert!(result.is_ok());
+        let (_res_prefix, _res_start, _res_end, res_count) = result.unwrap();
+        assert_eq!(res_count, Some(50));
     }
 }