using a callback id indirection

Author: fulmicoton

chitchat/src/listener.rs

@@ -3,10 +3,12 @@ use std::ops::Bound;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, RwLock, Weak};
 
+use tracing::error;
+
 use crate::KeyChangeEvent;
 
 pub struct ListenerHandle {
-    prefix: String,
+    // The prefix and listener_id are used for removal of the listener.
     listener_id: usize,
     listeners: Weak<RwLock<InnerListeners>>,
 }
@@ -25,7 +27,7 @@ impl Drop for ListenerHandle {
     fn drop(&mut self) {
         if let Some(listeners) = self.listeners.upgrade() {
             let mut listeners_guard = listeners.write().unwrap();
-            listeners_guard.remove_listener(&self.prefix, self.listener_id);
+            listeners_guard.remove_listener(self.listener_id);
         }
     }
 }
@@ -60,9 +62,15 @@ impl Listeners {
         let new_idx = inner_listener_guard
             .listener_idx
             .fetch_add(1, Ordering::Relaxed);
-        inner_listener_guard.subscribe_event(&key_prefix, new_idx, boxed_listener);
+        let callback_entry = CallbackEntry {
+            prefix: key_prefix.clone(),
+            callback: boxed_listener,
+        };
+        inner_listener_guard
+            .callbacks
+            .insert(new_idx, callback_entry);
+        inner_listener_guard.subscribe_event(&key_prefix, new_idx);
         ListenerHandle {
-            prefix: key_prefix,
             listener_id: new_idx,
             listeners: weak_listeners,
         }
@@ -73,33 +81,50 @@ impl Listeners {
     }
 }
 
+struct CallbackEntry {
+    prefix: String,
+    callback: BoxedListener,
+}
+
+type CallbackId = usize;
+
 #[derive(Default)]
 struct InnerListeners {
     // A trie would have been more efficient, but in reality we don't have
     // that many listeners.
-    listeners: BTreeMap<String, HashMap<usize, BoxedListener>>,
+    listeners: BTreeMap<String, Vec<CallbackId>>,
     listener_idx: AtomicUsize,
+    // Callbacks is a hashmap because as we delete listeners, we create "holes" in the
+    // callback_id -> callback mapping
+    callbacks: HashMap<usize, CallbackEntry>,
 }
 
 impl InnerListeners {
     // We don't inline this to make sure monomorphization generates as little code as possible.
-    fn subscribe_event(&mut self, key_prefix: &str, idx: usize, callback: BoxedListener) {
-        if let Some(callbacks) = self.listeners.get_mut(key_prefix) {
-            callbacks.insert(idx, callback);
-        } else {
-            let mut listener_map = HashMap::new();
-            listener_map.insert(idx, callback);
-            self.listeners.insert(key_prefix.to_string(), listener_map);
-        }
+    fn subscribe_event(&mut self, key_prefix: &str, idx: CallbackId) {
+        self.listeners
+            .entry(key_prefix.to_string())
+            .or_default()
+            .push(idx);
+    }
+
+    fn call(&self, callback_id: CallbackId, key_change_event: KeyChangeEvent) {
+        let Some(CallbackEntry { callback, .. }) = self.callbacks.get(&callback_id) else {
+            error!(
+                "callback {callback_id} not found upon call. this should not happen, please report"
+            );
+            return;
+        };
+        (*callback)(key_change_event);
     }
 
     fn trigger_event(&self, key_change_event: KeyChangeEvent) {
         // We treat the empty prefix a tiny bit separately to get able to at least
         // use the first character as a range bound, as if we were going to the first level of
         // a trie.
-        if let Some(listeners) = self.listeners.get("") {
-            for listener in listeners.values() {
-                (*listener)(key_change_event);
+        if let Some(callback_ids) = self.listeners.get("") {
+            for &callback_id in callback_ids {
+                self.call(callback_id, key_change_event);
             }
         }
         if key_change_event.key.is_empty() {
@@ -115,17 +140,32 @@ impl InnerListeners {
                 break;
             }
             if let Some(stripped_key_change_event) = key_change_event.strip_key_prefix(prefix_key) {
-                for listener in listeners.values() {
-                    (*listener)(stripped_key_change_event);
+                for &callback_id in listeners {
+                    self.call(callback_id, stripped_key_change_event)
                 }
             }
         }
     }
 
-    fn remove_listener(&mut self, key_prefix: &str, idx: usize) {
-        if let Some(callbacks) = self.listeners.get_mut(key_prefix) {
-            callbacks.remove(&idx);
-        }
+    fn remove_listener(&mut self, callback_id: CallbackId) {
+        let Some(CallbackEntry { prefix, .. }) = self.callbacks.remove(&callback_id) else {
+            error!(
+                "callback {callback_id} not found upon remove. this should not happen, please \
+                 report"
+            );
+            return;
+        };
+        let Some(callbacks) = self.listeners.get_mut(&prefix) else {
+            error!(
+                "callback prefix not foudn upon remove. this should never happen, please report"
+            );
+            return;
+        };
+        let position = callbacks
+            .iter()
+            .position(|x| *x == callback_id)
+            .expect("callback not found");
+        callbacks.swap_remove(position);
     }
 }
 

chitchat/src/state.rs

@@ -17,7 +17,7 @@ use tracing::{info, warn};
 use crate::delta::{Delta, DeltaSerializer, NodeDelta};
 use crate::digest::{Digest, NodeDigest};
 use crate::listener::Listeners;
-use crate::types::{DeletionStatus, DeletionStatusMutation};
+use crate::types::{DeletionStatus, DeletionStatusMutation, KeyValueMutation};
 use crate::{ChitchatId, Heartbeat, KeyChangeEvent, Version, VersionedValue};
 
 #[derive(Clone, Serialize, Deserialize)]
@@ -110,6 +110,7 @@ impl NodeState {
     }
 
     /// Returns the node's max version.
+    #[inline]
     pub fn max_version(&self) -> Version {
         self.max_version
     }
@@ -210,7 +211,11 @@ impl NodeState {
             current_last_gc_version=self.last_gc_version,
             "resetting node");
         *self = NodeState::new(node_delta.chitchat_id.clone(), self.listeners.clone());
-        // The node_delta max_version  whe
+        // It is possible for the node delta to not contain any KVs.
+        // (for instance they all have been GCed.)
+        //
+        // In that case, no KV are here to tell us what the max version is, so the
+        // node_delta itself holds a max_version.
         if let Some(max_version) = node_delta.max_version {
             if node_delta.key_values.is_empty() {
                 self.max_version = max_version;
@@ -226,6 +231,28 @@ impl NodeState {
         true
     }
 
+    fn applicable_key_value_mutations(
+        &self,
+        node_delta: NodeDelta,
+    ) -> impl Iterator<Item = KeyValueMutation> + '_ {
+        node_delta
+            .key_values
+            .into_iter()
+            .filter(move |key_value_mutation| {
+                if key_value_mutation.version <= self.max_version() {
+                    // We already know about this KV.
+                    return false;
+                }
+                if key_value_mutation.status.scheduled_for_deletion() {
+                    // This KV has already been GCed.
+                    if key_value_mutation.version <= self.last_gc_version {
+                        return false;
+                    }
+                }
+                true
+            })
+    }
+
     fn apply_delta(&mut self, node_delta: NodeDelta, now: Instant) {
         if !self.prepare_apply_delta(&node_delta) {
             return;