1use std::collections::BTreeMap;
12use std::future::Future;
13use std::pin::Pin;
14use std::sync::{Mutex, PoisonError};
15
16pub const DEFAULT_MAX_SESSION_BYTES: u64 = 16 * 1024 * 1024;
18pub const DEFAULT_MAX_VALUE_BYTES: u64 = 1024 * 1024;
20pub const DEFAULT_MAX_ENTRIES: u64 = 1024;
21pub const DEFAULT_MAX_KEY_UNITS: u64 = 256;
23
24const BYTES_PER_UNIT: u64 = 2;
26
27const 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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
52pub struct SessionKvUsage {
53 pub entries: u64,
54 pub bytes: u64,
56 pub replaced: Option<u64>,
58}
59
60impl 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 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 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#[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 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 LimitExceeded {
192 operation: &'static str,
193 key: String,
194 limit: SessionKvLimitKind,
195 },
196 InvalidKey {
198 operation: &'static str,
199 key: String,
200 reason: &'static str,
201 },
202 SessionClosed { operation: &'static str },
205 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#[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#[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
258pub type SessionKvFuture<'a, T> =
261 Pin<Box<dyn Future<Output = Result<T, SessionKvError>> + Send + 'a>>;
262
263pub trait SessionKvStore: Send + Sync {
270 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 fn remove<'a>(&'a self, key: &'a [u16]) -> SessionKvFuture<'a, bool>;
280
281 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
295pub const CURSOR_KEY_ID_LEN: usize = 8;
297
298#[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#[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
358pub 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 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 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
495pub 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
500pub fn value_bytes(payload: &[u16]) -> u64 {
502 payload.len() as u64 * BYTES_PER_UNIT
503}
504
505pub 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 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 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}