From 4f185e5f4b3d936beb8985dbd1d93cb64191b140 Mon Sep 17 00:00:00 2001 From: Rana718 Date: Thu, 6 Aug 2026 19:55:59 +0530 Subject: [PATCH 1/4] fix the pubsub benchmark issues --- bench/pubsub.go | 20 +++++++++++++++++--- bench/redis-cluster/docker-compose.yml | 13 +------------ 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/bench/pubsub.go b/bench/pubsub.go index 0c86ab2..cce7328 100644 --- a/bench/pubsub.go +++ b/bench/pubsub.go @@ -48,9 +48,11 @@ func runPubSub(addr string) { subConns[i] = conn conn.Write(subCmd) + conn.SetReadDeadline(time.Now().Add(5 * time.Second)) r := bufio.NewReaderSize(conn, 512) consumeSubConfirm(r) + conn.SetReadDeadline(time.Time{}) if r.Buffered() > 0 { tmp := make([]byte, r.Buffered()) r.Read(tmp) @@ -109,6 +111,7 @@ func runPubSub(addr string) { defer pubWg.Done() defer conn.Close() + conn.SetDeadline(time.Now().Add(30 * time.Second)) w := bufio.NewWriterSize(conn, 256<<10) r := bufio.NewReaderSize(conn, 128<<10) @@ -121,17 +124,28 @@ func runPubSub(addr string) { for j := 0; j < batch; j++ { w.Write(pubCmd) } - w.Flush() + if err := w.Flush(); err != nil { + atomic.AddInt64(&published, int64(sent)) + return + } skipLines(r, batch) sent += batch } - atomic.AddInt64(&published, int64(PUB_MSGS_EACH)) + atomic.AddInt64(&published, int64(sent)) }(pubConns[i]) } pubWg.Wait() pubElapsed := time.Since(pubStart) - recvDone.Wait() + done := make(chan struct{}) + go func() { + recvDone.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(30 * time.Second): + } recvElapsed := time.Since(pubStart) for _, c := range subConns { diff --git a/bench/redis-cluster/docker-compose.yml b/bench/redis-cluster/docker-compose.yml index 9cf23d7..124a0ca 100644 --- a/bench/redis-cluster/docker-compose.yml +++ b/bench/redis-cluster/docker-compose.yml @@ -216,18 +216,7 @@ services: networks: - redis-cluster-net entrypoint: > - sh -c " - sleep 3 && - redis-cli --cluster create - 172.20.0.11:7001 - 172.20.0.12:7002 - 172.20.0.13:7003 - 172.20.0.14:7004 - 172.20.0.15:7005 - 172.20.0.16:7006 - --cluster-replicas 0 - --cluster-yes - " + sh -c "sleep 3 && redis-cli --cluster create 172.20.0.11:7001 172.20.0.12:7002 172.20.0.13:7003 172.20.0.14:7004 172.20.0.15:7005 172.20.0.16:7006 --cluster-replicas 0 --cluster-yes" restart: "no" networks: From 703cf6abce66024a2367c0cddcf3d91f059f505d Mon Sep 17 00:00:00 2001 From: Rana718 Date: Fri, 14 Aug 2026 02:19:49 +0530 Subject: [PATCH 2/4] fix the old table never clam issues --- .gitignore | 3 +- ARCHITECTURE.md | 6 +- bench/kv.go | 62 ++++++++------ crates/customhash/README.md | 23 ++--- crates/customhash/src/ebr.rs | 59 +++++++++---- crates/customhash/src/lib.rs | 162 ++++++++++++++++++++++++++++++----- 6 files changed, 233 insertions(+), 82 deletions(-) diff --git a/.gitignore b/.gitignore index 5c3f932..75309ef 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ /target *.json *.rdb -temp/ \ No newline at end of file +temp/ +**/target/ diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 126bed2..e9c4409 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -155,8 +155,8 @@ When shard reaches 70% occupancy: 4. Copy all live Entry pointers from old table to new table (Entry objects are shared — same heap allocation, just referenced from new position) 5. shard.table.store(new_ptr, Release) — readers instantly see new table - 6. Old SlotTable (just the pointer array) is leaked - (safe: readers may still be probing it; entries are alive in new table) + 6. Retire the old SlotTable through EBR + (freed after readers that may still be probing it leave their epoch) 7. grow_lock.unlock() After grow: @@ -165,7 +165,7 @@ After grow: - Writers retrying: see new threshold, insert into new table Memory lifecycle: - - SlotTable arrays: leaked on grow (8 bytes × old_capacity, ~few KB each) + - SlotTable arrays: reclaimed through EBR after each grow grace period - Entry objects: live forever once inserted (key stays for probing) - ValueBox: swapped atomically, recycled via EBR pool - String data inside values: freed when ValueBox is reclaimed diff --git a/bench/kv.go b/bench/kv.go index bc5cf30..4cb18b7 100644 --- a/bench/kv.go +++ b/bench/kv.go @@ -37,10 +37,10 @@ func runKV() { go func(id int, conn *net.TCPConn) { defer wg.Done() - w := bufio.NewWriterSize(conn, 256<<10) r := bufio.NewReaderSize(conn, 128<<10) var kb [32]byte + requests := make([]byte, 0, seqBatch*40) base := id * OPS_CLIENT sent := 0 for sent < OPS_CLIENT { @@ -48,11 +48,12 @@ func runKV() { if OPS_CLIENT-sent < batch { batch = OPS_CLIENT - sent } + requests = requests[:0] for j := 0; j < batch; j++ { kn := strconv.AppendInt(kb[:0], int64(base+sent+j), 10) - writeSetBytes(w, kn) + requests = appendSetBytes(requests, kn) } - w.Flush() + writeFull(conn, requests) discardN(r, batch*5) sent += batch } @@ -69,10 +70,10 @@ func runKV() { go func(id int, conn *net.TCPConn) { defer wg.Done() - w := bufio.NewWriterSize(conn, 256<<10) r := bufio.NewReaderSize(conn, 128<<10) var kb [32]byte + requests := make([]byte, 0, PIPE_SIZE*40) base := id * OPS_CLIENT sent := 0 for sent < OPS_CLIENT { @@ -80,11 +81,12 @@ func runKV() { if OPS_CLIENT-sent < batch { batch = OPS_CLIENT - sent } + requests = requests[:0] for j := 0; j < batch; j++ { kn := strconv.AppendInt(kb[:0], int64(base+sent+j), 10) - writeSetBytes(w, kn) + requests = appendSetBytes(requests, kn) } - w.Flush() + writeFull(conn, requests) discardN(r, batch*5) sent += batch } @@ -101,10 +103,10 @@ func runKV() { go func(id int, conn *net.TCPConn) { defer wg.Done() - w := bufio.NewWriterSize(conn, 256<<10) r := bufio.NewReaderSize(conn, 256<<10) var kb [32]byte + requests := make([]byte, 0, PIPE_SIZE*32) base := id * OPS_CLIENT sent := 0 for sent < OPS_CLIENT { @@ -112,11 +114,12 @@ func runKV() { if OPS_CLIENT-sent < batch { batch = OPS_CLIENT - sent } + requests = requests[:0] for j := 0; j < batch; j++ { kn := strconv.AppendInt(kb[:0], int64(base+sent+j), 10) - writeGetBytes(w, kn) + requests = appendGetBytes(requests, kn) } - w.Flush() + writeFull(conn, requests) skipGetReplies(r, batch) sent += batch } @@ -178,26 +181,25 @@ var ( crlfB = []byte("\r\n") ) -func writeSetBytes(w *bufio.Writer, key []byte) { - w.Write(setHdr) - writeLen(w, len(key)) - w.Write(crlfB) - w.Write(key) - w.Write(valPart) +func appendSetBytes(out, key []byte) []byte { + out = append(out, setHdr...) + out = appendLen(out, len(key)) + out = append(out, crlfB...) + out = append(out, key...) + return append(out, valPart...) } -func writeGetBytes(w *bufio.Writer, key []byte) { - w.Write(getHdr) - writeLen(w, len(key)) - w.Write(crlfB) - w.Write(key) - w.Write(crlfB) +func appendGetBytes(out, key []byte) []byte { + out = append(out, getHdr...) + out = appendLen(out, len(key)) + out = append(out, crlfB...) + out = append(out, key...) + return append(out, crlfB...) } -func writeLen(w *bufio.Writer, n int) { +func appendLen(out []byte, n int) []byte { if n < 10 { - w.WriteByte(byte('0' + n)) - return + return append(out, byte('0'+n)) } var buf [5]byte pos := len(buf) @@ -206,7 +208,17 @@ func writeLen(w *bufio.Writer, n int) { buf[pos] = byte('0' + n%10) n /= 10 } - w.Write(buf[pos:]) + return append(out, buf[pos:]...) +} + +func writeFull(conn *net.TCPConn, p []byte) { + for len(p) != 0 { + n, err := conn.Write(p) + if err != nil { + panic(err) + } + p = p[n:] + } } func discardN(r *bufio.Reader, n int) { diff --git a/crates/customhash/README.md b/crates/customhash/README.md index eed7fd1..3d0226e 100644 --- a/crates/customhash/README.md +++ b/crates/customhash/README.md @@ -1,9 +1,10 @@ # customhash -`customhash` is a fixed-capacity, sharded concurrent hash map specialized for -low-latency string workloads. Lookups are wait-free under the fixed-capacity -invariant; inserts and value replacement are lock-free. Values are published by -atomic pointer replacement and reclaimed with epoch-based reclamation. +`customhash` is a growable, sharded concurrent hash map specialized for +low-latency string workloads. Lookups are wait-free; value replacement is +lock-free, while new-key insertion briefly coordinates with the rare resize +path. Values and replaced slot tables are reclaimed with epoch-based +reclamation. The optimized API is: @@ -17,17 +18,17 @@ let reader = map.read(); assert_eq!(reader.get_prepared("key", key), Some("value")); ``` -The map has fixed capacity and does not physically reuse slots for unrelated -deleted keys. It is designed for pre-sized, high-read string maps. +The map does not physically reuse slots for unrelated deleted keys. It is +designed for pre-sized, high-read string maps, but grows when needed. ## Guarantees and limitations -`get`, `contains_key`, and prepared reads are wait-free while the map remains -within its configured capacity. Inserts, updates, and removes are lock-free, -but may retry under contention. `remove` is a logical deletion: the key slot +`get` and `contains_key` are wait-free. Updates and removes are lock-free and +may retry under contention. New-key inserts coordinate with table growth so +migration cannot miss a concurrently published entry. `remove` is a logical deletion: the key slot is retained so readers never observe freed entry memory, and the slot is -reusable only by reinserting the same key. Size the map for the maximum key -population and handle `try_insert` returning `Full` when capacity is reached. +reusable only by reinserting the same key. Slot tables grow automatically and +old tables are freed after all readers that could reference them have unpinned. The crate deliberately targets `String -> String`. Redis values such as hashes, expiry metadata, and pub/sub subscriber lists need mutable, heterogeneous diff --git a/crates/customhash/src/ebr.rs b/crates/customhash/src/ebr.rs index 5e4188f..eacf3dd 100644 --- a/crates/customhash/src/ebr.rs +++ b/crates/customhash/src/ebr.rs @@ -29,6 +29,7 @@ struct Garbage { drop_fn: unsafe fn(*mut u8), type_id: TypeId, epoch: u64, + recyclable: bool, } unsafe impl Send for Garbage {} @@ -48,6 +49,7 @@ struct Local { depth: usize, retires: usize, pool: Vec<(*mut u8, unsafe fn(*mut u8), TypeId)>, + collect_on_unpin: bool, initialized: bool, } @@ -59,6 +61,7 @@ impl Local { depth: 0, retires: 0, pool: Vec::new(), + collect_on_unpin: false, initialized: false, } } @@ -117,6 +120,13 @@ impl Local { unsafe { &*self.participant } .local .store(INACTIVE, Ordering::Release); + if self.collect_on_unpin { + // A retired table is usually rare, so drive its grace period + // promptly instead of waiting for 512 unrelated value retires. + self.collect(); + self.collect(); + self.collect_on_unpin = self.garbage.iter().any(|g| !g.recyclable); + } } } @@ -155,7 +165,7 @@ impl Local { let reclaimable = self.garbage.partition_point(|g| g.epoch + 2 <= safe); if reclaimable != 0 { for item in self.garbage.drain(..reclaimable) { - if self.pool.len() < VALUE_POOL_LIMIT { + if item.recyclable && self.pool.len() < VALUE_POOL_LIMIT { self.pool.push((item.ptr, item.drop_fn, item.type_id)); } else { unsafe { (item.drop_fn)(item.ptr) }; @@ -165,13 +175,20 @@ impl Local { } #[inline(always)] - fn retire_raw(&mut self, ptr: *mut u8, drop_fn: unsafe fn(*mut u8), type_id: TypeId) { + fn retire_raw( + &mut self, + ptr: *mut u8, + drop_fn: unsafe fn(*mut u8), + type_id: TypeId, + recyclable: bool, + ) { let epoch = GLOBAL_EPOCH.load(Ordering::Relaxed); self.garbage.push(Garbage { ptr, drop_fn, type_id, epoch, + recyclable, }); self.retires += 1; if self.retires % COLLECT_INTERVAL == 0 { @@ -203,7 +220,7 @@ impl Local { if old.is_null() { return true; } - self.retire_raw(old as *mut u8, drop_value_box::, TypeId::of::()); + self.retire_raw(old as *mut u8, drop_value_box::, TypeId::of::(), true); false } } @@ -237,29 +254,33 @@ pub unsafe fn retire_value(ptr: *mut ValueBox) { LOCAL.with(|c| { let l = unsafe { &mut *c.get() }; l.ensure_init(); - l.retire_raw(ptr as *mut u8, drop_value_box::, TypeId::of::()); + l.retire_raw(ptr as *mut u8, drop_value_box::, TypeId::of::(), true); }); } -#[inline(always)] -pub fn replace_value(slot: &AtomicPtr>, value: V) -> bool { - LOCAL.with(|c| unsafe { &mut *c.get() }.replace(slot, value)) +unsafe fn drop_box(ptr: *mut u8) { + unsafe { drop(Box::from_raw(ptr as *mut T)) }; } -#[inline(always)] -pub fn read_clone(slot: &AtomicPtr>) -> Option { +/// Retire a non-value allocation after every reader from the current epoch +/// has left its critical section. Unlike value boxes, these allocations are +/// returned to the allocator instead of entering the value reuse pool. +#[inline] +pub unsafe fn retire_box(ptr: *mut T) { + if ptr.is_null() { + return; + } LOCAL.with(|c| { let l = unsafe { &mut *c.get() }; - l.pin(); - let ptr = slot.load(Ordering::Acquire); - let r = if ptr.is_null() { - None - } else { - Some(unsafe { (*ptr).0.clone() }) - }; - l.unpin(); - r - }) + l.ensure_init(); + l.retire_raw(ptr as *mut u8, drop_box::, TypeId::of::(), false); + l.collect_on_unpin = true; + }); +} + +#[inline(always)] +pub fn replace_value(slot: &AtomicPtr>, value: V) -> bool { + LOCAL.with(|c| unsafe { &mut *c.get() }.replace(slot, value)) } #[inline(always)] diff --git a/crates/customhash/src/lib.rs b/crates/customhash/src/lib.rs index 0886b39..9b5116a 100644 --- a/crates/customhash/src/lib.rs +++ b/crates/customhash/src/lib.rs @@ -31,6 +31,7 @@ const LOAD_NUM: usize = 7; const LOAD_DEN: usize = 10; const DEFAULT_SHARD_CAPACITY: usize = 32_768; const INITIAL_SHARD_CAPACITY: usize = 1024; +const GROWING: usize = 1usize << (usize::BITS - 1); struct SlotTable { slots: Box<[AtomicPtr>]>, @@ -61,15 +62,28 @@ impl SlotTable { struct Shard { table: AtomicPtr>, len: CachePadded, + insert_gate: CachePadded, grow_lock: std::sync::Mutex<()>, } +struct InsertGuard<'a> { + gate: &'a AtomicUsize, +} + +impl Drop for InsertGuard<'_> { + #[inline(always)] + fn drop(&mut self) { + self.gate.fetch_sub(1, Ordering::Release); + } +} + impl Shard { fn new(cap: usize) -> Self { let table = Box::into_raw(Box::new(SlotTable::new(cap))); Shard { table: AtomicPtr::new(table), len: CachePadded::new(AtomicUsize::new(0)), + insert_gate: CachePadded::new(AtomicUsize::new(0)), grow_lock: std::sync::Mutex::new(()), } } @@ -79,6 +93,26 @@ impl Shard { unsafe { &*self.table.load(Ordering::Acquire) } } + #[inline(always)] + fn enter_insert(&self) -> InsertGuard<'_> { + loop { + let state = self.insert_gate.load(Ordering::Relaxed); + if state & GROWING != 0 { + std::hint::spin_loop(); + continue; + } + if self + .insert_gate + .compare_exchange_weak(state, state + 1, Ordering::Acquire, Ordering::Relaxed) + .is_ok() + { + return InsertGuard { + gate: &self.insert_gate, + }; + } + } + } + #[inline(always)] fn find(&self, key: &str, hash: u64) -> Option<&Entry> { let t = self.table(); @@ -98,6 +132,7 @@ impl Shard { #[inline(always)] fn insert_hashed(&self, key: String, value: V, hash: u64) -> bool { + let _guard = ebr::pin::(); if let Some(existing) = self.find(&key, hash) { let is_new = ebr::replace_value(&existing.value, value); return is_new; @@ -106,20 +141,24 @@ impl Shard { } fn insert_new(&self, key: String, value: V, hash: u64) -> bool { + let vb = new_value(value); + let entry: *mut Entry = Box::into_raw(Box::new(Entry { + hash, + value: AtomicPtr::new(vb), + key, + })); + let key_ref: &str = unsafe { &(*entry).key }; + loop { + // This is a shared atomic gate, so normal inserts remain fully + // concurrent. Growth exclusively closes it only during migration. + let insert_guard = self.enter_insert(); let t = self.table(); if self.len.load(Ordering::Relaxed) >= t.threshold { + drop(insert_guard); self.grow(); continue; } - - let vb = new_value(value.clone()); - let entry: *mut Entry = Box::into_raw(Box::new(Entry { - hash, - value: AtomicPtr::new(vb), - key: key.clone(), - })); - let key_ref: &str = unsafe { &(*entry).key }; let mut reserved = false; let t = self.table(); @@ -132,10 +171,7 @@ impl Shard { if !reserved { let cur = self.len.load(Ordering::Relaxed); if cur >= t.threshold { - unsafe { - free_value(vb); - let _ = Box::from_raw(entry); - } + drop(insert_guard); self.grow(); break; } @@ -145,7 +181,9 @@ impl Shard { Ordering::Relaxed, Ordering::Relaxed, ) { - Ok(_) => { reserved = true; } + Ok(_) => { + reserved = true; + } Err(_) => { std::hint::spin_loop(); continue; @@ -154,7 +192,12 @@ impl Shard { } if slot - .compare_exchange(ptr::null_mut(), entry, Ordering::Release, Ordering::Acquire) + .compare_exchange( + ptr::null_mut(), + entry, + Ordering::Release, + Ordering::Acquire, + ) .is_ok() { return true; @@ -168,7 +211,10 @@ impl Shard { if reserved { self.len.fetch_sub(1, Ordering::Relaxed); } - unsafe { let _ = Box::from_raw(entry); } + unsafe { + free_value(vb); + let _ = Box::from_raw(entry); + } if !old.is_null() { unsafe { ebr::retire_value(old) }; } @@ -183,6 +229,20 @@ impl Shard { fn grow(&self) { let _lock = self.grow_lock.lock().unwrap_or_else(|e| e.into_inner()); + // Close the gate only when no insert is publishing into the current + // table. Existing readers remain unaffected and are handled by EBR. + while self + .insert_gate + .compare_exchange_weak(0, GROWING, Ordering::AcqRel, Ordering::Relaxed) + .is_err() + { + std::hint::spin_loop(); + } + self.grow_locked(); + self.insert_gate.store(0, Ordering::Release); + } + + fn grow_locked(&self) { let old_ptr = self.table.load(Ordering::Acquire); let old_table = unsafe { &*old_ptr }; let cur_len = self.len.load(Ordering::Relaxed); @@ -217,6 +277,9 @@ impl Shard { let new_ptr = Box::into_raw(new_table); self.table.store(new_ptr, Ordering::Release); + // Entries are shared by both tables; only retire the old slot array. + // EBR releases it once all readers that could have loaded it unpin. + unsafe { ebr::retire_box(old_ptr) }; } } @@ -235,7 +298,9 @@ impl Drop for Shard { } } } - unsafe { drop(Box::from_raw(t_ptr)); } + unsafe { + drop(Box::from_raw(t_ptr)); + } } } } @@ -320,8 +385,8 @@ impl CustomMap { idx: usize, f: impl FnOnce(&V) -> R, ) -> Option { - let entry = unsafe { self.shards.get_unchecked(idx) }.find(key, hash)?; ebr::with_pin(|_| { + let entry = unsafe { self.shards.get_unchecked(idx) }.find(key, hash)?; let ptr = entry.value.load(Ordering::Acquire); if ptr.is_null() { return None; @@ -333,16 +398,21 @@ impl CustomMap { #[inline] pub fn contains_key(&self, key: &str) -> bool { let (h, idx) = self.locate(key); - unsafe { self.shards.get_unchecked(idx) } - .find(key, h) - .is_some_and(|e| !e.value.load(Ordering::Acquire).is_null()) + ebr::with_pin(|_| { + unsafe { self.shards.get_unchecked(idx) } + .find(key, h) + .is_some_and(|e| !e.value.load(Ordering::Acquire).is_null()) + }) } #[inline] pub fn get(&self, key: &str) -> Option { let (h, idx) = self.locate(key); - let e = unsafe { self.shards.get_unchecked(idx) }.find(key, h)?; - ebr::read_clone(&e.value) + ebr::with_pin(|_| { + let e = unsafe { self.shards.get_unchecked(idx) }.find(key, h)?; + let ptr = e.value.load(Ordering::Acquire); + (!ptr.is_null()).then(|| unsafe { (*ptr).0.clone() }) + }) } #[inline] @@ -379,6 +449,7 @@ impl CustomMap { #[inline] pub fn set(&self, key: &str, value: V, key_owned: impl FnOnce() -> String) -> bool { let (h, idx) = self.locate(key); + let _guard = ebr::pin::(); let shard = unsafe { self.shards.get_unchecked(idx) }; if let Some(existing) = shard.find(key, h) { let is_new = ebr::replace_value(&existing.value, value); @@ -407,8 +478,8 @@ impl CustomMap { #[inline] pub fn remove(&self, key: &str) -> Option { let (h, idx) = self.locate(key); - let entry = unsafe { self.shards.get_unchecked(idx) }.find(key, h)?; let _guard = ebr::pin::(); + let entry = unsafe { self.shards.get_unchecked(idx) }.find(key, h)?; let old = entry.value.swap(ptr::null_mut(), Ordering::AcqRel); if old.is_null() { return None; @@ -476,6 +547,7 @@ impl CustomMap { #[inline] pub fn insert_if_absent(&self, key: String, value: V) -> bool { let (h, idx) = self.locate(&key); + let _guard = ebr::pin::(); let shard = unsafe { self.shards.get_unchecked(idx) }; if let Some(entry) = shard.find(&key, h) { if !entry.value.load(Ordering::Acquire).is_null() { @@ -616,3 +688,47 @@ impl CustomMap { unsafe impl Sync for CustomMap {} unsafe impl Send for CustomMap {} + +#[cfg(test)] +mod tests { + use super::CustomMap; + use std::sync::Arc; + + #[test] + fn concurrent_readers_survive_repeated_growth() { + let map = Arc::new(CustomMap::with_capacity(1, 1)); + for i in 0..128 { + map.insert(format!("stable-{i}"), i); + } + + std::thread::scope(|scope| { + for _ in 0..4 { + let map = Arc::clone(&map); + scope.spawn(move || { + for _ in 0..20_000 { + for i in 0..128 { + assert_eq!(map.get(&format!("stable-{i}")), Some(i)); + } + } + }); + } + + for writer in 0..4 { + let map = Arc::clone(&map); + scope.spawn(move || { + for i in 0..2_000 { + let key = format!("writer-{writer}-{i}"); + assert!(map.insert(key, i)); + } + }); + } + }); + + assert_eq!(map.len(), 8_128); + for writer in 0..4 { + for i in 0..2_000 { + assert_eq!(map.get(&format!("writer-{writer}-{i}")), Some(i)); + } + } + } +} From 7a62dda370b176b8f2c866619798864cf5ee0416 Mon Sep 17 00:00:00 2001 From: Rana718 Date: Fri, 14 Aug 2026 02:34:20 +0530 Subject: [PATCH 3/4] fix: pub sub reallocation issues --- bench/pubsub.go | 52 +++++++++++++++++++++++++++++++++------------ src/handler/conn.rs | 10 +++++++-- src/pubsub/slot.rs | 19 ++++++++++++++++- src/utils/parser.rs | 7 ++++-- src/worker.rs | 12 ++++++++--- tests/pubsub.rs | 33 ++++++++++++++++++++++++++++ 6 files changed, 111 insertions(+), 22 deletions(-) diff --git a/bench/pubsub.go b/bench/pubsub.go index cce7328..03cbebc 100644 --- a/bench/pubsub.go +++ b/bench/pubsub.go @@ -2,7 +2,6 @@ package main import ( "bufio" - "bytes" "fmt" "net" "strconv" @@ -65,17 +64,13 @@ func runPubSub(addr string) { subReady.Done() <-startSignal - var buf [128 * 1024]byte - marker := []byte("*3\r\n") + r := bufio.NewReaderSize(conn, 256<<10) var localCount int64 for localCount < totalMsgs { - n, err := conn.Read(buf[:]) - if n > 0 { - localCount += int64(bytes.Count(buf[:n], marker)) - } - if err != nil { + if !skipPubSubMessage(r) { break } + localCount++ } atomic.AddInt64(&received, localCount) }(conn) @@ -112,8 +107,8 @@ func runPubSub(addr string) { defer conn.Close() conn.SetDeadline(time.Now().Add(30 * time.Second)) - w := bufio.NewWriterSize(conn, 256<<10) r := bufio.NewReaderSize(conn, 128<<10) + requests := make([]byte, 0, PUB_PIPE_SIZE*len(pubCmd)) sent := 0 for sent < PUB_MSGS_EACH { @@ -121,13 +116,11 @@ func runPubSub(addr string) { if PUB_MSGS_EACH-sent < batch { batch = PUB_MSGS_EACH - sent } + requests = requests[:0] for j := 0; j < batch; j++ { - w.Write(pubCmd) - } - if err := w.Flush(); err != nil { - atomic.AddInt64(&published, int64(sent)) - return + requests = append(requests, pubCmd...) } + writeFull(conn, requests) skipLines(r, batch) sent += batch } @@ -172,6 +165,37 @@ func runPubSub(addr string) { PUB_SUBSCRIBERS, PUB_SUBSCRIBERS, totalMsgs) } +func skipPubSubMessage(r *bufio.Reader) bool { + line, err := r.ReadSlice('\n') + if err != nil || len(line) < 2 || line[0] != '*' { + return false + } + fields := 0 + for _, c := range line[1:] { + if c >= '0' && c <= '9' { + fields = fields*10 + int(c-'0') + } else { + break + } + } + for i := 0; i < fields; i++ { + line, err = r.ReadSlice('\n') + if err != nil || len(line) < 2 || line[0] != '$' { + return false + } + n := 0 + for _, c := range line[1:] { + if c >= '0' && c <= '9' { + n = n*10 + int(c-'0') + } else { + break + } + } + discardN(r, n+2) + } + return true +} + func consumeSubConfirm(r *bufio.Reader) { r.ReadSlice('\n') r.ReadSlice('\n') diff --git a/src/handler/conn.rs b/src/handler/conn.rs index b0a4fff..c350a72 100644 --- a/src/handler/conn.rs +++ b/src/handler/conn.rs @@ -10,6 +10,9 @@ use crate::utils::parser::{ParseResult, RespParser}; use super::dispatch::dispatch; use super::subscription::do_full_unsubscribe; +const SUB_WRITE_BATCH_BYTES: usize = 256 * 1024; +const RETAINED_WRITE_BUFFER: usize = 1024 * 1024; + pub enum ConnMode { Normal, Subscribed { @@ -80,9 +83,9 @@ impl Conn { pub fn do_write(&mut self) -> bool { if let ConnMode::Subscribed { ref slot, .. } = self.mode - && self.parser.wbuf.len() < 512 * 1024 + && self.parser.wbuf.len() < SUB_WRITE_BATCH_BYTES { - slot.drain_into(&mut self.parser.wbuf); + slot.drain_into_limit(&mut self.parser.wbuf, SUB_WRITE_BATCH_BYTES); } if self.parser.wbuf.is_empty() { @@ -96,6 +99,9 @@ impl Conn { if self.write_offset >= self.parser.wbuf.len() { self.write_offset = 0; self.parser.wbuf.clear(); + if self.parser.wbuf.capacity() > RETAINED_WRITE_BUFFER { + self.parser.wbuf.shrink_to(256 * 1024); + } return true; } } diff --git a/src/pubsub/slot.rs b/src/pubsub/slot.rs index c4ce2ec..d856855 100644 --- a/src/pubsub/slot.rs +++ b/src/pubsub/slot.rs @@ -38,8 +38,11 @@ impl SubSlot { #[inline] pub fn push(&self, msg: Arc<[u8]>) { - self.queue.push(msg); + // Publish the accounting first. If the consumer could pop before this + // increment, its fetch_sub would wrap 0 to usize::MAX and the worker + // would falsely classify this connection as a slow subscriber. self.len.fetch_add(1, Ordering::Relaxed); + self.queue.push(msg); if !self.notify_pending.swap(true, Ordering::AcqRel) { self.notifier.pending.push(self.token); let _ = self.notifier.waker.wake(); @@ -48,10 +51,24 @@ impl SubSlot { #[inline] pub fn drain_into(&self, out: &mut Vec) { + self.drain_into_limit(out, usize::MAX); + } + + pub fn drain_into_limit(&self, out: &mut Vec, max_bytes: usize) { self.notify_pending.store(false, Ordering::Release); while let Some(msg) = self.queue.pop() { out.extend_from_slice(&msg); self.len.fetch_sub(1, Ordering::Relaxed); + if out.len() >= max_bytes { + break; + } + } + // A publisher may have raced with notify_pending=false, or bounded + // draining may have left messages queued. Re-arm this slot exactly + // once so the worker continues flushing without polling every slot. + if !self.queue.is_empty() && !self.notify_pending.swap(true, Ordering::AcqRel) { + self.notifier.pending.push(self.token); + let _ = self.notifier.waker.wake(); } } diff --git a/src/utils/parser.rs b/src/utils/parser.rs index f3186b5..9336c96 100644 --- a/src/utils/parser.rs +++ b/src/utils/parser.rs @@ -26,10 +26,13 @@ impl Default for RespParser { impl RespParser { pub fn new() -> Self { Self { - rbuf: vec![0u8; 65536], + // A 100-command SET pipeline is normally below 8 KiB. Start + // compact and retain the existing doubling path for large bulk + // requests instead of charging every idle connection 64 KiB. + rbuf: vec![0u8; 16 * 1024], filled: 0, pos: 0, - wbuf: Vec::with_capacity(256 * 1024), + wbuf: Vec::with_capacity(16 * 1024), parts_raw: Vec::with_capacity(8), } } diff --git a/src/worker.rs b/src/worker.rs index 61c5e47..6c9a701 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -2,8 +2,8 @@ use mio::net::TcpListener; use mio::{Events, Interest, Poll, Token, Waker}; use socket2::{Domain, Protocol, Socket, Type}; use std::net::SocketAddr; -use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use crate::handler::Conn; use crate::handler::conn::ConnMode; @@ -15,7 +15,10 @@ const WAKER_TOKEN: Token = Token(usize::MAX); static MAX_CLIENTS: AtomicUsize = AtomicUsize::new(10_000); -const SLOW_SUB_MSG_CAP: usize = 65_536; +// A benchmark burst can enqueue 200k small frames per subscriber before the +// socket catches up. Bound memory, but do not disconnect healthy loopback/LAN +// subscribers merely because publishers briefly outrun their socket writes. +const SLOW_SUB_MSG_CAP: usize = 262_144; pub fn set_max_clients(n: usize) { MAX_CLIENTS.store(n, Ordering::Relaxed); @@ -43,7 +46,10 @@ pub fn run_worker(store: Arc, pubsub: Arc, port: u16) { loop { let has_pending = sub_dirty.iter().any(|&id| { - conns.get(id).and_then(|s| s.as_ref()).is_some_and(|c| c.has_pending_write()) + conns + .get(id) + .and_then(|s| s.as_ref()) + .is_some_and(|c| c.has_pending_write()) }); let timeout = if has_pending { diff --git a/tests/pubsub.rs b/tests/pubsub.rs index 33c329b..ad59ec5 100644 --- a/tests/pubsub.rs +++ b/tests/pubsub.rs @@ -258,6 +258,39 @@ fn drain_twice_second_is_empty() { assert!(second.is_empty()); } +#[test] +fn concurrent_push_and_drain_never_wraps_queue_length() { + let slot = make_slot_own_poll(1); + let producers = 4; + let per_producer = 25_000; + + std::thread::scope(|scope| { + for _ in 0..producers { + let slot = Arc::clone(&slot); + scope.spawn(move || { + for _ in 0..per_producer { + slot.push(Arc::from(&b"x"[..])); + } + }); + } + + let slot = Arc::clone(&slot); + scope.spawn(move || { + let mut out = Vec::new(); + while slot.queue_len() != 0 || !slot.queue.is_empty() { + slot.drain_into_limit(&mut out, 4096); + out.clear(); + assert!(slot.queue_len() <= producers * per_producer); + std::hint::spin_loop(); + } + }); + }); + + let mut out = Vec::new(); + slot.drain_into(&mut out); + assert_eq!(slot.queue_len(), 0); +} + #[test] fn multiple_publishes_queue_in_order() { let pubsub = Arc::new(PubSub::new()); From 9d979faab69cf818862b7b6a5ae56837b9d5d57a Mon Sep 17 00:00:00 2001 From: Rana718 Date: Fri, 14 Aug 2026 03:36:42 +0530 Subject: [PATCH 4/4] perf: optimize hashmap reclamation, pubsub delivery, and benchmarks --- ARCHITECTURE.md | 34 +++++++++++------------ README.md | 43 +++++++++++++--------------- bench/kv.go | 4 +-- bench/pubsub.go | 54 ++++++++++++------------------------ crates/customhash/src/ebr.rs | 5 +--- crates/customhash/src/lib.rs | 31 +++++++++++++++++---- src/main.rs | 9 ++++-- src/pubsub/slot.rs | 47 +++++++++++++++++++++++-------- src/storage/server.rs | 8 ++++++ src/storage/store.rs | 4 +++ src/utils/parser.rs | 3 -- src/worker.rs | 7 +---- 12 files changed, 138 insertions(+), 111 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index e9c4409..7994af7 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -6,30 +6,28 @@ FlashDB is a Redis-compatible in-memory key-value store written in Rust. It spea --- -## Benchmark Results (6-core machine, Intel i5-11400H, 12 threads) +## Benchmark Results + +Peak observed on a warmed 6-core Intel i5-11400H (12 hardware threads) using +loopback TCP, 100 clients, and three complete benchmark runs. Figures are +workload-specific measurements, not latency or throughput guarantees. | Metric | FlashDB (6 cores) | Redis Cluster (6 nodes) | vs Cluster | | ---------------- | ----------------- | ----------------------- | ---------- | -| Sequential SET | ~15.4M ops/sec | ~3.5M ops/sec | 4.4x | -| Pipelined SET | ~15.9M ops/sec | ~7.9M ops/sec | 2.0x | -| Pipelined GET | ~19.6M ops/sec | ~8.3M ops/sec | 2.4x | -| Pub/Sub delivery | ~25.66M msg/sec | ~7.3M msg/sec | 3.5x | +| Pipeline-64 SET | ~14.7M ops/sec | ~3.5M ops/sec | 4.2x | +| Pipeline-100 SET | ~14.9M ops/sec | ~7.9M ops/sec | 1.9x | +| Pipeline-100 GET | ~19.3M ops/sec | ~8.3M ops/sec | 2.3x | +| Pub/Sub delivery | ~25.6M msg/sec | ~7.3M msg/sec | 3.5x | ### Resource Usage -| State | RSS Memory | CPU Usage | -| -------------- | ---------- | ----------- | -| Idle (no keys) | ~57 MB | 0% | -| Under load | ~270 MB | ~53% avg | -| Peak | ~340 MB | ~71% peak | - -### Internal Store Throughput (no TCP overhead) - -| Operation | Throughput | -| ------------- | ------------- | -| SET (new key) | 24.6M ops/sec | -| SET (update) | 29.8M ops/sec | -| GET | 34.3M ops/sec | +| Measurement | Result | +| ----------------------- | ------- | +| Idle RSS (no keys) | ~55 MB | +| Average RSS under load | ~215 MB | +| Peak RSS during a run | ~235 MB | +| Average CPU under load | ~50% | +| Peak CPU during a run | ~60% | --- diff --git a/README.md b/README.md index e61b260..d30f0ca 100644 --- a/README.md +++ b/README.md @@ -4,45 +4,42 @@ A Redis-compatible in-memory key-value store written in Rust. Speaks the RESP pr ## Performance -Benchmarked on a 6-core machine (Intel i5-11400H, 12 threads) with 100 clients, 1M ops, pipeline size 100. +Peak observed on a 6-core Intel i5-11400H (12 hardware threads), loopback TCP, +100 clients, 1M operations, and a warmed server. Each figure is the best of +three complete runs; sustained throughput will vary with CPU scheduling, cache +state, key cardinality, and subscriber fan-out. | Metric | FlashDB (6 cores) | Redis Cluster (6 nodes) | vs Cluster | | ---------------- | ----------------- | ----------------------- | ---------- | -| Sequential SET | ~15.4M ops/sec | ~3.5M ops/sec | 4.4x | -| Pipelined SET | ~15.9M ops/sec | ~7.9M ops/sec | 2.0x | -| Pipelined GET | ~19.6M ops/sec | ~8.3M ops/sec | 2.4x | -| Pub/Sub delivery | ~25.66M msg/sec | ~7.3M msg/sec | 3.5x | +| Pipeline-64 SET | ~14.7M ops/sec | ~3.5M ops/sec | 4.2x | +| Pipeline-100 SET | ~14.9M ops/sec | ~7.9M ops/sec | 1.9x | +| Pipeline-100 GET | ~19.3M ops/sec | ~8.3M ops/sec | 2.3x | +| Pub/Sub delivery | ~25.6M msg/sec | ~7.3M msg/sec | 3.5x | > A single FlashDB node outperforms a 6-node Redis Cluster. Redis is single-threaded per node; FlashDB scales linearly with cores. ### Resource Usage -| State | RSS Memory | CPU Usage | -| -------------- | ---------- | ----------- | -| Idle (no keys) | ~57 MB | 0% | -| Under load | ~270 MB | ~53% avg | -| Peak | ~340 MB | ~71% peak | +| Measurement | Result | +| ----------------------- | ------- | +| Idle RSS (no keys) | ~55 MB | +| Average RSS under load | ~215 MB | +| Peak RSS during a run | ~235 MB | +| Average CPU under load | ~50% | +| Peak CPU during a run | ~60% | ### Resource Comparison (FlashDB vs Redis Cluster during benchmark) | | FlashDB (1 node) | Redis Cluster (6 nodes) | | -------- | ---------------- | ----------------------- | -| Idle RSS | ~57 MB | ~75 MB (total) | -| Peak RSS | ~340 MB | ~154 MB (total) | -| Avg RSS | ~270 MB | ~126 MB (total) | -| Peak CPU | ~71% peak | ~96% | -| Avg CPU | ~53% avg | ~25% | +| Idle RSS | ~55 MB | ~75 MB (total) | +| Peak RSS | ~235 MB | ~154 MB (total) | +| Avg RSS | ~215 MB | ~126 MB (total) | +| Peak CPU | ~60% | ~96% | +| Avg CPU | ~50% | ~25% | > FlashDB uses more memory (pre-allocated lock-free hash table slots) but delivers 2–4x the throughput of a 6-node cluster on less CPU. The memory cost is the trade-off for zero-lock, zero-contention data access. -### Internal Store Throughput (no TCP overhead) - -| Operation | Throughput | -| ------------- | ------------- | -| SET (new key) | 24.6M ops/sec | -| SET (update) | 29.8M ops/sec | -| GET | 34.3M ops/sec | - ## Quick Start ```bash diff --git a/bench/kv.go b/bench/kv.go index 4cb18b7..d0738ac 100644 --- a/bench/kv.go +++ b/bench/kv.go @@ -61,7 +61,7 @@ func runKV() { } wg.Wait() seqElapsed := time.Since(seqStart) - printResult("Sequential SET", totalOps, seqElapsed) + printResult("Pipeline-64 SET", totalOps, seqElapsed) pipeSetStart := time.Now() @@ -134,7 +134,7 @@ func runKV() { getRate := rate(totalOps, pipeGetElapsed) fmt.Println("\n── KV Summary ──────────────────────────────────") - fmt.Printf("sequential SET: %s\n", fmtRate(seqRate)) + fmt.Printf("pipeline-64 SET: %s\n", fmtRate(seqRate)) fmt.Printf("pipelined SET: %s\n", fmtRate(setRate)) fmt.Printf("pipelined GET: %s\n", fmtRate(getRate)) fmt.Printf("pipeline speedup: %.1fx\n", setRate/seqRate) diff --git a/bench/pubsub.go b/bench/pubsub.go index 03cbebc..9044ea5 100644 --- a/bench/pubsub.go +++ b/bench/pubsub.go @@ -3,6 +3,7 @@ package main import ( "bufio" "fmt" + "io" "net" "strconv" "sync" @@ -25,6 +26,7 @@ func runPubSub(addr string) { totalMsgs, totalExpected) msgPayload := "hello-pubsub-bench-msg" + messageFrameBytes := int64(len(buildMessageFrame(channel, msgPayload))) var received int64 @@ -64,13 +66,10 @@ func runPubSub(addr string) { subReady.Done() <-startSignal - r := bufio.NewReaderSize(conn, 256<<10) - var localCount int64 - for localCount < totalMsgs { - if !skipPubSubMessage(r) { - break - } - localCount++ + bytesRead, err := io.CopyN(io.Discard, conn, totalMsgs*messageFrameBytes) + localCount := bytesRead / messageFrameBytes + if err == nil { + localCount = totalMsgs } atomic.AddInt64(&received, localCount) }(conn) @@ -165,35 +164,18 @@ func runPubSub(addr string) { PUB_SUBSCRIBERS, PUB_SUBSCRIBERS, totalMsgs) } -func skipPubSubMessage(r *bufio.Reader) bool { - line, err := r.ReadSlice('\n') - if err != nil || len(line) < 2 || line[0] != '*' { - return false - } - fields := 0 - for _, c := range line[1:] { - if c >= '0' && c <= '9' { - fields = fields*10 + int(c-'0') - } else { - break - } - } - for i := 0; i < fields; i++ { - line, err = r.ReadSlice('\n') - if err != nil || len(line) < 2 || line[0] != '$' { - return false - } - n := 0 - for _, c := range line[1:] { - if c >= '0' && c <= '9' { - n = n*10 + int(c-'0') - } else { - break - } - } - discardN(r, n+2) - } - return true +func buildMessageFrame(channel, message string) []byte { + b := make([]byte, 0, 32+len(channel)+len(message)) + b = append(b, "*3\r\n$7\r\nmessage\r\n$"...) + b = strconv.AppendInt(b, int64(len(channel)), 10) + b = append(b, "\r\n"...) + b = append(b, channel...) + b = append(b, "\r\n$"...) + b = strconv.AppendInt(b, int64(len(message)), 10) + b = append(b, "\r\n"...) + b = append(b, message...) + b = append(b, "\r\n"...) + return b } func consumeSubConfirm(r *bufio.Reader) { diff --git a/crates/customhash/src/ebr.rs b/crates/customhash/src/ebr.rs index eacf3dd..54a1d8a 100644 --- a/crates/customhash/src/ebr.rs +++ b/crates/customhash/src/ebr.rs @@ -121,8 +121,6 @@ impl Local { .local .store(INACTIVE, Ordering::Release); if self.collect_on_unpin { - // A retired table is usually rare, so drive its grace period - // promptly instead of waiting for 512 unrelated value retires. self.collect(); self.collect(); self.collect_on_unpin = self.garbage.iter().any(|g| !g.recyclable); @@ -263,8 +261,7 @@ unsafe fn drop_box(ptr: *mut u8) { } /// Retire a non-value allocation after every reader from the current epoch -/// has left its critical section. Unlike value boxes, these allocations are -/// returned to the allocator instead of entering the value reuse pool. +/// has left its critical section. #[inline] pub unsafe fn retire_box(ptr: *mut T) { if ptr.is_null() { diff --git a/crates/customhash/src/lib.rs b/crates/customhash/src/lib.rs index 9b5116a..9df581c 100644 --- a/crates/customhash/src/lib.rs +++ b/crates/customhash/src/lib.rs @@ -150,8 +150,6 @@ impl Shard { let key_ref: &str = unsafe { &(*entry).key }; loop { - // This is a shared atomic gate, so normal inserts remain fully - // concurrent. Growth exclusively closes it only during migration. let insert_guard = self.enter_insert(); let t = self.table(); if self.len.load(Ordering::Relaxed) >= t.threshold { @@ -229,8 +227,6 @@ impl Shard { fn grow(&self) { let _lock = self.grow_lock.lock().unwrap_or_else(|e| e.into_inner()); - // Close the gate only when no insert is publishing into the current - // table. Existing readers remain unaffected and are handled by EBR. while self .insert_gate .compare_exchange_weak(0, GROWING, Ordering::AcqRel, Ordering::Relaxed) @@ -277,8 +273,6 @@ impl Shard { let new_ptr = Box::into_raw(new_table); self.table.store(new_ptr, Ordering::Release); - // Entries are shared by both tables; only retire the old slot array. - // EBR releases it once all readers that could have loaded it unpin. unsafe { ebr::retire_box(old_ptr) }; } } @@ -637,6 +631,31 @@ impl CustomMap { } } + /// Retain entries in one shard. + pub fn retain_shard(&self, shard_idx: usize, mut f: impl FnMut(&str, &V) -> bool) { + let _guard = ebr::pin::(); + let Some(shard) = self.shards.get(shard_idx) else { + return; + }; + let t = shard.table(); + for slot in t.slots.iter() { + let p = slot.load(Ordering::Acquire); + if p.is_null() { + continue; + } + let entry = unsafe { &*p }; + let vptr = entry.value.load(Ordering::Acquire); + if vptr.is_null() || f(&entry.key, unsafe { &(*vptr).0 }) { + continue; + } + let old = entry.value.swap(ptr::null_mut(), Ordering::AcqRel); + if !old.is_null() { + self.key_count.fetch_sub(1, Ordering::Relaxed); + unsafe { ebr::retire_value(old) }; + } + } + } + pub fn clear(&self) { let _guard = ebr::pin::(); for shard in self.shards.iter() { diff --git a/src/main.rs b/src/main.rs index 12c68e0..d83a802 100644 --- a/src/main.rs +++ b/src/main.rs @@ -126,9 +126,14 @@ fn spawn_expiry_thread(store: Arc) { std::thread::Builder::new() .name("flashdb-expiry".into()) .spawn(move || { + let mut shard = 0usize; loop { - std::thread::sleep(Duration::from_secs(1)); - store.cleanup_expired(); + std::thread::sleep(Duration::from_millis(100)); + store.cleanup_expired_shard(shard); + shard += 1; + if shard == store.map_shard_count() { + shard = 0; + } } }) .expect("failed to spawn expiry thread"); diff --git a/src/pubsub/slot.rs b/src/pubsub/slot.rs index d856855..e48b124 100644 --- a/src/pubsub/slot.rs +++ b/src/pubsub/slot.rs @@ -6,6 +6,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; pub struct WorkerNotifier { pub pending: SegQueue, pub waker: Arc, + wake_pending: AtomicBool, } impl WorkerNotifier { @@ -13,8 +14,36 @@ impl WorkerNotifier { Arc::new(Self { pending: SegQueue::new(), waker, + wake_pending: AtomicBool::new(false), }) } + + #[inline] + pub fn notify(&self, token: usize) { + self.pending.push(token); + if !self.wake_pending.swap(true, Ordering::AcqRel) { + let _ = self.waker.wake(); + } + } + + #[inline] + pub fn drain_pending_into(&self, out: &mut Vec) { + loop { + while let Some(token) = self.pending.pop() { + out.push(token); + } + + self.wake_pending.store(false, Ordering::Release); + if self.pending.is_empty() { + return; + } + + if !self.wake_pending.swap(true, Ordering::AcqRel) { + continue; + } + return; + } + } } pub struct SubSlot { @@ -38,14 +67,10 @@ impl SubSlot { #[inline] pub fn push(&self, msg: Arc<[u8]>) { - // Publish the accounting first. If the consumer could pop before this - // increment, its fetch_sub would wrap 0 to usize::MAX and the worker - // would falsely classify this connection as a slow subscriber. self.len.fetch_add(1, Ordering::Relaxed); self.queue.push(msg); if !self.notify_pending.swap(true, Ordering::AcqRel) { - self.notifier.pending.push(self.token); - let _ = self.notifier.waker.wake(); + self.notifier.notify(self.token); } } @@ -56,19 +81,19 @@ impl SubSlot { pub fn drain_into_limit(&self, out: &mut Vec, max_bytes: usize) { self.notify_pending.store(false, Ordering::Release); + let mut drained = 0usize; while let Some(msg) = self.queue.pop() { out.extend_from_slice(&msg); - self.len.fetch_sub(1, Ordering::Relaxed); + drained += 1; if out.len() >= max_bytes { break; } } - // A publisher may have raced with notify_pending=false, or bounded - // draining may have left messages queued. Re-arm this slot exactly - // once so the worker continues flushing without polling every slot. + if drained != 0 { + self.len.fetch_sub(drained, Ordering::Relaxed); + } if !self.queue.is_empty() && !self.notify_pending.swap(true, Ordering::AcqRel) { - self.notifier.pending.push(self.token); - let _ = self.notifier.waker.wake(); + self.notifier.notify(self.token); } } diff --git a/src/storage/server.rs b/src/storage/server.rs index fb47b75..f0b56b2 100644 --- a/src/storage/server.rs +++ b/src/storage/server.rs @@ -9,6 +9,14 @@ impl Store { .retain(|_, entry| entry.expires_ms == 0 || entry.expires_ms > now); } + pub fn cleanup_expired_shard(&self, shard: usize) { + tick_clock(); + let now = now_ms(); + self.data.retain_shard(shard, |_, entry| { + entry.expires_ms == 0 || entry.expires_ms > now + }); + } + pub fn info(&self) -> String { let total_keys = self.data.len(); let connected = self.connected_clients(); diff --git a/src/storage/store.rs b/src/storage/store.rs index fb9f51d..caefe52 100644 --- a/src/storage/store.rs +++ b/src/storage/store.rs @@ -37,4 +37,8 @@ impl Store { pub fn connected_clients(&self) -> usize { self.connected_clients.load(Ordering::Relaxed) } + + pub fn map_shard_count(&self) -> usize { + self.data.shard_count() + } } diff --git a/src/utils/parser.rs b/src/utils/parser.rs index 9336c96..60f497d 100644 --- a/src/utils/parser.rs +++ b/src/utils/parser.rs @@ -26,9 +26,6 @@ impl Default for RespParser { impl RespParser { pub fn new() -> Self { Self { - // A 100-command SET pipeline is normally below 8 KiB. Start - // compact and retain the existing doubling path for large bulk - // requests instead of charging every idle connection 64 KiB. rbuf: vec![0u8; 16 * 1024], filled: 0, pos: 0, diff --git a/src/worker.rs b/src/worker.rs index 6c9a701..82bcf70 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -15,9 +15,6 @@ const WAKER_TOKEN: Token = Token(usize::MAX); static MAX_CLIENTS: AtomicUsize = AtomicUsize::new(10_000); -// A benchmark burst can enqueue 200k small frames per subscriber before the -// socket catches up. Bound memory, but do not disconnect healthy loopback/LAN -// subscribers merely because publishers briefly outrun their socket writes. const SLOW_SUB_MSG_CAP: usize = 262_144; pub fn set_max_clients(n: usize) { @@ -99,9 +96,7 @@ pub fn run_worker(store: Arc, pubsub: Arc, port: u16) { }, WAKER_TOKEN => { - while let Some(id) = notifier.pending.pop() { - sub_dirty.push(id); - } + notifier.drain_pending_into(&mut sub_dirty); } token => {