fix(1115): Identify subscribers by Identity AND Address

Closes #1115.

Previously subscribers were only identified by their Identity.
However the same Identity can be associated to different Addresses.

Author: joshua-spacetime

crates/core/src/subscription/module_subscription_actor.rs

@@ -88,7 +88,7 @@ impl ModuleSubscriptions {
         // This also makes it possible for `broadcast_event` to get scheduled before the subsequent part here
         // but that should not pose an issue.
         let mut subscriptions = self.subscriptions.write();
-        subscriptions.remove_subscription(&sender.id.identity);
+        subscriptions.remove_subscription(&(sender.id.identity, sender.id.address));
         subscriptions.add_subscription(sender.clone(), execution_set.into_iter());
         let num_queries = subscriptions.num_queries();
 
@@ -118,7 +118,7 @@ impl ModuleSubscriptions {
 
     pub fn remove_subscriber(&self, client_id: ClientActorId) {
         let mut subscriptions = self.subscriptions.write();
-        subscriptions.remove_subscription(&client_id.identity);
+        subscriptions.remove_subscription(&(client_id.identity, client_id.address));
         WORKER_METRICS
             .subscription_queries
             .with_label_values(&self.relational_db.address())

crates/core/src/subscription/module_subscription_manager.rs

@@ -10,13 +10,14 @@ use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use smallvec::SmallVec;
 use spacetimedb_client_api_messages::client_api::{TableRowOperation, TableUpdate};
 use spacetimedb_data_structures::map::{Entry, HashMap, HashSet, IntMap};
-use spacetimedb_lib::Identity;
+use spacetimedb_lib::{Address, Identity};
 use spacetimedb_primitives::TableId;
 use std::ops::Deref;
 use std::sync::Arc;
 
 type Query = Arc<ExecutionUnit>;
 type Client = Arc<ClientConnectionSender>;
+type Id = (Identity, Address);
 
 /// Responsible for the efficient evaluation of subscriptions.
 /// It performs basic multi-query optimization,
@@ -26,17 +27,17 @@ type Client = Arc<ClientConnectionSender>;
 #[derive(Debug, Default)]
 pub struct SubscriptionManager {
     // Subscriber identities and their client connections.
-    clients: HashMap<Identity, Client>,
+    clients: HashMap<Id, Client>,
     // Queries for which there is at least one subscriber.
     queries: HashMap<QueryHash, Query>,
     // The subscribers for each query.
-    subscribers: HashMap<QueryHash, HashSet<Identity>>,
+    subscribers: HashMap<QueryHash, HashSet<Id>>,
     // Inverted index from tables to queries that read from them.
     tables: IntMap<TableId, HashSet<QueryHash>>,
 }
 
 impl SubscriptionManager {
-    pub fn client(&self, id: &Identity) -> Client {
+    pub fn client(&self, id: &Id) -> Client {
         self.clients[id].clone()
     }
 
@@ -54,7 +55,7 @@ impl SubscriptionManager {
     }
 
     #[cfg(test)]
-    fn contains_subscription(&self, subscriber: &Identity, query: &QueryHash) -> bool {
+    fn contains_subscription(&self, subscriber: &Id, query: &QueryHash) -> bool {
         self.subscribers.get(query).is_some_and(|ids| ids.contains(subscriber))
     }
 
@@ -68,7 +69,7 @@ impl SubscriptionManager {
     /// its table ids added to the inverted index.
     #[tracing::instrument(skip_all)]
     pub fn add_subscription(&mut self, client: Client, queries: impl IntoIterator<Item = Query>) {
-        let id = client.id.identity;
+        let id = (client.id.identity, client.id.address);
         self.clients.insert(id, client);
         for unit in queries {
             let hash = unit.hash();
@@ -83,7 +84,7 @@ impl SubscriptionManager {
     /// If a query no longer has any subscribers,
     /// it is removed from the index along with its table ids.
     #[tracing::instrument(skip_all)]
-    pub fn remove_subscription(&mut self, client: &Identity) {
+    pub fn remove_subscription(&mut self, client: &Id) {
         // Remove `hash` from the set of queries for `table_id`.
         // When the table has no queries, cleanup the map entry altogether.
         let mut remove_table_query = |table_id: TableId, hash: &QueryHash| {
@@ -188,7 +189,7 @@ impl SubscriptionManager {
                 // so we'll have either `TableUpdate` (`Protocol::Binary`)
                 // or `TableUpdateJson` (`Protocol::Text`).
                 .fold(
-                    HashMap::<(&Identity, TableId), Either<TableUpdate, TableUpdateJson>>::new(),
+                    HashMap::<(&Id, TableId), Either<TableUpdate, TableUpdateJson>>::new(),
                     |mut tables, (id, table_id, table_name, ops)| {
                         match tables.entry((id, table_id)) {
                             Entry::Occupied(mut entry) => match ops {
@@ -220,7 +221,7 @@ impl SubscriptionManager {
                 // So before sending the updates to each client,
                 // we must stitch together the `TableUpdate*`s into an aggregated list.
                 .fold(
-                    HashMap::<&Identity, Either<Vec<TableUpdate>, Vec<TableUpdateJson>>>::new(),
+                    HashMap::<&Id, Either<Vec<TableUpdate>, Vec<TableUpdateJson>>>::new(),
                     |mut updates, ((id, _), update)| {
                         let entry = updates.entry(id);
                         match update {
@@ -265,12 +266,12 @@ impl SubscriptionManager {
 mod tests {
     use std::sync::Arc;
 
-    use spacetimedb_lib::{error::ResultTest, AlgebraicType, Identity};
+    use spacetimedb_lib::{error::ResultTest, Address, AlgebraicType, Identity};
     use spacetimedb_primitives::TableId;
     use spacetimedb_vm::expr::CrudExpr;
 
     use crate::{
-        client::{ClientActorId, ClientConnectionSender, Protocol},
+        client::{ClientActorId, ClientConnectionSender, ClientName, Protocol},
         db::relational_db::{tests_utils::TestDB, RelationalDB},
         execution_context::ExecutionContext,
         sql::compiler::compile_sql,
@@ -300,6 +301,21 @@ mod tests {
         })
     }
 
+    fn id(address: u128) -> (Identity, Address) {
+        (Identity::ZERO, Address::from_u128(address))
+    }
+
+    fn client(address: u128) -> ClientConnectionSender {
+        ClientConnectionSender::dummy(
+            ClientActorId {
+                identity: Identity::ZERO,
+                address: Address::from_u128(address),
+                name: ClientName(0),
+            },
+            Protocol::Binary,
+        )
+    }
+
     #[test]
     fn test_subscribe() -> ResultTest<()> {
         let db = TestDB::durable()?;
@@ -309,9 +325,8 @@ mod tests {
         let plan = compile_plan(&db, sql)?;
         let hash = plan.hash();
 
-        let id = Identity::ZERO;
-        let client = ClientActorId::for_test(id);
-        let client = Arc::new(ClientConnectionSender::dummy(client, Protocol::Binary));
+        let id = id(0);
+        let client = Arc::new(client(0));
 
         let mut subscriptions = SubscriptionManager::default();
         subscriptions.add_subscription(client, [plan]);
@@ -332,9 +347,8 @@ mod tests {
         let plan = compile_plan(&db, sql)?;
         let hash = plan.hash();
 
-        let id = Identity::ZERO;
-        let client = ClientActorId::for_test(id);
-        let client = Arc::new(ClientConnectionSender::dummy(client, Protocol::Binary));
+        let id = id(0);
+        let client = Arc::new(client(0));
 
         let mut subscriptions = SubscriptionManager::default();
         subscriptions.add_subscription(client, [plan]);
@@ -356,9 +370,8 @@ mod tests {
         let plan = compile_plan(&db, sql)?;
         let hash = plan.hash();
 
-        let id = Identity::ZERO;
-        let client = ClientActorId::for_test(id);
-        let client = Arc::new(ClientConnectionSender::dummy(client, Protocol::Binary));
+        let id = id(0);
+        let client = Arc::new(client(0));
 
         let mut subscriptions = SubscriptionManager::default();
         subscriptions.add_subscription(client.clone(), [plan.clone()]);
@@ -386,13 +399,11 @@ mod tests {
         let plan = compile_plan(&db, sql)?;
         let hash = plan.hash();
 
-        let id0 = Identity::ZERO;
-        let client0 = ClientActorId::for_test(id0);
-        let client0 = Arc::new(ClientConnectionSender::dummy(client0, Protocol::Binary));
+        let id0 = id(0);
+        let client0 = Arc::new(client(0));
 
-        let id1 = Identity::from_byte_array([1; 32]);
-        let client1 = ClientActorId::for_test(id1);
-        let client1 = Arc::new(ClientConnectionSender::dummy(client1, Protocol::Binary));
+        let id1 = id(1);
+        let client1 = Arc::new(client(1));
 
         let mut subscriptions = SubscriptionManager::default();
         subscriptions.add_subscription(client0, [plan.clone()]);
@@ -433,13 +444,11 @@ mod tests {
         let hash_select0 = plan_select0.hash();
         let hash_select1 = plan_select1.hash();
 
-        let id0 = Identity::ZERO;
-        let client0 = ClientActorId::for_test(id0);
-        let client0 = Arc::new(ClientConnectionSender::dummy(client0, Protocol::Binary));
+        let id0 = id(0);
+        let client0 = Arc::new(client(0));
 
-        let id1 = Identity::from_byte_array([1; 32]);
-        let client1 = ClientActorId::for_test(id1);
-        let client1 = Arc::new(ClientConnectionSender::dummy(client1, Protocol::Binary));
+        let id1 = id(1);
+        let client1 = Arc::new(client(1));
 
         let mut subscriptions = SubscriptionManager::default();
         subscriptions.add_subscription(client0, [plan_scan.clone(), plan_select0.clone()]);