Skip to content

Commit 0ea9d7c

Browse files
authored
feat/store: prune account code (#1772)
1 parent 361e52c commit 0ea9d7c

File tree

4 files changed

+271
-8
lines changed

4 files changed

+271
-8
lines changed

crates/store/src/db/migrations/2025062000000_setup/up.sql

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ CREATE INDEX idx_accounts_created_at_block ON accounts(created_at_block);
4444
CREATE INDEX idx_accounts_block_num ON accounts(block_num);
4545
-- Index for joining with account_codes
4646
CREATE INDEX idx_accounts_code_commitment ON accounts(code_commitment) WHERE code_commitment IS NOT NULL;
47+
-- Covering index for the prune_account_codes subquery: filters rows by block_num/is_latest and projects code_commitment
48+
CREATE INDEX idx_accounts_prune_code ON accounts(block_num, is_latest, code_commitment) WHERE code_commitment IS NOT NULL;
4749

4850
CREATE TABLE notes (
4951
committed_at INTEGER NOT NULL, -- Block number when the note was committed

crates/store/src/db/models/queries/accounts.rs

Lines changed: 52 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1457,18 +1457,19 @@ pub(crate) struct AccountStorageMapRowInsert {
14571457
// CLEANUP FUNCTIONS
14581458
// ================================================================================================
14591459

1460-
/// Number of historical blocks to retain for vault assets and storage map values.
1460+
/// Number of historical blocks to retain for vault assets, storage map values, and account codes.
14611461
/// Entries older than `chain_tip - HISTORICAL_BLOCK_RETENTION` will be deleted,
14621462
/// except for entries marked with `is_latest=true` which are always retained.
14631463
pub const HISTORICAL_BLOCK_RETENTION: u32 = 50;
14641464

14651465
/// Clean up old entries for all accounts, deleting entries older than the retention window.
14661466
///
1467-
/// Deletes rows where `block_num < chain_tip - HISTORICAL_BLOCK_RETENTION` and `is_latest = false`.
1468-
/// This is a simple and efficient approach that doesn't require window functions.
1467+
/// Deletes rows where `block_num < chain_tip - HISTORICAL_BLOCK_RETENTION` and `is_latest = false`
1468+
/// for vault assets and storage map values. Also deletes account codes that are no longer
1469+
/// referenced by any account row within the retention window.
14691470
///
14701471
/// # Returns
1471-
/// A tuple of `(vault_assets_deleted, storage_map_values_deleted)`
1472+
/// A tuple of `(vault_assets_deleted, storage_map_values_deleted, account_codes_deleted)`
14721473
#[tracing::instrument(
14731474
target = COMPONENT,
14741475
skip_all,
@@ -1478,13 +1479,14 @@ pub const HISTORICAL_BLOCK_RETENTION: u32 = 50;
14781479
pub(crate) fn prune_history(
14791480
conn: &mut SqliteConnection,
14801481
chain_tip: BlockNumber,
1481-
) -> Result<(usize, usize), DatabaseError> {
1482+
) -> Result<(usize, usize, usize), DatabaseError> {
14821483
let cutoff_block = i64::from(chain_tip.as_u32().saturating_sub(HISTORICAL_BLOCK_RETENTION));
14831484
tracing::Span::current().record("cutoff_block", cutoff_block);
14841485
let vault_deleted = prune_account_vault_assets(conn, cutoff_block)?;
14851486
let storage_deleted = prune_account_storage_map_values(conn, cutoff_block)?;
1487+
let codes_deleted = prune_account_codes(conn, cutoff_block)?;
14861488

1487-
Ok((vault_deleted, storage_deleted))
1489+
Ok((vault_deleted, storage_deleted, codes_deleted))
14881490
}
14891491

14901492
#[tracing::instrument(
@@ -1528,3 +1530,47 @@ fn prune_account_storage_map_values(
15281530
.execute(conn)
15291531
.map_err(DatabaseError::Diesel)
15301532
}
1533+
1534+
/// Deletes account codes that are no longer referenced by any account row within the retention
1535+
/// window.
1536+
///
1537+
/// An account code is safe to delete when no `accounts` row with `block_num >= cutoff_block`
1538+
/// references its `code_commitment`. This covers both active accounts (`is_latest=true`) and
1539+
/// recent historical rows that still fall within the retention window.
1540+
///
1541+
/// # Raw SQL
1542+
///
1543+
/// ```sql
1544+
/// DELETE FROM account_codes
1545+
/// WHERE code_commitment NOT IN (
1546+
/// SELECT DISTINCT code_commitment
1547+
/// FROM accounts
1548+
/// WHERE code_commitment IS NOT NULL
1549+
/// AND (block_num >= ?1 OR is_latest = 1)
1550+
/// )
1551+
/// ```
1552+
#[tracing::instrument(
1553+
target = COMPONENT,
1554+
skip_all,
1555+
err,
1556+
fields(cutoff_block),
1557+
)]
1558+
fn prune_account_codes(
1559+
conn: &mut SqliteConnection,
1560+
cutoff_block: i64,
1561+
) -> Result<usize, DatabaseError> {
1562+
use diesel::sql_types::BigInt;
1563+
1564+
diesel::sql_query(
1565+
"DELETE FROM account_codes \
1566+
WHERE code_commitment NOT IN ( \
1567+
SELECT DISTINCT code_commitment \
1568+
FROM accounts \
1569+
WHERE code_commitment IS NOT NULL \
1570+
AND (block_num >= ?1 OR is_latest = 1 ) \
1571+
)",
1572+
)
1573+
.bind::<BigInt, _>(cutoff_block)
1574+
.execute(conn)
1575+
.map_err(DatabaseError::Diesel)
1576+
}

crates/store/src/db/models/queries/accounts/tests.rs

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1117,3 +1117,217 @@ fn test_select_account_vault_at_block_with_deletion() {
11171117
assert_eq!(assets_at_block_3.len(), 1, "Should have 1 asset at block 3");
11181118
assert_matches!(&assets_at_block_3[0], Asset::Fungible(f) if f.amount() == 2000);
11191119
}
1120+
1121+
// ACCOUNT CODE PRUNING TESTS
1122+
// ================================================================================================
1123+
1124+
/// Counts the number of rows in `account_codes`.
1125+
fn count_account_codes(conn: &mut SqliteConnection) -> usize {
1126+
use schema::account_codes;
1127+
1128+
let val =
1129+
SelectDsl::select(account_codes::table, diesel::dsl::count(account_codes::code_commitment))
1130+
.get_result::<i64>(conn)
1131+
.expect("Failed to count account_codes");
1132+
usize::try_from(u64::try_from(val).unwrap()).unwrap()
1133+
}
1134+
1135+
/// Returns whether a specific code commitment exists in `account_codes`.
1136+
fn account_code_exists(conn: &mut SqliteConnection, code_commitment: Word) -> bool {
1137+
use schema::account_codes;
1138+
1139+
let n =
1140+
SelectDsl::select(account_codes::table, diesel::dsl::count(account_codes::code_commitment))
1141+
.filter(account_codes::code_commitment.eq(code_commitment.to_bytes()))
1142+
.get_result::<i64>(conn)
1143+
.expect("Failed to query account_codes");
1144+
1145+
n == 1
1146+
}
1147+
1148+
/// Creates a full-state [`BlockAccountUpdate`] for the given account.
1149+
fn make_full_state_update(account: &Account) -> BlockAccountUpdate {
1150+
let delta = AccountDelta::try_from(account.clone()).unwrap();
1151+
assert!(delta.is_full_state(), "expected full-state delta");
1152+
BlockAccountUpdate::new(
1153+
account.id(),
1154+
account.to_commitment(),
1155+
AccountUpdateDetails::Delta(delta),
1156+
)
1157+
}
1158+
1159+
/// Builds a public account using a fixed account ID seed but a different component code.
1160+
///
1161+
/// All accounts produced here share the same [`AccountId`] because the same seed is used.
1162+
/// The `push_value` must be different for each variant to produce a distinct MAST root and thus a
1163+
/// distinct [`AccountCode::commitment`].
1164+
fn build_account_with_code(push_value: u32) -> Account {
1165+
let code_src = format!("pub proc variant push.{push_value} end");
1166+
let component_code = CodeBuilder::default()
1167+
.compile_component_code("test::code_prune", &code_src)
1168+
.unwrap();
1169+
let component = AccountComponent::new(
1170+
component_code,
1171+
vec![StorageSlot::with_value(
1172+
StorageSlotName::mock(0),
1173+
Word::from([Felt::new(1), Felt::ZERO, Felt::ZERO, Felt::ZERO]),
1174+
)],
1175+
AccountComponentMetadata::new("code_prune_test")
1176+
.with_supported_type(AccountType::RegularAccountUpdatableCode),
1177+
)
1178+
.unwrap();
1179+
1180+
// Seed [2u8; 32] keeps the account ID distinct from the other test helpers.
1181+
AccountBuilder::new([2u8; 32])
1182+
.account_type(AccountType::RegularAccountUpdatableCode)
1183+
.storage_mode(AccountStorageMode::Public)
1184+
.with_component(component)
1185+
.with_auth_component(AuthSingleSig::new(
1186+
PublicKeyCommitment::from(EMPTY_WORD),
1187+
AuthScheme::Falcon512Rpo,
1188+
))
1189+
.build_existing()
1190+
.unwrap()
1191+
}
1192+
1193+
/// Prune test 2: when an account's code changes, the old code must be pruned after the retention
1194+
/// window, while the new (latest) code is retained.
1195+
#[test]
1196+
fn test_prune_account_code_retains_latest_after_code_change() {
1197+
let mut conn = setup_test_db();
1198+
1199+
// Block 0: account created with code A.
1200+
// Block RETENTION+1 (=51): account updated to code B — within the retention window at prune
1201+
// time.
1202+
// Block 2*RETENTION+1 (=101): prune → cutoff is block 51; code A (last at block 0) is outside
1203+
// the window → pruned; code B (last at block 51) is within the window → retained.
1204+
let block_0 = BlockNumber::from(0u32);
1205+
let block_code_b = BlockNumber::from(HISTORICAL_BLOCK_RETENTION + 1);
1206+
let block_prunable = BlockNumber::from(2 * HISTORICAL_BLOCK_RETENTION + 1);
1207+
1208+
insert_block_header(&mut conn, block_0);
1209+
insert_block_header(&mut conn, block_code_b);
1210+
insert_block_header(&mut conn, block_prunable);
1211+
1212+
let account_a = build_account_with_code(1);
1213+
let account_b = build_account_with_code(2);
1214+
1215+
// Both accounts must have the same ID for this to test code replacement.
1216+
assert_eq!(account_a.id(), account_b.id(), "accounts must share the same ID");
1217+
assert_ne!(
1218+
account_a.code().commitment(),
1219+
account_b.code().commitment(),
1220+
"accounts must have different codes"
1221+
);
1222+
1223+
let account_id = account_a.id();
1224+
let code_commitment_a = account_a.code().commitment();
1225+
let code_commitment_b = account_b.code().commitment();
1226+
1227+
// Block 0: insert account with code A.
1228+
upsert_accounts(&mut conn, &[make_full_state_update(&account_a)], block_0)
1229+
.expect("initial upsert failed");
1230+
1231+
// Block RETENTION+1: update the same account ID to code B via a full-state delta.
1232+
upsert_accounts(&mut conn, &[make_full_state_update(&account_b)], block_code_b)
1233+
.expect("code-change upsert failed");
1234+
1235+
assert_eq!(count_account_codes(&mut conn), 2, "both codes must exist before pruning");
1236+
1237+
// Advance past retention window and prune.
1238+
// cutoff = block_prunable - RETENTION = 2*RETENTION+1 - RETENTION = RETENTION+1 = block_code_b
1239+
let (_, _, codes_deleted) =
1240+
prune_history(&mut conn, block_prunable).expect("prune_history failed");
1241+
1242+
// Only code A was dropped; code B is still referenced by the latest accounts row.
1243+
assert_eq!(codes_deleted, 1, "exactly one code (A) must be pruned");
1244+
assert!(!account_code_exists(&mut conn, code_commitment_a), "old code A must be pruned");
1245+
assert!(
1246+
account_code_exists(&mut conn, code_commitment_b),
1247+
"current code B must be retained"
1248+
);
1249+
1250+
// Confirm the latest account row still points to code B.
1251+
let (latest_header, _) =
1252+
select_account_header_with_storage_header_at_block(&mut conn, account_id, block_prunable)
1253+
.expect("query failed")
1254+
.expect("account must still exist");
1255+
assert_eq!(
1256+
latest_header.code_commitment(),
1257+
account_b.code().commitment(),
1258+
"latest account must reference code B"
1259+
);
1260+
}
1261+
1262+
/// Prune test 3: code A → code B → code A; after the retention window, code B must be pruned
1263+
/// but code A must be retained because it is still the latest.
1264+
#[test]
1265+
fn test_prune_account_code_retains_revisited_code() {
1266+
let mut conn = setup_test_db();
1267+
1268+
// Block 0: code A.
1269+
// Block RETENTION+1: code B (will be outside retention window at prune time).
1270+
// Block RETENTION+2: back to code A (within the retention window at prune time).
1271+
// Block 2*RETENTION+2: prune.
1272+
// cutoff = 2*RETENTION+2 - RETENTION = RETENTION+2.
1273+
// Code A: last referenced at block RETENTION+2 >= cutoff → retained.
1274+
// Code B: last referenced at block RETENTION+1 < cutoff → pruned.
1275+
let block_0 = BlockNumber::from(0u32);
1276+
let block_code_b = BlockNumber::from(HISTORICAL_BLOCK_RETENTION + 1);
1277+
let block_code_a_again = BlockNumber::from(HISTORICAL_BLOCK_RETENTION + 2);
1278+
let block_prunable = BlockNumber::from(2 * HISTORICAL_BLOCK_RETENTION + 2);
1279+
1280+
insert_block_header(&mut conn, block_0);
1281+
insert_block_header(&mut conn, block_code_b);
1282+
insert_block_header(&mut conn, block_code_a_again);
1283+
insert_block_header(&mut conn, block_prunable);
1284+
1285+
let account_a = build_account_with_code(1);
1286+
let account_b = build_account_with_code(2);
1287+
1288+
assert_eq!(account_a.id(), account_b.id(), "accounts must share the same ID");
1289+
assert_ne!(
1290+
account_a.code().commitment(),
1291+
account_b.code().commitment(),
1292+
"accounts must have different codes"
1293+
);
1294+
1295+
let account_id = account_a.id();
1296+
let code_commitment_a = account_a.code().commitment();
1297+
let code_commitment_b = account_b.code().commitment();
1298+
1299+
// Block 0: code A.
1300+
upsert_accounts(&mut conn, &[make_full_state_update(&account_a)], block_0)
1301+
.expect("block 0 upsert failed");
1302+
// Block RETENTION+1: code B.
1303+
upsert_accounts(&mut conn, &[make_full_state_update(&account_b)], block_code_b)
1304+
.expect("block code_b upsert failed");
1305+
// Block RETENTION+2: back to code A.
1306+
upsert_accounts(&mut conn, &[make_full_state_update(&account_a)], block_code_a_again)
1307+
.expect("block code_a_again upsert failed");
1308+
1309+
// Before pruning: both codes must be in account_codes (code A inserted once via ON CONFLICT DO
1310+
// NOTHING, code B inserted once).
1311+
assert_eq!(count_account_codes(&mut conn), 2, "both codes must exist before pruning");
1312+
1313+
// Advance past retention window and prune.
1314+
let (_, _, codes_deleted) =
1315+
prune_history(&mut conn, block_prunable).expect("prune_history failed");
1316+
1317+
// Code B is no longer referenced by any account row within the retention window → pruned.
1318+
// Code A is still referenced by the block_code_a_again accounts row (within cutoff) → retained.
1319+
assert_eq!(codes_deleted, 1, "exactly one code (B) must be pruned");
1320+
assert!(account_code_exists(&mut conn, code_commitment_a), "code A must be retained");
1321+
assert!(!account_code_exists(&mut conn, code_commitment_b), "code B must be pruned");
1322+
1323+
// Confirm the latest account row still points to code A.
1324+
let (latest_header, _) =
1325+
select_account_header_with_storage_header_at_block(&mut conn, account_id, block_prunable)
1326+
.expect("query failed")
1327+
.expect("account must still exist");
1328+
assert_eq!(
1329+
latest_header.code_commitment(),
1330+
account_a.code().commitment(),
1331+
"latest account must reference code A"
1332+
);
1333+
}

crates/store/src/db/tests.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2540,7 +2540,8 @@ fn test_prune_history() {
25402540

25412541
// Run cleanup with chain_tip = block_tip, cutoff will be block_tip - HISTORICAL_BLOCK_RETENTION
25422542
// = block_cutoff
2543-
let (vault_deleted, storage_deleted) = queries::prune_history(conn, block_tip).unwrap();
2543+
let (vault_deleted, storage_deleted, _codes_deleted) =
2544+
queries::prune_history(conn, block_tip).unwrap();
25442545

25452546
// Verify deletions occurred
25462547
assert_eq!(vault_deleted, 1, "should delete 1 old vault asset");
@@ -2620,7 +2621,7 @@ fn test_prune_history() {
26202621

26212622
// This entry at block 0 is marked as is_latest=true by insert_account_vault_asset
26222623
// Run cleanup again
2623-
let (vault_deleted_2, _) = queries::prune_history(conn, block_tip).unwrap();
2624+
let (vault_deleted_2, ..) = queries::prune_history(conn, block_tip).unwrap();
26242625

26252626
// The old latest entry should not be deleted (vault_deleted_2 should be 0)
26262627
assert_eq!(vault_deleted_2, 0, "should not delete any is_latest=true entries");

0 commit comments

Comments
 (0)