Skip to content

fix: Dereferencing freed memos when verifying provisional memos #788

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 9, 2025
Merged

fix: Dereferencing freed memos when verifying provisional memos #788

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0677f51
Add test
MichaReiser Apr 9, 2025
d33d824
Fix `fetch_hot` and `maybe_changed_after`
MichaReiser Apr 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions src/cycle.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,10 +132,6 @@ impl CycleHeads {
true
}

pub(crate) fn clear(&mut self) {
self.0.clear();
}

pub(crate) fn update_iteration_count(
&mut self,
cycle_head_index: DatabaseKeyIndex,
Expand Down
45 changes: 30 additions & 15 deletions src/function/fetch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,19 +79,24 @@ where
id: Id,
memo_ingredient_index: MemoIngredientIndex,
) -> Option<&'db Memo<C::Output<'db>>> {
let memo_guard = self.get_memo_from_table_for(zalsa, id, memo_ingredient_index);
if let Some(memo) = memo_guard {
let database_key_index = self.database_key_index(id);
if memo.value.is_some()
&& (self.validate_may_be_provisional(db, zalsa, database_key_index, memo)
|| self.validate_same_iteration(db, database_key_index, memo))
&& self.shallow_verify_memo(db, zalsa, database_key_index, memo)
{
// SAFETY: memo is present in memo_map and we have verified that it is
// still valid for the current revision.
return unsafe { Some(self.extend_memo_lifetime(memo)) };
}
let memo = self.get_memo_from_table_for(zalsa, id, memo_ingredient_index)?;

memo.value.as_ref()?;

let database_key_index = self.database_key_index(id);

let shallow_update = self.shallow_verify_memo(zalsa, database_key_index, memo)?;

if self.validate_may_be_provisional(db, zalsa, database_key_index, memo)
|| self.validate_same_iteration(db, database_key_index, memo)
{
self.update_shallow(db, zalsa, database_key_index, memo, shallow_update);

// SAFETY: memo is present in memo_map and we have verified that it is
// still valid for the current revision.
return unsafe { Some(self.extend_memo_lifetime(memo)) };
}

None
}

Expand Down Expand Up @@ -120,10 +125,20 @@ where
if let Some(memo) = memo_guard {
if memo.value.is_some()
&& memo.revisions.cycle_heads.contains(&database_key_index)
&& self.shallow_verify_memo(db, zalsa, database_key_index, memo)
{
// SAFETY: memo is present in memo_map.
return unsafe { Some(self.extend_memo_lifetime(memo)) };
if let Some(shallow_update) =
self.shallow_verify_memo(zalsa, database_key_index, memo)
{
self.update_shallow(
db,
zalsa,
database_key_index,
memo,
shallow_update,
);
// SAFETY: memo is present in memo_map.
return unsafe { Some(self.extend_memo_lifetime(memo)) };
}
}
}
// no provisional value; create/insert/return initial provisional value
Expand Down
97 changes: 67 additions & 30 deletions src/function/maybe_changed_after.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,16 @@ where

// Check if we have a verified version: this is the hot path.
let memo_guard = self.get_memo_from_table_for(zalsa, id, memo_ingredient_index);
if let Some(memo) = memo_guard {
if self.validate_may_be_provisional(db, zalsa, database_key_index, memo)
&& self.shallow_verify_memo(db, zalsa, database_key_index, memo)
{
let Some(memo) = memo_guard else {
// No memo? Assume has changed.
return VerifyResult::Changed;
};

if let Some(shallow_update) = self.shallow_verify_memo(zalsa, database_key_index, memo)
{
if self.validate_provisional(db, zalsa, database_key_index, memo) {
self.update_shallow(db, zalsa, database_key_index, memo, shallow_update);

return if memo.revisions.changed_at > revision {
VerifyResult::Changed
} else {
Expand All @@ -73,16 +79,14 @@ where
)
};
}
if let Some(mcs) =
self.maybe_changed_after_cold(zalsa, db, id, revision, memo_ingredient_index)
{
return mcs;
} else {
// We failed to claim, have to retry.
}
}

if let Some(mcs) =
self.maybe_changed_after_cold(zalsa, db, id, revision, memo_ingredient_index)
{
return mcs;
} else {
// No memo? Assume has changed.
return VerifyResult::Changed;
// We failed to claim, have to retry.
}
}
}
Expand Down Expand Up @@ -167,7 +171,7 @@ where
Some(VerifyResult::Changed)
}

