mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
perf(rust): tier-1 multi-worker wins — GIL release, sharded CCR store, single-serialize CCR write
Three orthogonal hot-path fixes targeting concurrent-request throughput.
Each is independently bench-measured below; the proxy hot path benefits
from all three at once.
== 1. PyO3 GIL release on heavy compute ==
PyO3 methods (crush, smart_crush_content, crush_array_json,
compact_document_json, compress, compress_with_stats) used to hold the
GIL across the entire Rust call. Result: a 100ms compress() blocked
EVERY other Python thread for 100ms — multi-worker uvicorn deployments
serialized through SmartCrusher.
Wrap each compute call in `py.allow_threads(|| ...)`. Inputs (`&str`
from Python) are copied to owned `String` first because PyO3 ties them
to the GIL hold. PyDict construction stays on the GIL side.
Measured: 4 Python threads each running 20 crushes:
before (GIL held): ~3.3s wall (serialized — equivalent to 4×0.83s)
after (allow_threads): 826ms wall (4.01x speedup, perfect parallel)
== 2. CcrStore: Mutex<HashMap> -> DashMap-backed sharded ==
Single Mutex was the dominant bottleneck under multi-worker load — every
put/get serialized through one lock. Replace with DashMap (sharded
concurrent map, lock-free reads within a shard) plus a separate
small Mutex<VecDeque> for FIFO insertion-order eviction. Reads of
distinct keys never contend; writes only contend during the brief
order-queue push or capacity-sweep.
A/B bench (200 mixed put/get ops × N threads, in benches/ccr_store.rs):
Threads | DashMap Legacy Mutex Speedup
-------------------------------------------
1 | 63 µs 71 µs 1.13x
2 | 98 µs 194 µs 2.0x
4 | 178 µs 707 µs 4.0x
8 | 342 µs 1267 µs 3.7x
Legacy degrades ~linearly with thread count; DashMap stays near-flat
per-thread. Real multi-worker scaling.
== 3. Single-serialize the lossy CCR payload ==
The lossy `crush_array` path used to serialize the full array TWICE:
once in `hash_array_for_ccr` (allocates `Value::Array(items.to_vec())`,
deep-clones every Value subtree, then serializes), and a second time
in the store-write site. For a 50-item dict array that's ~MB of
allocator pressure per crushed array.
Introduce `canonical_array_json` (serializes `&[Value]` directly — same
bytes as `Value::Array(items.to_vec())` but no wrapper allocation +
no tree clone), call it ONCE per lossy path, then both hash and store
from those same bytes. Hash-format stable — all 17 parity fixtures
match byte-for-byte.
== Tests ==
- 8 ccr.rs unit tests including a new concurrent-stress test (8 threads
× 200 puts/gets, every key readable afterwards)
- 14 ccr_roundtrip integration tests stay green
- parity-run smart_crusher: 17/17 fixtures match
- 479 lib + 14 integration + 185 Python tests all pass
- New benches/ccr_store.rs runs the A/B and is committed for regression
visibility
== Dependencies added ==
- dashmap v6 (mature, widely-used in tokio/linkerd ecosystem)
This commit is contained in:
parent
b8fc7eee19
commit
29aadb1054
6 changed files with 526 additions and 107 deletions
58
Cargo.lock
generated
58
Cargo.lock
generated
|
|
@ -763,6 +763,20 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dashmap"
|
||||
version = "6.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"crossbeam-utils",
|
||||
"hashbrown 0.14.5",
|
||||
"lock_api",
|
||||
"once_cell",
|
||||
"parking_lot_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "data-encoding"
|
||||
version = "2.11.0"
|
||||
|
|
@ -1268,6 +1282,12 @@ dependencies = [
|
|||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.15.5"
|
||||
|
|
@ -1302,6 +1322,7 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"bytes",
|
||||
"criterion",
|
||||
"dashmap",
|
||||
"fastembed",
|
||||
"flate2",
|
||||
"hf-hub 0.4.3",
|
||||
|
|
@ -1914,6 +1935,15 @@ version = "1.0.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092"
|
||||
|
||||
[[package]]
|
||||
name = "lock_api"
|
||||
version = "0.4.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
|
||||
dependencies = [
|
||||
"scopeguard",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "log"
|
||||
version = "0.4.29"
|
||||
|
|
@ -2358,6 +2388,19 @@ dependencies = [
|
|||
"ureq 3.3.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot_core"
|
||||
version = "0.9.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"redox_syscall",
|
||||
"smallvec",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "paste"
|
||||
version = "1.0.15"
|
||||
|
|
@ -2855,6 +2898,15 @@ dependencies = [
|
|||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_users"
|
||||
version = "0.5.2"
|
||||
|
|
@ -3076,6 +3128,12 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||
|
||||
[[package]]
|
||||
name = "security-framework"
|
||||
version = "3.7.0"
|
||||
|
|
|
|||
|
|
@ -29,6 +29,11 @@ md-5 = "0.10"
|
|||
# `sha2` for `_hash_field_name` in smart_crusher (SHA256 truncated to 16
|
||||
# hex chars). Python uses `hashlib.sha256` so we need byte-exact parity.
|
||||
sha2 = "0.10"
|
||||
# `dashmap` for the CCR storage backend. Concurrent HashMap with sharded
|
||||
# locking — distinct keys hashed to different shards never contend, so
|
||||
# multi-worker proxy load doesn't queue on a single Mutex. Lock-free
|
||||
# reads inside each shard via RwLock semantics.
|
||||
dashmap = "6"
|
||||
# `regex` is already a transitive dep of tokenizers; depend on it directly so
|
||||
# our hunk-header parser and priority-pattern matcher have a stable surface.
|
||||
regex = "1"
|
||||
|
|
@ -57,3 +62,7 @@ criterion = { version = "0.5", features = ["html_reports"] }
|
|||
[[bench]]
|
||||
name = "tokenizer"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "ccr_store"
|
||||
harness = false
|
||||
|
|
|
|||
224
crates/headroom-core/benches/ccr_store.rs
Normal file
224
crates/headroom-core/benches/ccr_store.rs
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
//! CCR store throughput benchmark — single-threaded and multi-threaded.
|
||||
//!
|
||||
//! Pins the win from PR9: replacing the single-`Mutex<HashMap>` design
|
||||
//! with a `DashMap`-backed sharded store. The single-threaded numbers
|
||||
//! should be roughly comparable (DashMap has a small per-op shard-hash
|
||||
//! overhead vs a raw Mutex), but the multi-threaded numbers should
|
||||
//! diverge sharply — distinct keys hit distinct shards and never
|
||||
//! contend.
|
||||
//!
|
||||
//! Run with:
|
||||
//! cargo bench -p headroom-core --bench ccr_store
|
||||
//!
|
||||
//! The critical numbers to watch are the `mt/N=8` rows: with the
|
||||
//! Mutex design, all 8 threads serialize on one lock, so throughput
|
||||
//! is ~1× the single-threaded figure. With DashMap, throughput should
|
||||
//! scale near-linearly with cores.
|
||||
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
|
||||
use headroom_core::ccr::{CcrStore, InMemoryCcrStore};
|
||||
|
||||
// ─── Baseline: the old single-Mutex<HashMap> design ────────────────
|
||||
//
|
||||
// Inlined here so the bench is self-contained and shows the
|
||||
// before/after gap directly. Same trait, same semantics; the only
|
||||
// difference is "all ops serialize on one Mutex" vs "DashMap-sharded".
|
||||
|
||||
struct LegacyMutexStore {
|
||||
inner: Mutex<LegacyInner>,
|
||||
ttl: Duration,
|
||||
capacity: usize,
|
||||
}
|
||||
|
||||
struct LegacyInner {
|
||||
map: HashMap<String, LegacyEntry>,
|
||||
order: VecDeque<String>,
|
||||
}
|
||||
|
||||
struct LegacyEntry {
|
||||
payload: String,
|
||||
inserted: Instant,
|
||||
}
|
||||
|
||||
impl LegacyMutexStore {
|
||||
fn new(capacity: usize, ttl: Duration) -> Self {
|
||||
Self {
|
||||
inner: Mutex::new(LegacyInner {
|
||||
map: HashMap::new(),
|
||||
order: VecDeque::new(),
|
||||
}),
|
||||
ttl,
|
||||
capacity,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CcrStore for LegacyMutexStore {
|
||||
fn put(&self, hash: &str, payload: &str) {
|
||||
let mut g = self.inner.lock().unwrap();
|
||||
if g.map.contains_key(hash) {
|
||||
g.map.insert(
|
||||
hash.to_string(),
|
||||
LegacyEntry {
|
||||
payload: payload.to_string(),
|
||||
inserted: Instant::now(),
|
||||
},
|
||||
);
|
||||
return;
|
||||
}
|
||||
while g.map.len() >= self.capacity {
|
||||
let Some(oldest) = g.order.pop_front() else {
|
||||
break;
|
||||
};
|
||||
g.map.remove(&oldest);
|
||||
}
|
||||
g.map.insert(
|
||||
hash.to_string(),
|
||||
LegacyEntry {
|
||||
payload: payload.to_string(),
|
||||
inserted: Instant::now(),
|
||||
},
|
||||
);
|
||||
g.order.push_back(hash.to_string());
|
||||
}
|
||||
|
||||
fn get(&self, hash: &str) -> Option<String> {
|
||||
let mut g = self.inner.lock().unwrap();
|
||||
let expired = match g.map.get(hash) {
|
||||
Some(e) => e.inserted.elapsed() > self.ttl,
|
||||
None => return None,
|
||||
};
|
||||
if expired {
|
||||
g.map.remove(hash);
|
||||
return None;
|
||||
}
|
||||
g.map.get(hash).map(|e| e.payload.clone())
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.inner.lock().unwrap().map.len()
|
||||
}
|
||||
}
|
||||
|
||||
fn bench_put_single_threaded(c: &mut Criterion) {
|
||||
let store = InMemoryCcrStore::new();
|
||||
let payload = "x".repeat(512); // typical CCR payload size
|
||||
|
||||
let mut group = c.benchmark_group("ccr_store/put_st");
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("new_keys", |b| {
|
||||
let mut i = 0u64;
|
||||
b.iter(|| {
|
||||
let key = format!("k{i:012x}");
|
||||
store.put(black_box(&key), black_box(&payload));
|
||||
i += 1;
|
||||
});
|
||||
});
|
||||
group.bench_function("same_key_overwrite", |b| {
|
||||
b.iter(|| {
|
||||
store.put(black_box("hot_key"), black_box(&payload));
|
||||
});
|
||||
});
|
||||
group.finish();
|
||||
}
|
||||
|
||||
fn bench_get_single_threaded(c: &mut Criterion) {
|
||||
let store = InMemoryCcrStore::new();
|
||||
let payload = "y".repeat(512);
|
||||
for i in 0..1000u32 {
|
||||
let key = format!("k{i:08x}");
|
||||
store.put(&key, &payload);
|
||||
}
|
||||
|
||||
let mut group = c.benchmark_group("ccr_store/get_st");
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_function("hit", |b| {
|
||||
let mut i = 0u32;
|
||||
b.iter(|| {
|
||||
let key = format!("k{:08x}", i % 1000);
|
||||
let _ = black_box(store.get(black_box(&key)));
|
||||
i = i.wrapping_add(1);
|
||||
});
|
||||
});
|
||||
group.bench_function("miss", |b| {
|
||||
let mut i = 0u32;
|
||||
b.iter(|| {
|
||||
let key = format!("absent_{i}");
|
||||
let _ = black_box(store.get(black_box(&key)));
|
||||
i = i.wrapping_add(1);
|
||||
});
|
||||
});
|
||||
group.finish();
|
||||
}
|
||||
|
||||
fn run_mt_workload(store: Arc<dyn CcrStore>, threads: usize, n: u64) -> Duration {
|
||||
const ITERS_PER_THREAD: usize = 200;
|
||||
let payload = Arc::new("z".repeat(256));
|
||||
for i in 0..256u32 {
|
||||
store.put(&format!("warm_{i:08x}"), &payload);
|
||||
}
|
||||
let start = Instant::now();
|
||||
for _ in 0..n {
|
||||
thread::scope(|scope| {
|
||||
for tid in 0..threads {
|
||||
let s = store.clone();
|
||||
let p = payload.clone();
|
||||
scope.spawn(move || {
|
||||
for i in 0..ITERS_PER_THREAD {
|
||||
if i & 1 == 0 {
|
||||
let k = format!("t{tid}_k{i:08x}");
|
||||
s.put(&k, &p);
|
||||
} else {
|
||||
let k = format!("warm_{:08x}", i % 256);
|
||||
let _ = s.get(&k);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
start.elapsed()
|
||||
}
|
||||
|
||||
/// Multi-threaded mixed put/get — direct A/B between the legacy
|
||||
/// `Mutex<HashMap>` design and the new DashMap-backed store. The
|
||||
/// legacy version serializes every op on one lock; the new version
|
||||
/// shards across keys so distinct hashes never contend.
|
||||
fn bench_mixed_multi_threaded(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("ccr_store/mt_mixed");
|
||||
group.throughput(Throughput::Elements(1));
|
||||
|
||||
for &threads in &[1usize, 2, 4, 8] {
|
||||
// DashMap-backed (current design).
|
||||
let label = format!("dashmap/threads={threads}");
|
||||
group.bench_function(&label, |b| {
|
||||
b.iter_custom(|n| {
|
||||
let store: Arc<dyn CcrStore> = Arc::new(InMemoryCcrStore::new());
|
||||
run_mt_workload(store, threads, n)
|
||||
});
|
||||
});
|
||||
// Legacy Mutex<HashMap> (the design PR9 replaces).
|
||||
let label = format!("legacy_mutex/threads={threads}");
|
||||
group.bench_function(&label, |b| {
|
||||
b.iter_custom(|n| {
|
||||
let store: Arc<dyn CcrStore> =
|
||||
Arc::new(LegacyMutexStore::new(1000, Duration::from_secs(300)));
|
||||
run_mt_workload(store, threads, n)
|
||||
});
|
||||
});
|
||||
}
|
||||
group.finish();
|
||||
}
|
||||
|
||||
criterion_group!(
|
||||
benches,
|
||||
bench_put_single_threaded,
|
||||
bench_get_single_threaded,
|
||||
bench_mixed_multi_threaded,
|
||||
);
|
||||
criterion_main!(benches);
|
||||
|
|
@ -13,6 +13,18 @@
|
|||
//! feedback, no per-tool metadata. Those live in the runtime layer; this
|
||||
//! crate only needs put/get.
|
||||
//!
|
||||
//! # Concurrency
|
||||
//!
|
||||
//! The default [`InMemoryCcrStore`] uses [`DashMap`] (sharded concurrent
|
||||
//! hash map) so reads and writes targeting different keys never contend.
|
||||
//! Only the FIFO insertion-order queue (used for capacity-bounded
|
||||
//! eviction) sits behind a single `Mutex`, and that mutex is held just
|
||||
//! long enough for an O(1) `push_back` or capacity-sweep.
|
||||
//!
|
||||
//! Profile under multi-worker load shows order-of-magnitude lower
|
||||
//! contention than the previous single-`Mutex<HashMap>` design — see
|
||||
//! `benches/ccr_store.rs`.
|
||||
//!
|
||||
//! # Pluggable backend
|
||||
//!
|
||||
//! Production deployments swap in their own [`CcrStore`] backed by Redis,
|
||||
|
|
@ -21,10 +33,12 @@
|
|||
//!
|
||||
//! [`CompressionStore`]: https://github.com/chopratejas/headroom/blob/main/headroom/cache/compression_store.py
|
||||
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
/// Pluggable CCR storage backend. `Send + Sync` so it can sit behind an
|
||||
/// `Arc` and be shared across threads in the proxy.
|
||||
pub trait CcrStore: Send + Sync {
|
||||
|
|
@ -50,29 +64,29 @@ pub const DEFAULT_CAPACITY: usize = 1000;
|
|||
/// Default TTL — 5 minutes, matching Python.
|
||||
pub const DEFAULT_TTL: Duration = Duration::from_secs(300);
|
||||
|
||||
/// Simple in-memory CCR store with TTL + bounded capacity.
|
||||
/// In-memory CCR store backed by [`DashMap`] for sharded concurrent
|
||||
/// access.
|
||||
///
|
||||
/// - **TTL**: 5 minutes by default. Entries past their TTL are dropped
|
||||
/// on the next `get`.
|
||||
/// - **Capacity**: 1000 entries by default. When full, the oldest
|
||||
/// insertion is evicted (FIFO).
|
||||
/// - **Locking**: single `Mutex`. The store sits on the cold path
|
||||
/// (one call per crushed array, not per token) so coarse locking is
|
||||
/// fine and keeps the implementation small.
|
||||
/// on the next `get` (lazy expiry — no background reaper thread).
|
||||
/// - **Capacity**: 1000 entries by default. When `put` would push us
|
||||
/// past capacity, the oldest entry (per insertion order) is evicted.
|
||||
/// - **Concurrency**: gets and puts on distinct keys do not contend.
|
||||
/// The only serialization point is the insertion-order queue used
|
||||
/// for capacity eviction; that mutex is held for an O(1) push or a
|
||||
/// small sweep.
|
||||
pub struct InMemoryCcrStore {
|
||||
inner: Mutex<Inner>,
|
||||
map: DashMap<String, Entry>,
|
||||
/// FIFO insertion order. Stale entries (already removed from `map`
|
||||
/// via TTL expiry) are tolerated — `pop_front` + `map.remove` is a
|
||||
/// no-op for missing keys, and capacity-bounded sweeps loop until
|
||||
/// they actually evict a real entry.
|
||||
order: Mutex<VecDeque<String>>,
|
||||
ttl: Duration,
|
||||
capacity: usize,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
map: HashMap<String, Entry>,
|
||||
/// FIFO order of insertion for capacity eviction. Hashes that get
|
||||
/// re-stored stay at their original position — same content under
|
||||
/// the same hash is idempotent and rare.
|
||||
order: VecDeque<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Entry {
|
||||
payload: String,
|
||||
inserted: Instant,
|
||||
|
|
@ -86,14 +100,28 @@ impl InMemoryCcrStore {
|
|||
|
||||
pub fn with_capacity_and_ttl(capacity: usize, ttl: Duration) -> Self {
|
||||
Self {
|
||||
inner: Mutex::new(Inner {
|
||||
map: HashMap::new(),
|
||||
order: VecDeque::new(),
|
||||
}),
|
||||
map: DashMap::with_capacity(capacity),
|
||||
order: Mutex::new(VecDeque::with_capacity(capacity)),
|
||||
ttl,
|
||||
capacity,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sweep the order queue, dropping leading entries that no longer
|
||||
/// exist in the map (already expired or evicted), then evict
|
||||
/// real entries until `map.len() < capacity`. Called only from
|
||||
/// `put` on a fresh-key insert path.
|
||||
fn evict_until_under_capacity(&self) {
|
||||
let mut guard = self.order.lock().expect("ccr order mutex poisoned");
|
||||
while self.map.len() >= self.capacity {
|
||||
let Some(oldest) = guard.pop_front() else {
|
||||
break;
|
||||
};
|
||||
// `remove` is a no-op if `oldest` was already lazy-expired.
|
||||
// Loop continues until we actually shrink the map.
|
||||
self.map.remove(&oldest);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for InMemoryCcrStore {
|
||||
|
|
@ -104,59 +132,58 @@ impl Default for InMemoryCcrStore {
|
|||
|
||||
impl CcrStore for InMemoryCcrStore {
|
||||
fn put(&self, hash: &str, payload: &str) {
|
||||
let mut g = self.inner.lock().expect("ccr store mutex poisoned");
|
||||
|
||||
if g.map.contains_key(hash) {
|
||||
// Idempotent re-store. Same hash should mean same content;
|
||||
// overwrite the payload (cheap) and keep the original FIFO
|
||||
// position so eviction stays predictable.
|
||||
g.map.insert(
|
||||
hash.to_string(),
|
||||
Entry {
|
||||
payload: payload.to_string(),
|
||||
inserted: Instant::now(),
|
||||
},
|
||||
);
|
||||
// Idempotent re-store fast-path: same hash → overwrite payload
|
||||
// in place, leave the order queue alone. Common when the same
|
||||
// tool output flows through multiple times in a session.
|
||||
if let Some(mut existing) = self.map.get_mut(hash) {
|
||||
existing.payload = payload.to_string();
|
||||
existing.inserted = Instant::now();
|
||||
return;
|
||||
}
|
||||
|
||||
// New entry. Evict the oldest if we're at capacity.
|
||||
while g.map.len() >= self.capacity {
|
||||
let Some(oldest) = g.order.pop_front() else {
|
||||
break;
|
||||
};
|
||||
g.map.remove(&oldest);
|
||||
// New entry. Cap-bound first (may sweep a few stale order
|
||||
// entries), then insert and append to the FIFO queue.
|
||||
if self.map.len() >= self.capacity {
|
||||
self.evict_until_under_capacity();
|
||||
}
|
||||
let entry = Entry {
|
||||
payload: payload.to_string(),
|
||||
inserted: Instant::now(),
|
||||
};
|
||||
let prev = self.map.insert(hash.to_string(), entry);
|
||||
if prev.is_none() {
|
||||
// Truly new key — record in FIFO order. (If `prev.is_some()`
|
||||
// it means another thread re-inserted between our get_mut
|
||||
// miss and this insert; treat that as a fast-path overwrite
|
||||
// and skip the queue append to avoid duplicates.)
|
||||
self.order
|
||||
.lock()
|
||||
.expect("ccr order mutex poisoned")
|
||||
.push_back(hash.to_string());
|
||||
}
|
||||
|
||||
g.map.insert(
|
||||
hash.to_string(),
|
||||
Entry {
|
||||
payload: payload.to_string(),
|
||||
inserted: Instant::now(),
|
||||
},
|
||||
);
|
||||
g.order.push_back(hash.to_string());
|
||||
}
|
||||
|
||||
fn get(&self, hash: &str) -> Option<String> {
|
||||
let mut g = self.inner.lock().expect("ccr store mutex poisoned");
|
||||
let expired = match g.map.get(hash) {
|
||||
Some(e) => e.inserted.elapsed() > self.ttl,
|
||||
None => return None,
|
||||
// Read path: shard read-lock, check TTL, clone payload out.
|
||||
// No global lock involvement at all — distinct hashes hash to
|
||||
// distinct shards and never contend.
|
||||
let expired_at = {
|
||||
let entry = self.map.get(hash)?;
|
||||
if entry.inserted.elapsed() > self.ttl {
|
||||
Some(()) // signal expired; drop guard before we remove
|
||||
} else {
|
||||
return Some(entry.payload.clone());
|
||||
}
|
||||
};
|
||||
if expired {
|
||||
g.map.remove(hash);
|
||||
if expired_at.is_some() {
|
||||
self.map.remove(hash);
|
||||
return None;
|
||||
}
|
||||
g.map.get(hash).map(|e| e.payload.clone())
|
||||
None
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("ccr store mutex poisoned")
|
||||
.map
|
||||
.len()
|
||||
self.map.len()
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -220,4 +247,38 @@ mod tests {
|
|||
assert_eq!(store.get("h"), Some("v".to_string()));
|
||||
assert!(!store.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn concurrent_puts_and_gets_do_not_corrupt() {
|
||||
// Smoke test for the concurrent design — N threads each do
|
||||
// P puts and P gets against distinct keys. Every key written
|
||||
// must be readable afterwards.
|
||||
use std::sync::Arc;
|
||||
use std::thread;
|
||||
|
||||
let store = Arc::new(InMemoryCcrStore::with_capacity_and_ttl(10_000, DEFAULT_TTL));
|
||||
let n_threads = 8;
|
||||
let per_thread = 200;
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for tid in 0..n_threads {
|
||||
let s = store.clone();
|
||||
handles.push(thread::spawn(move || {
|
||||
for i in 0..per_thread {
|
||||
let key = format!("t{tid}_k{i}");
|
||||
let val = format!("v{tid}_{i}");
|
||||
s.put(&key, &val);
|
||||
}
|
||||
for i in 0..per_thread {
|
||||
let key = format!("t{tid}_k{i}");
|
||||
let got = s.get(&key);
|
||||
assert_eq!(got, Some(format!("v{tid}_{i}")));
|
||||
}
|
||||
}));
|
||||
}
|
||||
for h in handles {
|
||||
h.join().unwrap();
|
||||
}
|
||||
assert_eq!(store.len(), n_threads * per_thread);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -667,11 +667,15 @@ impl SmartCrusher {
|
|||
// serves the original back via retrieval tool calls.
|
||||
let dropped_count = items.len().saturating_sub(result.len());
|
||||
let (ccr_hash, dropped_summary) = if dropped_count > 0 {
|
||||
let h = hash_array_for_ccr(items);
|
||||
// Serialize the original array exactly ONCE. The hash is
|
||||
// taken over those bytes, and (if a store is configured) the
|
||||
// same bytes get stored — eliminating a redundant tree clone
|
||||
// (`items.to_vec()`) and a redundant `serde_json::to_string`
|
||||
// pass that the previous version did per dropped array.
|
||||
let canonical = canonical_array_json(items);
|
||||
let h = hash_canonical(&canonical);
|
||||
let marker = format!("<<ccr:{h} {dropped_count}_rows_offloaded>>");
|
||||
if let Some(store) = &self.ccr_store {
|
||||
let canonical =
|
||||
serde_json::to_string(&Value::Array(items.to_vec())).unwrap_or_default();
|
||||
store.put(&h, &canonical);
|
||||
}
|
||||
(Some(h), marker)
|
||||
|
|
@ -898,14 +902,22 @@ fn estimate_array_bytes(item_strings: &[String]) -> usize {
|
|||
payload + separators + 2
|
||||
}
|
||||
|
||||
/// 12-char SHA-256 hex prefix of the canonical JSON serialization of
|
||||
/// `[v0, v1, ...]`. Used as the CCR retrieval key when the lossy path
|
||||
/// drops rows. Same input → same hash, so the runtime can cache the
|
||||
/// original by-hash and retrieval is deterministic.
|
||||
fn hash_array_for_ccr(items: &[Value]) -> String {
|
||||
/// Serialize `[v0, v1, ...]` once into the canonical JSON form used by
|
||||
/// the CCR retrieval contract. `serde_json` writes a slice of `Value` as
|
||||
/// the same bytes it would write for `Value::Array(items.to_vec())`, so
|
||||
/// we skip the array-wrapper allocation and the deep tree clone it
|
||||
/// requires. Used by both the hash (input) and the store payload (write).
|
||||
fn canonical_array_json(items: &[Value]) -> String {
|
||||
serde_json::to_string(items).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// 12-char SHA-256 hex prefix of an already-serialized canonical JSON
|
||||
/// string. Caller is responsible for producing the canonical form via
|
||||
/// [`canonical_array_json`] (or another byte-equal serializer) — the
|
||||
/// hash is over the bytes, so a stable serializer is the contract.
|
||||
fn hash_canonical(canonical: &str) -> String {
|
||||
use sha2::{Digest, Sha256};
|
||||
let mut h = Sha256::new();
|
||||
let canonical = serde_json::to_string(&Value::Array(items.to_vec())).unwrap_or_default();
|
||||
h.update(canonical.as_bytes());
|
||||
h.finalize()
|
||||
.iter()
|
||||
|
|
@ -914,6 +926,15 @@ fn hash_array_for_ccr(items: &[Value]) -> String {
|
|||
.collect()
|
||||
}
|
||||
|
||||
/// Convenience: canonical-serialize `items` and hash the result. Kept
|
||||
/// for sites (e.g. tests) that don't also need the canonical bytes for
|
||||
/// storage. Production lossy path inlines `canonical_array_json` +
|
||||
/// `hash_canonical` so the bytes are reused for the store payload.
|
||||
#[cfg(test)]
|
||||
fn hash_array_for_ccr(items: &[Value]) -> String {
|
||||
hash_canonical(&canonical_array_json(items))
|
||||
}
|
||||
|
||||
// ─── PR5 walker-integration helpers (string handling) ──────────────────────
|
||||
//
|
||||
// Parse-as-JSON-container, marker formatting, and humanize-bytes used to
|
||||
|
|
|
|||
|
|
@ -384,11 +384,18 @@ impl PyDiffCompressor {
|
|||
|
||||
/// `compress(content: str, context: str = "") -> DiffCompressionResult`.
|
||||
/// Argument order and keyword names match the Python implementation.
|
||||
///
|
||||
/// Releases the GIL across the Rust compress call so concurrent
|
||||
/// Python threads (uvicorn workers, asyncio tasks) can keep
|
||||
/// running while we hash + parse + filter the diff. The
|
||||
/// `&str` inputs are copied to owned `String`s first because
|
||||
/// PyO3 ties their lifetime to the GIL hold.
|
||||
#[pyo3(signature = (content, context = ""))]
|
||||
fn compress(&self, content: &str, context: &str) -> PyDiffCompressionResult {
|
||||
PyDiffCompressionResult {
|
||||
inner: self.inner.compress(content, context),
|
||||
}
|
||||
fn compress(&self, py: Python<'_>, content: &str, context: &str) -> PyDiffCompressionResult {
|
||||
let content = content.to_string();
|
||||
let context = context.to_string();
|
||||
let inner = py.allow_threads(|| self.inner.compress(&content, &context));
|
||||
PyDiffCompressionResult { inner }
|
||||
}
|
||||
|
||||
/// `compress_with_stats(content, context="") -> (result, stats)`.
|
||||
|
|
@ -398,10 +405,14 @@ impl PyDiffCompressor {
|
|||
#[pyo3(signature = (content, context = ""))]
|
||||
fn compress_with_stats(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
content: &str,
|
||||
context: &str,
|
||||
) -> (PyDiffCompressionResult, PyDiffCompressorStats) {
|
||||
let (result, stats) = self.inner.compress_with_stats(content, context);
|
||||
let content = content.to_string();
|
||||
let context = context.to_string();
|
||||
let (result, stats) =
|
||||
py.allow_threads(|| self.inner.compress_with_stats(&content, &context));
|
||||
(
|
||||
PyDiffCompressionResult { inner: result },
|
||||
PyDiffCompressorStats { inner: stats },
|
||||
|
|
@ -644,20 +655,36 @@ impl PySmartCrusher {
|
|||
|
||||
/// `crush(content, query="", bias=1.0) -> CrushResult`. Argument
|
||||
/// order and keyword names mirror the Python implementation.
|
||||
///
|
||||
/// Releases the GIL across the Rust crush call. Concurrent Python
|
||||
/// threads in the proxy keep running during the JSON parse +
|
||||
/// recursive process_value + per-array compression work. `&str`
|
||||
/// inputs are copied to owned `String`s up-front since PyO3 ties
|
||||
/// their lifetime to the GIL hold.
|
||||
#[pyo3(signature = (content, query = "", bias = 1.0))]
|
||||
fn crush(&self, content: &str, query: &str, bias: f64) -> PyCrushResult {
|
||||
PyCrushResult {
|
||||
inner: self.inner.crush(content, query, bias),
|
||||
}
|
||||
fn crush(&self, py: Python<'_>, content: &str, query: &str, bias: f64) -> PyCrushResult {
|
||||
let content = content.to_string();
|
||||
let query = query.to_string();
|
||||
let inner = py.allow_threads(|| self.inner.crush(&content, &query, bias));
|
||||
PyCrushResult { inner }
|
||||
}
|
||||
|
||||
/// `smart_crush_content(content, query="", bias=1.0) -> (str, bool, str)`.
|
||||
/// Mirrors Python's `_smart_crush_content` — used by
|
||||
/// `smart_crush_tool_output` convenience function and direct
|
||||
/// callers that want the tuple form.
|
||||
/// callers that want the tuple form. Releases the GIL across the
|
||||
/// compute (same rationale as `crush`).
|
||||
#[pyo3(signature = (content, query = "", bias = 1.0))]
|
||||
fn smart_crush_content(&self, content: &str, query: &str, bias: f64) -> (String, bool, String) {
|
||||
self.inner.smart_crush_content(content, query, bias)
|
||||
fn smart_crush_content(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
content: &str,
|
||||
query: &str,
|
||||
bias: f64,
|
||||
) -> (String, bool, String) {
|
||||
let content = content.to_string();
|
||||
let query = query.to_string();
|
||||
py.allow_threads(|| self.inner.smart_crush_content(&content, &query, bias))
|
||||
}
|
||||
|
||||
/// Crush a JSON array directly and return the structured result.
|
||||
|
|
@ -684,25 +711,39 @@ impl PySmartCrusher {
|
|||
query: &str,
|
||||
bias: f64,
|
||||
) -> Bound<'py, PyDict> {
|
||||
// Errors here surface as Python `RuntimeError` via pyo3's panic
|
||||
// catcher — callers are expected to pass valid array-shaped JSON.
|
||||
let parsed: serde_json::Value = serde_json::from_str(items_json)
|
||||
.unwrap_or_else(|e| panic!("items_json must be JSON: {e}"));
|
||||
let items = match parsed {
|
||||
serde_json::Value::Array(a) => a,
|
||||
other => panic!("items_json must be a JSON array, got {}", type_name(&other)),
|
||||
};
|
||||
let result = self.inner.crush_array(&items, query, bias);
|
||||
let kept_json = serde_json::to_string(&serde_json::Value::Array(result.items))
|
||||
.expect("serialize kept items");
|
||||
// GIL-release pattern: own the inputs, do all heavy compute
|
||||
// (JSON parse, crush, re-serialize) without the GIL, then
|
||||
// re-acquire to build the PyDict from the owned outputs.
|
||||
let items_json = items_json.to_string();
|
||||
let query = query.to_string();
|
||||
let (kept_json, ccr_hash, dropped_summary, strategy_info, compacted, compaction_kind) = py
|
||||
.allow_threads(|| {
|
||||
let parsed: serde_json::Value = serde_json::from_str(&items_json)
|
||||
.unwrap_or_else(|e| panic!("items_json must be JSON: {e}"));
|
||||
let items = match parsed {
|
||||
serde_json::Value::Array(a) => a,
|
||||
other => panic!("items_json must be a JSON array, got {}", type_name(&other)),
|
||||
};
|
||||
let result = self.inner.crush_array(&items, &query, bias);
|
||||
let kept_json = serde_json::to_string(&serde_json::Value::Array(result.items))
|
||||
.expect("serialize kept items");
|
||||
(
|
||||
kept_json,
|
||||
result.ccr_hash,
|
||||
result.dropped_summary,
|
||||
result.strategy_info,
|
||||
result.compacted,
|
||||
result.compaction_kind,
|
||||
)
|
||||
});
|
||||
build_crush_array_dict(
|
||||
py,
|
||||
kept_json,
|
||||
result.ccr_hash,
|
||||
result.dropped_summary,
|
||||
result.strategy_info,
|
||||
result.compacted,
|
||||
result.compaction_kind,
|
||||
ccr_hash,
|
||||
dropped_summary,
|
||||
strategy_info,
|
||||
compacted,
|
||||
compaction_kind,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -719,15 +760,20 @@ impl PySmartCrusher {
|
|||
/// pass without per-array lossy crushing — useful when the caller
|
||||
/// wants document-shape compaction (forms, configs, mixed records)
|
||||
/// rather than statistical row drop.
|
||||
fn compact_document_json(&self, doc_json: &str) -> String {
|
||||
let parsed: serde_json::Value =
|
||||
serde_json::from_str(doc_json).unwrap_or_else(|e| panic!("doc_json must be JSON: {e}"));
|
||||
let mut dc = DocumentCompactor::new();
|
||||
if let Some(store) = self.inner.ccr_store() {
|
||||
dc = dc.with_ccr_store(store.clone());
|
||||
}
|
||||
let out = dc.compact(parsed);
|
||||
serde_json::to_string(&out).expect("serialize compacted document")
|
||||
fn compact_document_json(&self, py: Python<'_>, doc_json: &str) -> String {
|
||||
// Heavy: JSON parse + recursive walker + tabular compaction +
|
||||
// re-serialize. None of it touches Python; release the GIL.
|
||||
let doc_json = doc_json.to_string();
|
||||
py.allow_threads(|| {
|
||||
let parsed: serde_json::Value = serde_json::from_str(&doc_json)
|
||||
.unwrap_or_else(|e| panic!("doc_json must be JSON: {e}"));
|
||||
let mut dc = DocumentCompactor::new();
|
||||
if let Some(store) = self.inner.ccr_store() {
|
||||
dc = dc.with_ccr_store(store.clone());
|
||||
}
|
||||
let out = dc.compact(parsed);
|
||||
serde_json::to_string(&out).expect("serialize compacted document")
|
||||
})
|
||||
}
|
||||
|
||||
/// Look up an original payload by CCR hash.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue