Skip to content

Fix query overwrites in the subscription manager #2905

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
joshua-spacetime merged 1 commit into from
Jun 26, 2025
Merged

Fix query overwrites in the subscription manager #2905

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
124 changes: 112 additions & 12 deletions crates/core/src/subscription/module_subscription_actor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -554,7 +554,7 @@ impl ModuleSubscriptions {
fn compile_queries(
&self,
sender: Identity,
queries: impl IntoIterator<Item = Box<str>>,
queries: &[Box<str>],
num_queries: usize,
metrics: &SubscriptionMetrics,
) -> Result<(Vec<Arc<Plan>>, AuthCtx, TxId, HistogramTimer), DBError> {
Expand All @@ -563,12 +563,13 @@ impl ModuleSubscriptions {
let mut query_hashes = Vec::with_capacity(num_queries);

for sql in queries {
if is_subscribe_to_all_tables(&sql) {
let sql = sql.trim();
if is_subscribe_to_all_tables(sql) {
subscribe_to_all_tables = true;
continue;
}
let hash = QueryHash::from_string(&sql, sender, false);
let hash_with_param = QueryHash::from_string(&sql, sender, true);
let hash = QueryHash::from_string(sql, sender, false);
let hash_with_param = QueryHash::from_string(sql, sender, true);
query_hashes.push((sql, hash, hash_with_param));
}

Expand Down Expand Up @@ -606,10 +607,10 @@ impl ModuleSubscriptions {
plans.push(unit);
} else {
plans.push(Arc::new(
compile_query_with_hashes(&auth, &tx, &sql, hash, hash_with_param).map_err(|err| {
compile_query_with_hashes(&auth, &tx, sql, hash, hash_with_param).map_err(|err| {
DBError::WithSql {
error: Box::new(DBError::Other(err.into())),
sql,
sql: sql.into(),
}
})?,
));
Expand Down Expand Up @@ -670,7 +671,7 @@ impl ModuleSubscriptions {
let (queries, auth, tx, compile_timer) = return_on_err!(
self.compile_queries(
sender.id.identity,
request.query_strings,
&request.query_strings,
num_queries,
&subscription_metrics
),
Expand Down Expand Up @@ -767,7 +768,7 @@ impl ModuleSubscriptions {

let (queries, auth, tx, compile_timer) = self.compile_queries(
sender.id.identity,
subscription.query_strings,
&subscription.query_strings,
num_queries,
&subscription_metrics,
)?;
Expand Down Expand Up @@ -2331,48 +2332,147 @@ mod tests {
Ok(())
}

/// Test that one client unsubscribing does not affect another
/// Test that one client subscribing does not affect another
#[tokio::test]
async fn test_unsubscribe() -> anyhow::Result<()> {
async fn test_subscribe_distinct_queries_same_plan() -> anyhow::Result<()> {
// Establish a connection for each client
let (tx_for_a, mut rx_for_a) = client_connection(client_id_from_u8(1));
let (tx_for_b, mut rx_for_b) = client_connection(client_id_from_u8(2));

let db = relational_db()?;
let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());

let u_id = db.create_table_for_test(
let u_id = db.create_table_for_test_with_the_works(
"u",
&[
("i", AlgebraicType::U64),
("a", AlgebraicType::U64),
("b", AlgebraicType::U64),
],
&[0.into()],
// The join column for this table does not have to be unique,
// because pruning only requires us to probe the join index on `v`.
&[],
StAccess::Public,
)?;
let v_id = db.create_table_for_test(
let v_id = db.create_table_for_test_with_the_works(
"v",
&[
("i", AlgebraicType::U64),
("x", AlgebraicType::U64),
("y", AlgebraicType::U64),
],
&[0.into(), 1.into()],
&[0.into()],
StAccess::Public,
)?;

commit_tx(&db, &subs, [], [(v_id, product![1u64, 1u64, 1u64])])?;

let mut query_ids = 0;

// Both clients subscribe to the same query modulo whitespace
subscribe_multi(
&subs,
&["select u.* from u join v on u.i = v.i where v.x = 1"],
tx_for_a,
&mut query_ids,
)?;
subscribe_multi(
&subs,
&["select u.* from u join v on u.i = v.i where v.x = 1"],
tx_for_b.clone(),
&mut query_ids,
)?;

// Wait for both subscriptions
assert_matches!(
rx_for_a.recv().await,
Some(SerializableMessage::Subscription(SubscriptionMessage {
result: SubscriptionResult::SubscribeMulti(_),
..
}))
);
assert_matches!(
rx_for_b.recv().await,
Some(SerializableMessage::Subscription(SubscriptionMessage {
result: SubscriptionResult::SubscribeMulti(_),
..
}))
);

// Insert a new row into `u`
commit_tx(&db, &subs, [], [(u_id, product![1u64, 0u64, 0u64])])?;

assert_tx_update_for_table(
&mut rx_for_a,
u_id,
&ProductType::from([AlgebraicType::U64, AlgebraicType::U64, AlgebraicType::U64]),
[product![1u64, 0u64, 0u64]],
[],
)
.await;

assert_tx_update_for_table(
&mut rx_for_b,
u_id,
&ProductType::from([AlgebraicType::U64, AlgebraicType::U64, AlgebraicType::U64]),
[product![1u64, 0u64, 0u64]],
[],
)
.await;

Ok(())
}

/// Test that one client unsubscribing does not affect another
#[tokio::test]
async fn test_unsubscribe_distinct_queries_same_plan() -> anyhow::Result<()> {
// Establish a connection for each client
let (tx_for_a, mut rx_for_a) = client_connection(client_id_from_u8(1));
let (tx_for_b, mut rx_for_b) = client_connection(client_id_from_u8(2));

let db = relational_db()?;
let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());

let u_id = db.create_table_for_test_with_the_works(
"u",
&[
("i", AlgebraicType::U64),
("a", AlgebraicType::U64),
("b", AlgebraicType::U64),
],
&[0.into()],
// The join column for this table does not have to be unique,
// because pruning only requires us to probe the join index on `v`.
&[],
StAccess::Public,
)?;
let v_id = db.create_table_for_test_with_the_works(
"v",
&[
("i", AlgebraicType::U64),
("x", AlgebraicType::U64),
("y", AlgebraicType::U64),
],
&[0.into(), 1.into()],
&[0.into()],
StAccess::Public,
)?;

commit_tx(&db, &subs, [], [(v_id, product![1u64, 1u64, 1u64])])?;

let mut query_ids = 0;

subscribe_multi(
&subs,
&["select u.* from u join v on u.i = v.i where v.x = 1"],
tx_for_a,
&mut query_ids,
)?;
subscribe_multi(
&subs,
&["select u.* from u join v on u.i = v.i where v.x = 1"],
tx_for_b.clone(),
&mut query_ids,
)?;
Expand Down
23 changes: 17 additions & 6 deletions crates/core/src/subscription/module_subscription_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -399,7 +399,7 @@ impl QueriedTableIndexIds {
/// See [`JoinEdge`] for more details.
#[derive(Debug, Default)]
pub struct JoinEdges {
edges: BTreeMap<JoinEdge, HashMap<AlgebraicValue, QueryHash>>,
edges: BTreeMap<JoinEdge, HashMap<AlgebraicValue, HashSet<QueryHash>>>,
}

impl JoinEdges {
Expand All @@ -408,18 +408,28 @@ impl JoinEdges {
let mut inserted = false;
for (edge, rhs_val) in qs.query.join_edges() {
inserted = true;
self.edges.entry(edge).or_default().insert(rhs_val, qs.query.hash);
self.edges
.entry(edge)
.or_default()
.entry(rhs_val)
.or_default()
.insert(qs.query.hash);
}
inserted
}

/// If this query has any join edges, remove them from the map.
fn remove_query(&mut self, query: &Query) {
for (edge, rhs_val) in query.join_edges() {
if let Some(hashes) = self.edges.get_mut(&edge) {
hashes.remove(&rhs_val);
if hashes.is_empty() {
self.edges.remove(&edge);
if let Some(values) = self.edges.get_mut(&edge) {
if let Some(hashes) = values.get_mut(&rhs_val) {
hashes.remove(&query.hash);
if hashes.is_empty() {
values.remove(&rhs_val);
if values.is_empty() {
self.edges.remove(&edge);
}
}
}
}
}
Expand All @@ -436,6 +446,7 @@ impl JoinEdges {
self.edges
.range(JoinEdge::range_for_table(table_id))
.filter_map(move |(edge, hashes)| find_rhs_val(edge, row).as_ref().and_then(|rhs_val| hashes.get(rhs_val)))
.flatten()
}
}

Expand Down
Loading