Skip to main content

submilli_engine/runtime/
session_kv.rs

1//! Session-scoped key-value store: the trait an embedder implements and a
2//! bounded in-memory implementation.
3//!
4//! Keys and payloads are UTF-16 code units end to end. A payload can contain
5//! lone surrogates — a value a Submilli `string` may legally hold — so nothing
6//! here round-trips through a Rust `String`, which would replace them.
7//!
8//! The store outlives the Wasm store it is read from, so it never retains guest
9//! GC references: serialization to code units happens above this layer.
10
11use std::collections::BTreeMap;
12use std::future::Future;
13use std::pin::Pin;
14use std::sync::{Mutex, PoisonError};
15
16/// Retained serialized payload per session.
17pub const DEFAULT_MAX_SESSION_BYTES: u64 = 16 * 1024 * 1024;
18/// Serialized payload of a single value.
19pub const DEFAULT_MAX_VALUE_BYTES: u64 = 1024 * 1024;
20pub const DEFAULT_MAX_ENTRIES: u64 = 1024;
21/// UTF-16 code units in a nonempty key.
22pub const DEFAULT_MAX_KEY_UNITS: u64 = 256;
23
24/// Bytes charged per stored UTF-16 code unit, for both keys and payloads.
25const BYTES_PER_UNIT: u64 = 2;
26
27/// Per-key implementation overhead, bounded separately from payload bytes so a
28/// session holding many tiny values cannot outgrow its map accounting.
29const ENTRY_OVERHEAD_BYTES: u64 = 64;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct SessionKvLimits {
33    pub max_session_bytes: u64,
34    pub max_value_bytes: u64,
35    pub max_entries: u64,
36    pub max_key_units: u64,
37}
38
39impl Default for SessionKvLimits {
40    fn default() -> Self {
41        Self {
42            max_session_bytes: DEFAULT_MAX_SESSION_BYTES,
43            max_value_bytes: DEFAULT_MAX_VALUE_BYTES,
44            max_entries: DEFAULT_MAX_ENTRIES,
45            max_key_units: DEFAULT_MAX_KEY_UNITS,
46        }
47    }
48}
49
50/// What a session holds when a write is checked against its limits.
51#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
52pub struct SessionKvUsage {
53    pub entries: u64,
54    /// Retained bytes, as [`entry_bytes`] counts them.
55    pub bytes: u64,
56    /// The retained size of the entry the write replaces, `None` for a new key.
57    pub replaced: Option<u64>,
58}
59
60/// The checks every store applies, in one place so a durable store refuses
61/// exactly what the in-memory one does.
62impl SessionKvLimits {
63    pub fn check_key(&self, operation: &'static str, key: &[u16]) -> Result<(), SessionKvError> {
64        if key.is_empty() {
65            return Err(SessionKvError::InvalidKey {
66                operation,
67                key: String::new(),
68                reason: "key must not be empty",
69            });
70        }
71        let units = key.len() as u64;
72        if units > self.max_key_units {
73            return Err(SessionKvError::LimitExceeded {
74                operation,
75                key: render_key(key),
76                limit: SessionKvLimitKind::KeyUnits {
77                    units,
78                    limit: self.max_key_units,
79                },
80            });
81        }
82        Ok(())
83    }
84
85    /// The key and value checks `set` makes before looking at what the session
86    /// holds.
87    pub fn check_write(&self, key: &[u16], payload: &[u16]) -> Result<(), SessionKvError> {
88        self.check_key("set", key)?;
89        let bytes = value_bytes(payload);
90        if bytes > self.max_value_bytes {
91            return Err(SessionKvError::LimitExceeded {
92                operation: "set",
93                key: render_key(key),
94                limit: SessionKvLimitKind::ValueBytes {
95                    bytes,
96                    limit: self.max_value_bytes,
97                },
98            });
99        }
100        Ok(())
101    }
102
103    /// Whether `set` fits beside what the session already holds. The count
104    /// bound only gates a new key. Returns the session's retained bytes after
105    /// the write.
106    pub fn check_capacity(
107        &self,
108        key: &[u16],
109        payload: &[u16],
110        held: SessionKvUsage,
111    ) -> Result<u64, SessionKvError> {
112        if held.replaced.is_none() && held.entries >= self.max_entries {
113            return Err(SessionKvError::LimitExceeded {
114                operation: "set",
115                key: render_key(key),
116                limit: SessionKvLimitKind::EntryCount {
117                    limit: self.max_entries,
118                },
119            });
120        }
121        let next = held
122            .bytes
123            .saturating_sub(held.replaced.unwrap_or(0))
124            .saturating_add(entry_bytes(key, payload));
125        if next > self.max_session_bytes {
126            return Err(SessionKvError::LimitExceeded {
127                operation: "set",
128                key: render_key(key),
129                limit: SessionKvLimitKind::PerSessionBytes {
130                    requested: next,
131                    limit: self.max_session_bytes,
132                },
133            });
134        }
135        Ok(next)
136    }
137}
138
139/// Which bound a rejected operation ran into. Carries the numbers, never the
140/// value — an error message about a quota must not leak session contents.
141#[derive(Debug, Clone, Copy, PartialEq, Eq)]
142pub enum SessionKvLimitKind {
143    KeyUnits {
144        units: u64,
145        limit: u64,
146    },
147    ValueBytes {
148        bytes: u64,
149        limit: u64,
150    },
151    /// This session's own retained bytes.
152    PerSessionBytes {
153        requested: u64,
154        limit: u64,
155    },
156    EntryCount {
157        limit: u64,
158    },
159}
160
161impl std::fmt::Display for SessionKvLimitKind {
162    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
163        match self {
164            Self::KeyUnits { units, limit } => write!(
165                f,
166                "key length limit: {units} UTF-16 code units exceeds the {limit} allowed"
167            ),
168            Self::ValueBytes { bytes, limit } => write!(
169                f,
170                "value size limit: {bytes} serialized bytes exceeds the {limit} allowed"
171            ),
172            Self::PerSessionBytes { requested, limit } => write!(
173                f,
174                "session payload limit: {requested} retained bytes exceeds the {limit} this \
175                 session may hold — remove entries this session no longer needs, or store \
176                 less per key"
177            ),
178            Self::EntryCount { limit } => write!(
179                f,
180                "entry count limit: the session already holds {limit} — remove entries this \
181                 session no longer needs"
182            ),
183        }
184    }
185}
186
187#[derive(Debug, Clone, PartialEq, Eq)]
188pub enum SessionKvError {
189    /// A bound was exceeded. `key` is rendered lossily for the message only;
190    /// the value never appears.
191    LimitExceeded {
192        operation: &'static str,
193        key: String,
194        limit: SessionKvLimitKind,
195    },
196    /// An empty key, or one the implementation otherwise refuses.
197    InvalidKey {
198        operation: &'static str,
199        key: String,
200        reason: &'static str,
201    },
202    /// The session ended while the run still held its store. Nothing was read
203    /// or written: an ended session's state is gone and cannot be refilled.
204    SessionClosed { operation: &'static str },
205    /// The backing store failed. The operation may not have happened, so the
206    /// run cannot carry on as though it did.
207    Backend {
208        operation: &'static str,
209        message: String,
210    },
211}
212
213impl std::fmt::Display for SessionKvError {
214    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
215        match self {
216            Self::LimitExceeded {
217                operation,
218                key,
219                limit,
220            } => write!(f, "session.{operation}(\"{key}\") exceeded the {limit}"),
221            Self::InvalidKey {
222                operation,
223                key,
224                reason,
225            } => write!(f, "session.{operation}(\"{key}\"): {reason}"),
226            Self::SessionClosed { operation } => write!(
227                f,
228                "session.{operation} failed: the session was closed while this run was using it"
229            ),
230            Self::Backend { operation, message } => {
231                write!(f, "session.{operation} failed: {message}")
232            }
233        }
234    }
235}
236
237impl std::error::Error for SessionKvError {}
238
239/// Key plus retained payload size — what `list` may disclose. Never contents.
240#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
241pub struct SessionKvEntry {
242    pub key: Vec<u16>,
243    pub size_bytes: u64,
244}
245
246/// One page of a key scan.
247///
248/// `scanned_all` is false when the scan stopped on its bound rather than on
249/// exhausting the keyspace, so a paginating caller can tell "no more matching
250/// entries" from "look again past `last_scanned`".
251#[derive(Debug, Clone, PartialEq, Eq)]
252pub struct SessionKvPage {
253    pub entries: Vec<SessionKvEntry>,
254    pub last_scanned: Option<Vec<u16>>,
255    pub scanned_all: bool,
256}
257
258/// What a [`SessionKvStore`] operation returns: a store may reach storage that
259/// must not block the caller's thread.
260pub type SessionKvFuture<'a, T> =
261    Pin<Box<dyn Future<Output = Result<T, SessionKvError>> + Send + 'a>>;
262
263/// Session-scoped key-value storage.
264///
265/// Every operation is atomic and commits only after validation and capacity
266/// checks succeed, so a rejected write leaves the previous entry intact.
267/// Operations on one session serialize: two writers cannot both spend the same
268/// headroom, and concurrent writes to one key resolve last-committed-write.
269pub trait SessionKvStore: Send + Sync {
270    /// `Ok(None)` is an absent key. A present entry holding a serialized
271    /// `null` returns `Ok(Some(..))` — callers distinguish the two.
272    fn get<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, Option<Vec<u16>>>;
273
274    fn has<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, bool>;
275
276    fn set<'a>(&'a self, key: &'a [u16], payload: &'a [u16]) -> SessionKvFuture<'a, ()>;
277
278    /// `Ok(true)` when an entry existed and was removed.
279    fn remove<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, bool>;
280
281    /// Keys in UTF-16 code-unit order, strictly after `after` when given and
282    /// starting with `prefix`. Scans at most `max_scan` keys; matching is exact
283    /// code units, with no normalization or path semantics. One page is one
284    /// consistent read.
285    fn scan<'a>(
286        &'a self,
287        after: Option<&'a [u16]>,
288        prefix: &'a [u16],
289        max_scan: usize,
290    ) -> SessionKvFuture<'a, SessionKvPage>;
291
292    fn limits(&self) -> SessionKvLimits;
293}
294
295/// Bytes in a [`CursorKey`] id.
296pub const CURSOR_KEY_ID_LEN: usize = 8;
297
298/// A secret that seals `session.list` cursors, named by an id each cursor
299/// carries so a server holding several keys knows which one sealed it.
300#[derive(Clone)]
301pub struct CursorKey {
302    id: [u8; CURSOR_KEY_ID_LEN],
303    secret: [u8; 32],
304}
305
306impl CursorKey {
307    pub fn new(id: [u8; CURSOR_KEY_ID_LEN], secret: [u8; 32]) -> Self {
308        Self { id, secret }
309    }
310
311    pub fn id(&self) -> &[u8; CURSOR_KEY_ID_LEN] {
312        &self.id
313    }
314
315    pub(crate) fn secret(&self) -> &[u8; 32] {
316        &self.secret
317    }
318}
319
320impl std::fmt::Debug for CursorKey {
321    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
322        f.debug_struct("CursorKey")
323            .field("id", &self.id)
324            .finish_non_exhaustive()
325    }
326}
327
328/// The keys `session.list` cursors are sealed and opened with. New cursors are
329/// sealed with `current`; a cursor sealed with a `retired` key still opens, so
330/// replacing the key does not break listings already under way.
331///
332/// An embedder that keeps state across restarts installs keys that survive
333/// them too, or every cursor it handed out stops working when it restarts. With
334/// none installed, the runtime uses a random key of its own for the life of the
335/// process.
336#[derive(Debug, Clone)]
337pub struct CursorKeys {
338    current: CursorKey,
339    retired: Vec<CursorKey>,
340}
341
342impl CursorKeys {
343    pub fn new(current: CursorKey, retired: Vec<CursorKey>) -> Self {
344        Self { current, retired }
345    }
346
347    pub fn current(&self) -> &CursorKey {
348        &self.current
349    }
350
351    pub fn find(&self, id: &[u8]) -> Option<&CursorKey> {
352        std::iter::once(&self.current)
353            .chain(&self.retired)
354            .find(|key| key.id == id)
355    }
356}
357
358/// Bounded in-memory store, for embedders without durable storage and for
359/// state that must not outlive its owner.
360pub struct InMemorySessionKv {
361    limits: SessionKvLimits,
362    state: Mutex<KvState>,
363}
364
365#[derive(Default)]
366struct KvState {
367    entries: BTreeMap<Vec<u16>, Vec<u16>>,
368    used_bytes: u64,
369}
370
371impl Default for InMemorySessionKv {
372    fn default() -> Self {
373        Self::new(SessionKvLimits::default())
374    }
375}
376
377impl InMemorySessionKv {
378    pub fn new(limits: SessionKvLimits) -> Self {
379        Self {
380            limits,
381            state: Mutex::new(KvState::default()),
382        }
383    }
384
385    /// Retained bytes: payloads plus keys plus per-entry overhead.
386    pub fn used_bytes(&self) -> u64 {
387        self.locked().map_or(0, |s| s.used_bytes)
388    }
389
390    pub fn len(&self) -> usize {
391        self.locked().map_or(0, |s| s.entries.len())
392    }
393
394    pub fn is_empty(&self) -> bool {
395        self.len() == 0
396    }
397
398    fn locked(&self) -> Result<std::sync::MutexGuard<'_, KvState>, SessionKvError> {
399        self.state.lock().map_err(|e: PoisonError<_>| {
400            let _ = e;
401            SessionKvError::Backend {
402                operation: "state",
403                message: "session store lock poisoned".into(),
404            }
405        })
406    }
407
408    pub fn get_now(&self, key: &[u16]) -> Result<Option<Vec<u16>>, SessionKvError> {
409        self.limits.check_key("get", key)?;
410        Ok(self.locked()?.entries.get(key).cloned())
411    }
412
413    pub fn has_now(&self, key: &[u16]) -> Result<bool, SessionKvError> {
414        self.limits.check_key("has", key)?;
415        Ok(self.locked()?.entries.contains_key(key))
416    }
417
418    pub fn set_now(&self, key: &[u16], payload: &[u16]) -> Result<(), SessionKvError> {
419        self.limits.check_write(key, payload)?;
420        let mut state = self.locked()?;
421        let held = SessionKvUsage {
422            entries: state.entries.len() as u64,
423            bytes: state.used_bytes,
424            replaced: state
425                .entries
426                .get(key)
427                .map(|existing| entry_bytes(key, existing)),
428        };
429        let next_used = self.limits.check_capacity(key, payload, held)?;
430        state.entries.insert(key.to_vec(), payload.to_vec());
431        state.used_bytes = next_used;
432        Ok(())
433    }
434
435    pub fn remove_now(&self, key: &[u16]) -> Result<bool, SessionKvError> {
436        self.limits.check_key("remove", key)?;
437        let mut state = self.locked()?;
438        let Some(previous) = state.entries.remove(key) else {
439            return Ok(false);
440        };
441        let released = entry_bytes(key, &previous);
442        state.used_bytes = state.used_bytes.saturating_sub(released);
443        Ok(true)
444    }
445
446    pub fn scan_now(
447        &self,
448        after: Option<&[u16]>,
449        prefix: &[u16],
450        max_scan: usize,
451    ) -> Result<SessionKvPage, SessionKvError> {
452        let state = self.locked()?;
453        let lower = match after {
454            Some(key) => std::ops::Bound::Excluded(key.to_vec()),
455            None => std::ops::Bound::Unbounded,
456        };
457
458        let candidates = state
459            .entries
460            .range((lower, std::ops::Bound::Unbounded))
461            .take(max_scan);
462
463        let mut page = SessionKvPage {
464            entries: Vec::new(),
465            last_scanned: None,
466            scanned_all: true,
467        };
468        for (key, payload) in candidates {
469            page.last_scanned = Some(key.clone());
470            if key.starts_with(prefix) {
471                page.entries.push(SessionKvEntry {
472                    key: key.clone(),
473                    size_bytes: value_bytes(payload),
474                });
475            }
476        }
477        // `take` ended the walk either because the keyspace ran out or because the
478        // bound did; only the second leaves unvisited keys behind `last_scanned`.
479        if let Some(last) = &page.last_scanned
480            && state
481                .entries
482                .range((
483                    std::ops::Bound::Excluded(last.clone()),
484                    std::ops::Bound::Unbounded,
485                ))
486                .next()
487                .is_some()
488        {
489            page.scanned_all = false;
490        }
491        Ok(page)
492    }
493}
494
495/// Retained bytes for one entry: what the per-session limit counts.
496pub fn entry_bytes(key: &[u16], payload: &[u16]) -> u64 {
497    (key.len() as u64 + payload.len() as u64) * BYTES_PER_UNIT + ENTRY_OVERHEAD_BYTES
498}
499
500/// Serialized bytes of a value: what the value limit and `list` count.
501pub fn value_bytes(payload: &[u16]) -> u64 {
502    payload.len() as u64 * BYTES_PER_UNIT
503}
504
505/// Only ever used to name the key in a diagnostic. Lone surrogates become
506/// replacement characters here, which is acceptable in a message and is why the
507/// stored key itself is kept as code units.
508pub fn render_key(key: &[u16]) -> String {
509    String::from_utf16_lossy(key)
510}
511
512impl SessionKvStore for InMemorySessionKv {
513    fn get<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, Option<Vec<u16>>> {
514        Box::pin(std::future::ready(self.get_now(key)))
515    }
516
517    fn has<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, bool> {
518        Box::pin(std::future::ready(self.has_now(key)))
519    }
520
521    fn set<'a>(&'a self, key: &'a [u16], payload: &'a [u16]) -> SessionKvFuture<'a, ()> {
522        Box::pin(std::future::ready(self.set_now(key, payload)))
523    }
524
525    fn remove<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, bool> {
526        Box::pin(std::future::ready(self.remove_now(key)))
527    }
528
529    fn scan<'a>(
530        &'a self,
531        after: Option<&'a [u16]>,
532        prefix: &'a [u16],
533        max_scan: usize,
534    ) -> SessionKvFuture<'a, SessionKvPage> {
535        Box::pin(std::future::ready(self.scan_now(after, prefix, max_scan)))
536    }
537
538    fn limits(&self) -> SessionKvLimits {
539        self.limits
540    }
541}
542
543#[cfg(test)]
544mod tests {
545    use super::*;
546    use std::sync::Arc;
547
548    fn u(s: &str) -> Vec<u16> {
549        s.encode_utf16().collect()
550    }
551
552    fn store() -> InMemorySessionKv {
553        InMemorySessionKv::default()
554    }
555
556    #[test]
557    fn set_then_get_returns_the_value_and_missing_keys_are_none() {
558        let kv = store();
559        kv.set_now(&u("a"), &u("{\"x\":1}")).expect("set");
560        assert_eq!(kv.get_now(&u("a")).expect("get"), Some(u("{\"x\":1}")));
561        assert_eq!(kv.get_now(&u("b")).expect("get"), None);
562    }
563
564    #[test]
565    fn stored_null_is_distinguishable_from_a_missing_key() {
566        let kv = store();
567        kv.set_now(&u("n"), &u("null")).expect("set");
568        assert_eq!(kv.get_now(&u("n")).expect("get"), Some(u("null")));
569        assert!(kv.has_now(&u("n")).expect("has"));
570        assert!(!kv.has_now(&u("missing")).expect("has"));
571        assert_eq!(kv.get_now(&u("missing")).expect("get"), None);
572    }
573
574    #[test]
575    fn empty_payload_is_a_present_value() {
576        let kv = store();
577        kv.set_now(&u("e"), &[]).expect("set");
578        assert!(kv.has_now(&u("e")).expect("has"));
579        assert_eq!(kv.get_now(&u("e")).expect("get"), Some(Vec::new()));
580    }
581
582    #[test]
583    fn replacement_charges_only_the_size_delta() {
584        let kv = store();
585        kv.set_now(&u("k"), &u("aaaa")).expect("set");
586        let after_first = kv.used_bytes();
587        kv.set_now(&u("k"), &u("aaaaaa")).expect("replace");
588        assert_eq!(kv.len(), 1);
589        assert_eq!(kv.used_bytes(), after_first + 4);
590    }
591
592    #[test]
593    fn remove_reports_existence_and_releases_capacity() {
594        let kv = store();
595        kv.set_now(&u("k"), &u("value")).expect("set");
596        let held = kv.used_bytes();
597        assert!(held > 0);
598        assert!(kv.remove_now(&u("k")).expect("remove"));
599        assert_eq!(kv.used_bytes(), 0);
600        assert!(!kv.remove_now(&u("k")).expect("remove"));
601    }
602
603    #[test]
604    fn oversized_value_leaves_the_previous_entry_intact() {
605        let kv = store();
606        kv.set_now(&u("k"), &u("kept")).expect("set");
607        let before = kv.used_bytes();
608        let huge = vec![b'x' as u16; (DEFAULT_MAX_VALUE_BYTES as usize / 2) + 1];
609        let err = kv.set_now(&u("k"), &huge).expect_err("must reject");
610        assert!(matches!(
611            err,
612            SessionKvError::LimitExceeded {
613                operation: "set",
614                limit: SessionKvLimitKind::ValueBytes { .. },
615                ..
616            }
617        ));
618        assert_eq!(kv.get_now(&u("k")).expect("get"), Some(u("kept")));
619        assert_eq!(kv.used_bytes(), before);
620    }
621
622    #[test]
623    fn entry_count_limit_rejects_without_partial_mutation() {
624        let kv = InMemorySessionKv::new(SessionKvLimits {
625            max_entries: 4,
626            ..SessionKvLimits::default()
627        });
628        for i in 0..4 {
629            kv.set_now(&u(&format!("k{i}")), &u("v")).expect("set");
630        }
631        let before = kv.used_bytes();
632        let err = kv
633            .set_now(&u("overflow"), &u("v"))
634            .expect_err("must reject");
635        assert!(matches!(
636            err,
637            SessionKvError::LimitExceeded {
638                limit: SessionKvLimitKind::EntryCount { limit: 4 },
639                ..
640            }
641        ));
642        assert_eq!(
643            err.to_string(),
644            "session.set(\"overflow\") exceeded the entry count limit: the session already \
645             holds 4 — remove entries this session no longer needs"
646        );
647        assert_eq!(kv.len(), 4);
648        assert_eq!(kv.used_bytes(), before);
649        assert!(!kv.has_now(&u("overflow")).expect("has"));
650        // An existing key still writes: the count bound only gates insertion.
651        kv.set_now(&u("k0"), &u("vv")).expect("replace");
652    }
653
654    #[test]
655    fn session_payload_limit_names_operation_key_and_limit_but_not_the_value() {
656        let kv = InMemorySessionKv::new(SessionKvLimits {
657            max_session_bytes: 512,
658            ..SessionKvLimits::default()
659        });
660        let filler = vec![b'y' as u16; 300];
661        let err = kv.set_now(&u("bigkey"), &filler).expect_err("must reject");
662        let rendered = err.to_string();
663        assert!(rendered.contains("set"), "{rendered}");
664        assert!(rendered.contains("bigkey"), "{rendered}");
665        assert!(rendered.contains("512"), "{rendered}");
666        assert!(!rendered.contains("yyyy"), "{rendered}");
667        assert_eq!(kv.len(), 0);
668
669        kv.set_now(&u("held"), &vec![b'z' as u16; 200])
670            .expect("first write fits");
671        let secret = u("SUPER_SECRET_PAYLOAD");
672        let err = kv.set_now(&u("other"), &secret).expect_err("must reject");
673        assert!(!err.to_string().contains("SUPER_SECRET_PAYLOAD"));
674        assert_eq!(kv.len(), 1);
675    }
676
677    #[test]
678    fn key_at_the_unit_limit_is_accepted_and_one_over_is_rejected() {
679        let kv = store();
680        let at_limit = vec![b'k' as u16; DEFAULT_MAX_KEY_UNITS as usize];
681        kv.set_now(&at_limit, &u("v")).expect("256 units accepted");
682
683        let over = vec![b'k' as u16; DEFAULT_MAX_KEY_UNITS as usize + 1];
684        let err = kv.set_now(&over, &u("v")).expect_err("257 units rejected");
685        assert!(matches!(
686            err,
687            SessionKvError::LimitExceeded {
688                limit: SessionKvLimitKind::KeyUnits {
689                    units: 257,
690                    limit: 256
691                },
692                ..
693            }
694        ));
695        assert_eq!(kv.len(), 1);
696    }
697
698    #[test]
699    fn empty_key_is_rejected() {
700        let kv = store();
701        assert!(matches!(
702            kv.set_now(&[], &u("v")),
703            Err(SessionKvError::InvalidKey { .. })
704        ));
705    }
706
707    #[test]
708    fn lone_surrogates_survive_a_round_trip() {
709        let kv = store();
710        let key = vec![0x0061, 0xD800, 0x0062];
711        let payload = vec![0x0022, 0xDC00, 0xD83D, 0x0022];
712        kv.set_now(&key, &payload).expect("set");
713        assert_eq!(kv.get_now(&key).expect("get"), Some(payload.clone()));
714        // A Rust String round trip would have replaced these.
715        assert_ne!(u(&String::from_utf16_lossy(&payload)), payload);
716    }
717
718    #[test]
719    fn scan_walks_keys_in_code_unit_order_and_reports_its_bound() {
720        let kv = store();
721        for key in ["a", "ab", "b", "ba", "c"] {
722            kv.set_now(&u(key), &u("v")).expect("set");
723        }
724        let all = kv.scan_now(None, &[], 10).expect("scan");
725        assert!(all.scanned_all);
726        let keys: Vec<Vec<u16>> = all.entries.iter().map(|e| e.key.clone()).collect();
727        assert_eq!(keys, vec![u("a"), u("ab"), u("b"), u("ba"), u("c")]);
728        assert_eq!(all.entries[0].size_bytes, 2);
729
730        let bounded = kv.scan_now(None, &u("z"), 2).expect("scan");
731        assert!(bounded.entries.is_empty());
732        assert!(!bounded.scanned_all);
733        assert_eq!(bounded.last_scanned, Some(u("ab")));
734
735        let after = kv.scan_now(Some(&u("ab")), &u("b"), 10).expect("scan");
736        assert!(after.scanned_all);
737        let keys: Vec<Vec<u16>> = after.entries.iter().map(|e| e.key.clone()).collect();
738        assert_eq!(keys, vec![u("b"), u("ba")]);
739    }
740
741    #[test]
742    fn concurrent_writes_to_one_key_resolve_last_committed_write() {
743        let kv = Arc::new(store());
744        let writers: Vec<_> = (0..8)
745            .map(|i| {
746                let kv = Arc::clone(&kv);
747                std::thread::spawn(move || {
748                    for _ in 0..200 {
749                        kv.set_now(&u("hot"), &u(&format!("writer-{i}")))
750                            .expect("set");
751                    }
752                })
753            })
754            .collect();
755        for w in writers {
756            w.join().expect("writer thread");
757        }
758        let final_value = kv.get_now(&u("hot")).expect("get").expect("present");
759        let rendered = String::from_utf16_lossy(&final_value);
760        assert!(rendered.starts_with("writer-"), "{rendered}");
761        assert_eq!(kv.len(), 1);
762        assert_eq!(kv.used_bytes(), entry_bytes(&u("hot"), &final_value));
763    }
764}