/// True if the memo's value and `changed_at` time is still valid in this revision.
/// `Some` if the memo's value and `changed_at` time is still valid in this revision.
/// Does only a shallow O(1) check, doesn't walk the dependencies.
///
/// In general, a provisional memo (from cycle iteration) does not verify. Since we don't
Expand All @@ -177,11 +181,10 @@ where
#[inline]
pub(super) fn shallow_verify_memo(
&self,
db: &C::DbView,
zalsa: &Zalsa,
database_key_index: DatabaseKeyIndex,
memo: &Memo<C::Output<'_>>,
) -> bool {
) -> Option<ShallowUpdate> {
tracing::debug!(
"{database_key_index:?}: shallow_verify_memo(memo = {memo:#?})",
memo = memo.tracing_debug()
Expand All @@ -191,7 +194,7 @@ where

if verified_at == revision_now {
// Already verified.
return true;
return Some(ShallowUpdate::Verified);
}

let last_changed = zalsa.last_changed_revision(memo.revisions.durability);
Expand All @@ -204,17 +207,31 @@ where
);
if last_changed <= verified_at {
// No input of the suitable durability has changed since last verified.
Some(ShallowUpdate::HigherDurability(revision_now))
} else {
None
}
}

#[inline]
pub(super) fn update_shallow(
&self,
db: &C::DbView,
zalsa: &Zalsa,
database_key_index: DatabaseKeyIndex,
memo: &Memo<C::Output<'_>>,
update: ShallowUpdate,
) {
if let ShallowUpdate::HigherDurability(revision_now) = update {
memo.mark_as_verified(
db,
revision_now,
database_key_index,
memo.revisions.accumulated_inputs.load(),
);

memo.mark_outputs_as_verified(zalsa, db.as_dyn_database(), database_key_index);
return true;
}

false
}

/// Validates this memo if it is a provisional memo. Returns true for non provisional memos or
Expand Down Expand Up @@ -311,10 +328,15 @@ where
old_memo = old_memo.tracing_debug()
);

if self.validate_may_be_provisional(db, zalsa, database_key_index, old_memo)
&& self.shallow_verify_memo(db, zalsa, database_key_index, old_memo)
{
return VerifyResult::unchanged();
let shallow_update = self.shallow_verify_memo(zalsa, database_key_index, old_memo);
let shallow_update_possible = shallow_update.is_some();

if let Some(shallow_update) = shallow_update {
if self.validate_provisional(db, zalsa, database_key_index, old_memo) {
self.update_shallow(db, zalsa, database_key_index, old_memo, shallow_update);

return VerifyResult::unchanged();
}
}

match &old_memo.revisions.origin {
Expand All @@ -339,7 +361,9 @@ where
VerifyResult::Changed
}
QueryOrigin::Derived(edges) => {
if old_memo.may_be_provisional() {
let is_provisional = old_memo.may_be_provisional();
// If the value is from the same revision but is still provisional, consider it changed
if shallow_update_possible && is_provisional {
return VerifyResult::Changed;
}

Expand Down Expand Up @@ -428,15 +452,18 @@ where
inputs,
);

if is_provisional {
old_memo
.revisions
.verified_final
.store(true, Ordering::Relaxed);
}

if in_heads {
// Iterate our dependency graph again, starting from the top. We clear the
// cycle heads here because we are starting a fresh traversal. (It might be
// logically clearer to create a new HashSet each time, but clearing the
// existing one is more efficient.)
cycle_heads.clear();
continue 'cycle;
}
}

break 'cycle VerifyResult::Unchanged(
InputAccumulatedValues::Empty,
cycle_heads,
Expand All @@ -446,3 +473,13 @@ where
}
}
}

#[derive(Copy, Clone, Eq, PartialEq)]
pub(super) enum ShallowUpdate {
/// The memo is from this revision and has already been verified
Verified,

/// The revision for the memo's durability hasn't changed. It can be marked as verified
/// in this revision.
HigherDurability(Revision),
}
Loading