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

submilli_engine/runtime/
session_kv.rs
