Skip to content

Commit c849165

Browse files
ShuoWangNSLShuoWangNSL
authored
chore: [MR-620] Keep extra in-memory states (#2061)
The is the first PR towards removing in-memory state more eagerly at checkpointed heights. Currently, in-memory state at the previous checkpoint height is kept until the next CUP, roughly the whole checkpointing interval. We plan to remove it more eagerly and only keep it until the current CUP. Currently,`remove_inmemory_states_below` purges states based on the certified height referenced in the finalized tip. As certified height could go beyond the summary block height, we need to keep certain states for several rounds more until the current CUP is created. Thus Consensus needs to inform state manager of the extra heights to keep through a new API . This PR extends the interface, adds implementation and tests in state manager. Next, Consensus will use this API with proper arguments, i.e. filling in all the heights required for CUP creation and validation, and run end-to-end tests. After these two PRs, we can proceed with removing in-memory states more eagerly at checkpointing heights.
1 parent 218fdbc commit c849165

File tree

6 files changed

+165
-24
lines changed

6 files changed

+165
-24
lines changed

‎rs/consensus/src/consensus/purger.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
//!
1919
//! 4. Replicated states below the certified height recorded in the block
2020
//! in the latest CatchUpPackage can be purged.
21+
use super::{bounds::validated_pool_within_bounds, MINIMUM_CHAIN_LENGTH};
2122
use crate::consensus::metrics::PurgerMetrics;
2223
use ic_consensus_utils::pool_reader::PoolReader;
2324
use ic_interfaces::{
@@ -34,10 +35,9 @@ use ic_types::{
3435
replica_config::ReplicaConfig,
3536
Height,
3637
};
38+
use std::collections::BTreeSet;
3739
use std::{cell::RefCell, sync::Arc};
3840

39-
use super::{bounds::validated_pool_within_bounds, MINIMUM_CHAIN_LENGTH};
40-
4141
pub(crate) const VALIDATED_POOL_BOUNDS_CHECK_FREQUENCY: u64 = 10;
4242

4343
/// The Purger sub-component.
@@ -324,7 +324,8 @@ impl Purger {
324324
.certified_height
325325
.min(self.state_manager.latest_state_height());
326326

327-
self.state_manager.remove_inmemory_states_below(height);
327+
self.state_manager
328+
.remove_inmemory_states_below(height, &BTreeSet::new());
328329
trace!(
329330
self.log,
330331
"Purge replicated states below [memory] {:?}",
@@ -502,7 +503,9 @@ mod tests {
502503
state_manager
503504
.get_mut()
504505
.expect_remove_inmemory_states_below()
505-
.withf(move |height| *height == *inmemory_purge_height_clone.read().unwrap())
506+
.withf(move |height, _extra_heights| {
507+
*height == *inmemory_purge_height_clone.read().unwrap()
508+
})
506509
.return_const(());
507510

508511
state_manager

‎rs/interfaces/state_manager/mocks/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ mock! {
4848

4949
fn remove_states_below(&self, height: Height);
5050

51-
fn remove_inmemory_states_below(&self, height: Height);
51+
fn remove_inmemory_states_below(&self, height: Height, extra_heights_to_keep: &std::collections::BTreeSet<Height>);
5252

5353
fn commit_and_certify(
5454
&self,

‎rs/interfaces/state_manager/src/lib.rs

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
//! The state manager public interface.
2+
23
use ic_crypto_tree_hash::{LabeledTree, MixedHashTree};
34
use ic_types::{
45
batch::BatchSummary, consensus::certification::Certification, CryptoHashOfPartialState,
56
CryptoHashOfState, Height,
67
};
78
use phantom_newtype::BitMask;
9+
use std::collections::BTreeSet;
810
use std::sync::Arc;
911
use thiserror::Error;
1012

@@ -233,15 +235,20 @@ pub trait StateManager: StateReader {
233235
fn remove_states_below(&self, height: Height);
234236

235237
/// Notify the state manager that states committed with partial certification
236-
/// state and heights strictly less than specified `height` can be removed.
238+
/// state and heights strictly less than the specified `height` can be removed, except
239+
/// for any heights provided in `extra_heights_to_keep`, which will still be retained.
237240
///
238241
/// Note that:
239242
/// * The initial state (height = 0) cannot be removed.
240243
/// * Some states matching the removal criteria might be kept alive. For
241244
/// example, the last fully persisted state might be preserved to
242245
/// optimize future operations.
243246
/// * No checkpoints are removed, see also `remove_states_below()`
244-
fn remove_inmemory_states_below(&self, height: Height);
247+
fn remove_inmemory_states_below(
248+
&self,
249+
height: Height,
250+
extra_heights_to_keep: &BTreeSet<Height>,
251+
);
245252

246253
/// Commits the `state` at given `height`, limits the certification to
247254
/// `scope`. The `state` must be the mutable state obtained via a call to

‎rs/state_manager/src/lib.rs

Lines changed: 39 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2181,7 +2181,8 @@ impl StateManagerImpl {
21812181
.expect("Failed to receive deallocation notification");
21822182
}
21832183

2184-
/// Remove any inmemory state at height h with h < last_height_to_keep, and
2184+
/// Remove any inmemory state at height h with h < last_height_to_keep
2185+
/// except for any heights provided in `extra_inmemory_heights_to_keep`, and
21852186
/// any checkpoint at height h < last_checkpoint_to_keep
21862187
///
21872188
/// Shared inner function of the public functions remove_states_below
@@ -2190,6 +2191,7 @@ impl StateManagerImpl {
21902191
&self,
21912192
last_height_to_keep: Height,
21922193
last_checkpoint_to_keep: Height,
2194+
extra_inmemory_heights_to_keep: &BTreeSet<Height>,
21932195
) {
21942196
debug_assert!(
21952197
last_height_to_keep >= last_checkpoint_to_keep,
@@ -2223,6 +2225,19 @@ impl StateManagerImpl {
22232225
state_metadata.bundled_manifest.as_ref().map(|_| *height)
22242226
});
22252227

2228+
// The `extra_inmemory_heights_to_keep` is used for preserving in-memory states,
2229+
// but it can safely be included in the `heights_to_keep` set, which retains both
2230+
// in-memory states and checkpoints. This is safe because:
2231+
//
2232+
// 1. When called by `remove_inmemory_states_below`, checkpoints are never removed,
2233+
// regardless of the inclusion of `extra_inmemory_heights_to_keep`, so no harm
2234+
// or unnecessary preservation occurs.
2235+
//
2236+
// 2. When called by `remove_states_below`, `extra_inmemory_heights_to_keep` is always
2237+
// an empty set, having no effect on the outcome.
2238+
//
2239+
// In the future, separating these sets could clarify their distinct purposes and
2240+
// simplify reasoning about correctness without relying heavily on input behavior.
22262241
let heights_to_keep: BTreeSet<Height> = states
22272242
.states_metadata
22282243
.keys()
@@ -2232,6 +2247,7 @@ impl StateManagerImpl {
22322247
})
22332248
.chain(std::iter::once(latest_certified_height))
22342249
.chain(latest_manifest_height)
2250+
.chain(extra_inmemory_heights_to_keep.iter().copied())
22352251
.collect();
22362252

22372253
// Send object to deallocation thread if it has capacity.
@@ -2368,9 +2384,11 @@ impl StateManagerImpl {
23682384

23692385
let state_heights = self.list_state_heights(CERT_ANY);
23702386

2371-
debug_assert!(heights_to_keep.iter().all(|h| unfiltered_checkpoint_heights
2372-
.contains(h)
2373-
|| *h == latest_certified_height));
2387+
debug_assert!(heights_to_keep
2388+
.iter()
2389+
.all(|h| unfiltered_checkpoint_heights.contains(h)
2390+
|| extra_inmemory_heights_to_keep.contains(h)
2391+
|| *h == latest_certified_height));
23742392

23752393
debug_assert!(state_heights.contains(&latest_state_height));
23762394
debug_assert!(state_heights.contains(&latest_certified_height));
@@ -3298,7 +3316,12 @@ impl StateManager for StateManagerImpl {
32983316
.min(oldest_checkpoint_to_keep)
32993317
};
33003318

3301-
self.remove_states_below_impl(oldest_height_to_keep, oldest_checkpoint_to_keep);
3319+
// The public interface does not protect extra states, so we pass an empty set here.
3320+
self.remove_states_below_impl(
3321+
oldest_height_to_keep,
3322+
oldest_checkpoint_to_keep,
3323+
&BTreeSet::new(),
3324+
);
33023325
}
33033326

33043327
/// Variant of `remove_states_below()` that only removes states committed with
@@ -3310,7 +3333,12 @@ impl StateManager for StateManagerImpl {
33103333
/// * The latest state
33113334
/// * The latest certified state
33123335
/// * State 0
3313-
fn remove_inmemory_states_below(&self, requested_height: Height) {
3336+
/// * Specified extra heights to keep
3337+
fn remove_inmemory_states_below(
3338+
&self,
3339+
requested_height: Height,
3340+
extra_heights_to_keep: &BTreeSet<Height>,
3341+
) {
33143342
let _timer = self
33153343
.metrics
33163344
.api_call_duration
@@ -3323,7 +3351,11 @@ impl StateManager for StateManagerImpl {
33233351
.min(requested_height)
33243352
.max(Height::new(1));
33253353

3326-
self.remove_states_below_impl(oldest_height_to_keep, Self::INITIAL_STATE_HEIGHT);
3354+
self.remove_states_below_impl(
3355+
oldest_height_to_keep,
3356+
Self::INITIAL_STATE_HEIGHT,
3357+
extra_heights_to_keep,
3358+
);
33273359
}
33283360

33293361
fn commit_and_certify(

‎rs/state_manager/tests/state_manager.rs

Lines changed: 97 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ use ic_types::{
7171
CanisterId, CryptoHashOfPartialState, CryptoHashOfState, Height, NodeId, NumBytes, PrincipalId,
7272
};
7373
use ic_types::{epoch_from_height, QueryStatsEpoch};
74-
use maplit::btreemap;
74+
use maplit::{btreemap, btreeset};
7575
use nix::sys::time::TimeValLike;
7676
use nix::sys::{
7777
stat::{utimensat, UtimensatFlags},
@@ -1434,7 +1434,7 @@ fn cannot_remove_height_zero() {
14341434

14351435
state_manager.remove_states_below(height(0));
14361436
state_manager.flush_deallocation_channel();
1437-
state_manager.remove_inmemory_states_below(height(0));
1437+
state_manager.remove_inmemory_states_below(height(0), &BTreeSet::new());
14381438

14391439
assert_eq!(state_manager.list_state_heights(CERT_ANY), vec![height(0),],);
14401440

@@ -1448,7 +1448,7 @@ fn cannot_remove_height_zero() {
14481448

14491449
state_manager.remove_states_below(height(0));
14501450
state_manager.flush_deallocation_channel();
1451-
state_manager.remove_inmemory_states_below(height(0));
1451+
state_manager.remove_inmemory_states_below(height(0), &BTreeSet::new());
14521452

14531453
assert_eq!(
14541454
state_manager.list_state_heights(CERT_ANY),
@@ -1481,7 +1481,7 @@ fn cannot_remove_latest_height_or_checkpoint() {
14811481
// checkpoint can be retained until the hashing is complete.
14821482
state_manager.flush_tip_channel();
14831483
state_manager.remove_states_below(height(20));
1484-
state_manager.remove_inmemory_states_below(height(20));
1484+
state_manager.remove_inmemory_states_below(height(20), &BTreeSet::new());
14851485
state_manager.flush_deallocation_channel();
14861486

14871487
assert_eq!(
@@ -1504,7 +1504,7 @@ fn cannot_remove_latest_height_or_checkpoint() {
15041504

15051505
state_manager.flush_tip_channel();
15061506
state_manager.remove_states_below(height(20));
1507-
state_manager.remove_inmemory_states_below(height(20));
1507+
state_manager.remove_inmemory_states_below(height(20), &BTreeSet::new());
15081508
state_manager.flush_deallocation_channel();
15091509

15101510
assert_eq!(
@@ -1539,7 +1539,7 @@ fn can_remove_checkpoints_and_noncheckpoints_separately() {
15391539
state_manager.flush_tip_channel();
15401540

15411541
assert_eq!(state_manager.list_state_heights(CERT_ANY), heights);
1542-
state_manager.remove_inmemory_states_below(height(6));
1542+
state_manager.remove_inmemory_states_below(height(6), &BTreeSet::new());
15431543

15441544
// Only odd heights should have been removed
15451545
assert_eq!(
@@ -1579,6 +1579,97 @@ fn can_remove_checkpoints_and_noncheckpoints_separately() {
15791579
});
15801580
}
15811581

1582+
#[test]
1583+
fn remove_inmemory_states_below_can_keep_extra_states() {
1584+
state_manager_restart_test(|state_manager, restart_fn| {
1585+
let mut heights = vec![height(0)];
1586+
for i in 1..10 {
1587+
let (_height, state) = state_manager.take_tip();
1588+
heights.push(height(i));
1589+
1590+
let scope = if i % 2 == 0 {
1591+
CertificationScope::Full
1592+
} else {
1593+
CertificationScope::Metadata
1594+
};
1595+
1596+
state_manager.commit_and_certify(state, height(i), scope.clone(), None);
1597+
}
1598+
// We need to wait for hashing to complete, otherwise the
1599+
// checkpoint can be retained until the hashing is complete.
1600+
state_manager.flush_tip_channel();
1601+
1602+
assert_eq!(state_manager.list_state_heights(CERT_ANY), heights);
1603+
1604+
state_manager.remove_inmemory_states_below(height(5), &btreeset![height(1), height(7)]);
1605+
1606+
// State at height 1 is kept because of it is included in `extra_heights_to_keep`.
1607+
// The additional protection on the state at height 7 has no effect since it is above the requested height.
1608+
assert_eq!(
1609+
state_manager.list_state_heights(CERT_ANY),
1610+
vec![
1611+
height(0),
1612+
height(1),
1613+
height(2),
1614+
height(4),
1615+
height(5),
1616+
height(6),
1617+
height(7),
1618+
height(8),
1619+
height(9)
1620+
],
1621+
);
1622+
1623+
state_manager.remove_inmemory_states_below(height(9), &btreeset![height(7), height(8)]);
1624+
1625+
// State at height 7 is kept because of it is included in `extra_heights_to_keep`.
1626+
// The additional protection on the state at height 8 currently has no effect since it is a checkpoint.
1627+
// However, this may be subject to change in the future.
1628+
assert_eq!(
1629+
state_manager.list_state_heights(CERT_ANY),
1630+
vec![
1631+
height(0),
1632+
height(2),
1633+
height(4),
1634+
height(6),
1635+
height(7),
1636+
height(8),
1637+
height(9)
1638+
],
1639+
);
1640+
1641+
state_manager.remove_inmemory_states_below(height(9), &BTreeSet::new());
1642+
1643+
// State at height 7 is removed.
1644+
assert_eq!(
1645+
state_manager.list_state_heights(CERT_ANY),
1646+
vec![
1647+
height(0),
1648+
height(2),
1649+
height(4),
1650+
height(6),
1651+
height(8),
1652+
height(9)
1653+
],
1654+
);
1655+
1656+
state_manager.remove_states_below(height(8));
1657+
state_manager.flush_deallocation_channel();
1658+
1659+
assert_eq!(
1660+
state_manager.list_state_heights(CERT_ANY),
1661+
vec![height(0), height(8), height(9)],
1662+
);
1663+
1664+
let state_manager = restart_fn(state_manager, Some(height(8)));
1665+
1666+
assert_eq!(
1667+
state_manager.list_state_heights(CERT_ANY),
1668+
vec![height(0), height(8)],
1669+
);
1670+
});
1671+
}
1672+
15821673
#[test]
15831674
fn can_keep_last_checkpoint_and_higher_states_after_removal() {
15841675
state_manager_restart_test(|state_manager, restart_fn| {

‎rs/test_utilities/src/state_manager.rs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ use ic_types::{
2525
CryptoHashOfPartialState, CryptoHashOfState, Height, RegistryVersion, SubnetId,
2626
};
2727
use serde::{Deserialize, Serialize};
28-
use std::collections::VecDeque;
28+
use std::collections::{BTreeSet, VecDeque};
2929
use std::sync::{Arc, Barrier, RwLock};
3030

3131
#[derive(Clone)]
@@ -196,7 +196,11 @@ impl StateManager for FakeStateManager {
196196
.retain(|snap| snap.height == Height::new(0) || snap.height >= height)
197197
}
198198

199-
fn remove_inmemory_states_below(&self, _height: Height) {
199+
fn remove_inmemory_states_below(
200+
&self,
201+
_height: Height,
202+
_extra_heights_to_keep: &BTreeSet<Height>,
203+
) {
200204
// All heights are checkpoints
201205
}
202206

@@ -666,11 +670,15 @@ impl StateManager for RefMockStateManager {
666670
self.mock.read().unwrap().remove_states_below(height)
667671
}
668672

669-
fn remove_inmemory_states_below(&self, height: Height) {
673+
fn remove_inmemory_states_below(
674+
&self,
675+
height: Height,
676+
extra_heights_to_keep: &BTreeSet<Height>,
677+
) {
670678
self.mock
671679
.read()
672680
.unwrap()
673-
.remove_inmemory_states_below(height)
681+
.remove_inmemory_states_below(height, extra_heights_to_keep)
674682
}
675683

676684
fn commit_and_certify(

0 commit comments

Comments
 (0)