Skip to content

perf: defer database attachment for cached queries #1177

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

Draft
MichaReiser wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

perf: defer database attachment for cached queries #1177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 2 additions & 1 deletion components/salsa-macro-rules/src/setup_interned_struct.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,8 @@ macro_rules! setup_interned_struct {
)*
{
let (zalsa, zalsa_local) = db.zalsas();
$Configuration::ingredient(zalsa).intern(zalsa, zalsa_local,
$zalsa::assert_current_database(db);
$Configuration::ingredient(zalsa).intern(db, zalsa, zalsa_local,
StructKey::<$db_lt>($($field_id,)* ::std::marker::PhantomData::default()), |_, data| ($($zalsa::Lookup::into_owned(data.$field_index),)*))
}

Expand Down
33 changes: 17 additions & 16 deletions components/salsa-macro-rules/src/setup_tracked_fn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,23 +95,23 @@ macro_rules! setup_tracked_fn {
) -> ::salsa::plumbing::return_mode_ty!(($return_mode, __), $db_lt, $output_ty) {
use ::salsa::plumbing as $zalsa;

$zalsa::attach($db, || {
let (zalsa, zalsa_local) = $db.zalsas();
let result = $zalsa::macro_if! {
if $needs_interner {
{
let key = $fn_name::intern_ingredient_(zalsa).intern_id(zalsa, zalsa_local, ($($input_id),*), |_, data| data);
$fn_name::fn_ingredient_($db, zalsa).fetch($db, zalsa, zalsa_local, key)
}
} else {
{
$fn_name::fn_ingredient_($db, zalsa).fetch($db, zalsa, zalsa_local, $zalsa::AsId::as_id(&($($input_id),*)))
}
let (zalsa, zalsa_local) = $db.zalsas();
$zalsa::assert_current_database_or_attached($db, zalsa_local);

let result = $zalsa::macro_if! {
if $needs_interner {
{
let key = $fn_name::intern_ingredient_(zalsa).intern_id($db, zalsa, zalsa_local, ($($input_id),*), |_, data| data);
$fn_name::fn_ingredient_($db, zalsa).fetch($db, zalsa, zalsa_local, key)
}
};
} else {
{
$fn_name::fn_ingredient_($db, zalsa).fetch($db, zalsa, zalsa_local, $zalsa::AsId::as_id(&($($input_id),*)))
}
}
};

$zalsa::return_mode_expression!(($return_mode, __), $output_ty, result,)
})
$zalsa::return_mode_expression!(($return_mode, __), $output_ty, result,)
}

// The module needs be last in the macro expansion in order to make the tracked
Expand Down Expand Up @@ -458,10 +458,11 @@ macro_rules! setup_tracked_fn {
$($input_id: $interned_input_ty,)*
) -> Vec<&$db_lt A> {
use ::salsa::plumbing as $zalsa;
$zalsa::assert_current_database($db);
let key = $zalsa::macro_if! {
if $needs_interner {{
let (zalsa, zalsa_local) = $db.zalsas();
$Configuration::intern_ingredient(zalsa).intern_id(zalsa, zalsa_local, ($($input_id),*), |_, data| data)
$Configuration::intern_ingredient(zalsa).intern_id($db, zalsa, zalsa_local, ($($input_id),*), |_, data| data)
}} else {
$zalsa::AsId::as_id(&($($input_id),*))
}
Expand Down
145 changes: 126 additions & 19 deletions src/attach.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,23 @@ impl Attached {
}
}

#[inline]
fn is_current_database<Db>(&self, db: &Db) -> bool
where
Db: ?Sized + Database,
{
let Some(current_db) = self.database.get() else {
return false;
};

let new_db = NonNull::from(db.as_dyn_database());
if !std::ptr::addr_eq(current_db.as_ptr(), new_db.as_ptr()) {
panic!("Cannot change database mid-query. current: {current_db:?}, new: {new_db:?}");
}

true
}

#[inline]
fn attach<Db, R>(&self, db: &Db, op: impl FnOnce() -> R) -> R
where
Expand All @@ -49,23 +66,13 @@ impl Attached {
impl<'s> DbGuard<'s> {
#[inline]
fn new(attached: &'s Attached, db: &dyn Database) -> Self {
match attached.database.get() {
// A database is already attached, make sure it's the same as the new one.
Some(current_db) => {
let new_db = NonNull::from(db);
if !std::ptr::addr_eq(current_db.as_ptr(), new_db.as_ptr()) {
panic!(
"Cannot change database mid-query. current: {current_db:?}, new: {new_db:?}"
);
}
Self { state: None }
}
// No database is attached, attach the new one.
None => {
attached.database.set(Some(NonNull::from(db)));
Self {
state: Some(attached),
}
if attached.is_current_database(db) {
Self { state: None }
} else {
attached.database.set(Some(NonNull::from(db)));
db.zalsa_local().set_attached(true);
Self {
state: Some(attached),
}
}
}
Expand All @@ -78,7 +85,11 @@ impl Attached {
if let Some(attached) = self.state {
if let Some(prev) = attached.database.replace(None) {
// SAFETY: `prev` is a valid pointer to a database.
unsafe { prev.as_ref().zalsa_local().uncancel() };
unsafe {
let zalsa_local = prev.as_ref().zalsa_local();
zalsa_local.set_attached(false);
zalsa_local.uncancel();
}
}
}
}
Expand Down Expand Up @@ -119,6 +130,11 @@ impl Attached {
}
} else {
// and it was the a different one from ours, record the state changes.
// SAFETY: Both pointers remain valid for their attachment scopes.
unsafe {
prev.as_ref().zalsa_local().set_attached(false);
db.as_ref().zalsa_local().set_attached(true);
}
Self {
state: Some(attached),
prev: Some(prev),
Expand All @@ -128,6 +144,8 @@ impl Attached {
// No database is attached, attach the new one.
None => {
attached.database.set(Some(db));
// SAFETY: `db` remains valid for the attachment scope.
unsafe { db.as_ref().zalsa_local().set_attached(true) };
Self {
state: Some(attached),
prev: None,
Expand All @@ -144,7 +162,15 @@ impl Attached {
if let Some(attached) = self.state {
if let Some(prev) = attached.database.replace(self.prev) {
// SAFETY: `prev` is a valid pointer to a database.
unsafe { prev.as_ref().zalsa_local().uncancel() };
unsafe {
let zalsa_local = prev.as_ref().zalsa_local();
zalsa_local.set_attached(false);
zalsa_local.uncancel();
}
}
if let Some(prev) = self.prev {
// SAFETY: `prev` is valid for its enclosing attachment scope.
unsafe { prev.as_ref().zalsa_local().set_attached(true) };
}
}
}
Expand Down Expand Up @@ -179,6 +205,87 @@ where
)
}

/// Panics if a database other than `db` is currently attached.
#[doc(hidden)]
#[inline]
pub fn assert_current_database<Db>(db: &Db)
where
Db: ?Sized + Database,
{
ATTACHED.with(
#[inline]
|attached| {
attached.is_current_database(db);
},
)
}

/// Panics if a database other than `db` is currently attached, unless `zalsa_local` already
/// records that `db` is attached.
#[doc(hidden)]
#[inline(always)]
pub fn assert_current_database_or_attached<Db>(
db: &Db,
zalsa_local: &crate::zalsa_local::ZalsaLocal,
) where
Db: ?Sized + Database,
{
if !zalsa_local.is_attached() {
assert_current_database_unattached(db);
}
}

#[cold]
#[inline(never)]
fn assert_current_database_unattached<Db>(db: &Db)
where
Db: ?Sized + Database,
{
assert_current_database(db);
}

#[inline]
pub(crate) fn attach_if_needed<R, Db>(db: &Db, op: impl FnOnce() -> R) -> R
where
Db: ?Sized + Database,
{
ATTACHED.with(
#[inline]
|attached| {
if attached.is_current_database(db) {
op()
} else {
attach_cold(attached, db, op)
}
},
)
}

#[cold]
#[inline(never)]
fn attach_cold<R, Db>(attached: &Attached, db: &Db, op: impl FnOnce() -> R) -> R
where
Db: ?Sized + Database,
{
attached.attach(db, op)
}

#[inline]
pub(crate) fn is_attached<Db>(db: &Db) -> bool
where
Db: ?Sized + Database,
{
ATTACHED.with(
#[inline]
|attached| {
attached.database.get().is_some_and(|current_db| {
let db = NonNull::from(db.as_dyn_database());
std::ptr::addr_eq(current_db.as_ptr(), db.as_ptr())
})
},
)
}

/// Attach the database to the current thread and execute `op`.
/// Allows a different database than currently attached. The original database
/// will be restored on return.
Expand Down
3 changes: 2 additions & 1 deletion src/database.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,9 @@ pub trait Database: Send + ZalsaDatabase + AsDynDatabase {
/// `salsa_event` is emitted when this method is called, so that should be
/// used instead.
fn unwind_if_revision_cancelled(&self) {
crate::attach::assert_current_database(self);
let (zalsa, zalsa_local) = self.zalsas();
zalsa.unwind_if_revision_cancelled(zalsa_local);
zalsa.unwind_if_revision_cancelled(self, zalsa_local);
}

/// Execute `op` with the database in thread-local storage for debug print-outs.
Expand Down
2 changes: 2 additions & 0 deletions src/function/accumulated.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ where
let mut output = vec![];

// First ensure the result is up to date
crate::attach::assert_current_database(db);
self.fetch(db, zalsa, zalsa_local, key);

let db_key = self.database_key_index(key);
Expand Down Expand Up @@ -94,6 +95,7 @@ where
) -> (Option<&'db AccumulatedMap>, InputAccumulatedValues) {
let (zalsa, zalsa_local) = db.zalsas();
// NEXT STEP: stash and refactor `fetch` to return an `&Memo` so we can make this work
crate::attach::assert_current_database(db);
let memo = self.refresh_memo(db, zalsa, zalsa_local, key);
(
memo.header.revisions.accumulated(),
Expand Down
42 changes: 30 additions & 12 deletions src/function/fetch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::cycle::{CycleRecoveryStrategy, IterationStamp};
use crate::database::AsDynDatabase;
use crate::function::eviction::EvictionPolicy;
use crate::function::memo::Memo;
use crate::function::sync::ClaimResult;
Expand All @@ -19,12 +20,17 @@ where
zalsa_local: &'db ZalsaLocal,
id: Id,
) -> &'db C::Output<'db> {
zalsa.unwind_if_revision_cancelled(zalsa_local);
zalsa.unwind_if_revision_cancelled(db, zalsa_local);

let database_key_index = self.database_key_index(id);

#[cfg(debug_assertions)]
let _span = crate::tracing::debug_span!("fetch", query = ?database_key_index).entered();
let _span = crate::tracing::debug_span_with_db!(
db,
"fetch",
query = ?database_key_index
)
.entered();

let memo = self.refresh_memo(db, zalsa, zalsa_local, id);

Expand All @@ -35,6 +41,7 @@ where

let revisions = &memo.header.revisions;
zalsa_local.report_tracked_read(
db,
database_key_index,
revisions.durability,
revisions.changed_at,
Expand Down Expand Up @@ -62,11 +69,13 @@ where
// Keep the hot and cold probes in distinct control-flow blocks. Using `or_else`
// here can outline both into one function, making hot hits pay for the cold path's
// stack frame.
if let Some(memo) = self.fetch_hot(zalsa, id, memo_ingredient_index) {
if let Some(memo) = self.fetch_hot(db, zalsa, id, memo_ingredient_index) {
return memo;
}

if let Some(memo) = self.fetch_cold(zalsa, zalsa_local, db, id, memo_ingredient_index) {
if let Some(memo) = crate::attach::attach_if_needed(db, || {
self.fetch_cold(zalsa, zalsa_local, db, id, memo_ingredient_index)
}) {
return memo;
}
}
Expand All @@ -75,6 +84,7 @@ where
#[inline(always)]
fn fetch_hot<'db>(
&'db self,
db: &'db C::DbView,
zalsa: &'db Zalsa,
id: Id,
memo_ingredient_index: MemoIngredientIndex,
Expand All @@ -85,13 +95,17 @@ where

let database_key_index = self.database_key_index(id);

let can_shallow_update = memo
.header
.shallow_verify_memo(zalsa, database_key_index, true);
let can_shallow_update =
memo.header
.shallow_verify_memo(db.as_dyn_database(), zalsa, database_key_index, true);

if can_shallow_update.yes() && !memo.header.may_be_provisional() {
memo.header
.update_shallow(zalsa, database_key_index, can_shallow_update);
memo.header.update_shallow(
db.as_dyn_database(),
zalsa,
database_key_index,
can_shallow_update,
);

// SAFETY: memo is present in memo_map and we have verified that it is
// still valid for the current revision.
Expand Down Expand Up @@ -137,9 +151,13 @@ where

if let Some(old_memo) = opt_old_memo {
if old_memo.value.is_some()
&& old_memo
.header
.verify_memo(db.into(), &claim_guard, C::CYCLE_STRATEGY, true)
&& old_memo.header.verify_memo(
db.into(),
db.as_dyn_database(),
&claim_guard,
C::CYCLE_STRATEGY,
true,
)
{
// SAFETY: memo is present in memo_map and we have verified that it is
// still valid for the current revision.
Expand Down
Loading
Loading