Replace IntervalMap with a BBST-based implementation (#39326)

This replaces the naive IntervalMap implementation (which queries each
subscription against each index key separately) with a treap-based data
structure that supports querying an index key in O((k+1) log n) time, where k is the
number of _matching_ subscriptions and n is the total number of subscribed
intervals. The net effect is that long-lived subscriptions get much cheaper.

GitOrigin-RevId: f39ee382b2ca4441b71b64715c52d27a2556bcc5

Author: goffrie

Cargo.lock

@@ -59,6 +59,7 @@ encoding_rs = "0.8.32"
 enum-iterator = "2.1.0"
 event-listener = "5.3.1"
 fastrace = { git = "https://github.com/fast/fastrace", rev = "eacc377a8b3435e04f1d7a68085ce6eedb0d1d4a", version = "0.7", features = [ "enable" ] }
+fastrand = "2.3.0"
 flate2 = { version = "1", features = [ "zlib-ng" ] }
 flexbuffers = "25"
 float_next_after = "1.0.0"

Cargo.toml

@@ -60,12 +60,18 @@ impl End {
         }
     }
 
+    #[inline]
     pub const fn as_ref(&self) -> EndRef<'_> {
         match self {
             End::Excluded(binary_key) => EndRef::Excluded(binary_key.as_slice()),
             End::Unbounded => EndRef::Unbounded,
         }
     }
+
+    #[inline]
+    pub fn greater_than(&self, key: &[u8]) -> bool {
+        self.as_ref().greater_than(key)
+    }
 }
 
 impl HeapSize for End {
@@ -90,6 +96,14 @@ impl EndRef<'_> {
             EndRef::Unbounded => End::Unbounded,
         }
     }
+
+    #[inline]
+    pub fn greater_than(&self, key: &[u8]) -> bool {
+        match *self {
+            EndRef::Excluded(end) => key < end,
+            EndRef::Unbounded => true,
+        }
+    }
 }
 
 #[cfg(test)]

crates/common/src/interval/bounds.rs

@@ -50,6 +50,7 @@ futures-async-stream = { workspace = true }
 governor = { workspace = true }
 imbl = { workspace = true }
 indexing = { path = "../indexing" }
+interval_map = { path = "../interval_map" }
 itertools = { workspace = true }
 keybroker = { path = "../keybroker" }
 maplit = { workspace = true }

crates/database/Cargo.toml

@@ -51,7 +51,7 @@ use futures::{
     FutureExt as _,
     StreamExt as _,
 };
-use indexing::interval::IntervalMap;
+use interval_map::IntervalMap;
 use parking_lot::Mutex;
 use prometheus::VMHistogram;
 use search::query::TextSearchSubscriptions;
@@ -390,7 +390,7 @@ impl SubscriptionManager {
                         .subscriptions
                         .indexed
                         .iter()
-                        .map(|(key, (_fields, range_map))| (key, range_map.len()))
+                        .map(|(key, (_fields, range_map))| (key, range_map.subscriber_len()))
                         .collect();
                     let total_subscribers: usize = subscribers_by_index.values().sum();
                     let search_len = self.subscriptions.search.filter_len();
@@ -471,10 +471,7 @@ impl SubscriptionManager {
                     metrics::log_missing_index_key_subscriptions();
                     continue;
                 };
-
-                for subscriber_id in range_map.query(index_key) {
-                    notify(subscriber_id);
-                }
+                range_map.query(index_key, &mut *notify);
             }
         }
 
@@ -623,7 +620,8 @@ impl Subscription {
 
 /// Tracks every subscriber for a given read-set.
 struct SubscriptionMap {
-    indexed: BTreeMap<TabletIndexName, (IndexedFields, IntervalMap<SubscriberId>)>,
+    // TODO: remove nesting, merge all IntervalMaps into one big data structure
+    indexed: BTreeMap<TabletIndexName, (IndexedFields, IntervalMap)>,
     search: TextSearchSubscriptions,
 }
 
@@ -641,7 +639,9 @@ impl SubscriptionMap {
                 .indexed
                 .entry(index.clone())
                 .or_insert_with(|| (index_reads.fields.clone(), IntervalMap::new()));
-            interval_map.insert(id, index_reads.intervals.clone());
+            interval_map
+                .insert(id, index_reads.intervals.iter())
+                .expect("stored more than u32::MAX intervals?");
         }
         for (index, reads) in reads.iter_search() {
             self.search.insert(id, index, reads);
@@ -654,7 +654,7 @@ impl SubscriptionMap {
                 .indexed
                 .get_mut(index)
                 .unwrap_or_else(|| panic!("Missing index entry for {}", index));
-            assert!(range_map.remove(id).is_some());
+            range_map.remove(id);
             if range_map.is_empty() {
                 self.indexed.remove(index);
             }

crates/database/src/subscription.rs

@@ -386,6 +386,7 @@ impl EqFields {
 }
 
 impl From<EqFields> for Interval {
+    #[inline]
     fn from(value: EqFields) -> Self {
         Interval::prefix(value.prefix.into())
     }

crates/database/src/system_query.rs

@@ -8,7 +8,6 @@
 
 pub mod backend_in_memory_indexes;
 pub mod index_registry;
-pub mod interval;
 mod metrics;
 
 #[cfg(test)]

crates/indexing/src/interval.rs

@@ -0,0 +1,19 @@
+[package]
+name = "interval_map"
+version = "0.1.0"
+edition = "2021"
+license = "LicenseRef-FSL-1.1-Apache-2.0"
+
+[dependencies]
+common = { path = "../common" }
+fastrand = { workspace = true }
+slab = { workspace = true }
+
+[dev-dependencies]
+cmd_util = { path = "../cmd_util" }
+common = { path = "../common", features = ["testing"] }
+proptest = { workspace = true }
+proptest-derive = { workspace = true }
+
+[lints]
+workspace = true