Skip to main content

fast_cache/storage/
embedded_store_sharded.rs

1use std::cell::RefCell;
2use std::fmt;
3use std::marker::PhantomData;
4use std::ops::ControlFlow;
5use std::rc::Rc;
6
7use crate::cuda::{
8    CudaChunkTransferHit, CudaSessionChunkEvent, CudaSessionTransferRequest,
9    CudaSessionTransferStats,
10};
11use crate::storage::{
12    Bytes, EmbeddedKeyRoute, EmbeddedReadSlice, EmbeddedRouteMode, EmbeddedSessionRoute,
13    EmbeddedStore, OwnedEmbeddedBatchReadView, OwnedEmbeddedReadView,
14    OwnedEmbeddedSessionPackedView, OwnedEmbeddedWorkerShards, PackedBatch, PackedSessionWrite,
15    PreparedPointKey, TierStatsSnapshot,
16};
17
18thread_local! {
19    static THREAD_LOCAL_EMBEDDED_STORE: RefCell<Option<WorkerLocalEmbeddedStore>> = const { RefCell::new(None) };
20}
21
22/// Worker-owned view of the shards assigned to one local server worker.
23///
24/// This is the sharded counterpart to [`EmbeddedStore`]. An `EmbeddedStore` can
25/// be split into one `WorkerLocalEmbeddedStore` per worker. Each handle owns
26/// only its assigned shards and exposes local operations through `&mut self`,
27/// letting hot paths skip shared `RwLock` traffic when the key or session route
28/// belongs to this worker.
29///
30/// Methods ending in `_if_local` return [`LocalRouteError`] when a key/session
31/// routes to another worker. Methods ending in `_local` assume the caller has
32/// already checked routing and use debug assertions for that contract.
33#[derive(Debug)]
34pub struct WorkerLocalEmbeddedStore {
35    inner: OwnedEmbeddedWorkerShards,
36}
37
38/// Temporary container produced while partitioning an [`EmbeddedStore`] across workers.
39///
40/// `from_embedded` consumes the shared store and produces one
41/// [`WorkerLocalEmbeddedStore`] per requested worker. Call [`Self::into_stores`]
42/// to hand those stores to worker threads or runtimes.
43#[derive(Debug)]
44pub struct WorkerLocalEmbeddedStoreBootstrap {
45    stores: Vec<WorkerLocalEmbeddedStore>,
46}
47mod bootstrap;
48mod core;
49mod errors;
50mod lifecycle;
51mod read;
52#[cfg(test)]
53mod tests;
54mod thread_local;
55mod views;
56mod write;
57
58pub use errors::{LocalRouteError, LocalStoreAccessError, LocalStoreInstallError};
59pub use thread_local::{take_thread_local_embedded_store, with_thread_local_embedded_store};
60use views::worker_local_batch_view_from_embedded;
61pub use views::{
62    WorkerLocalBatchReadView, WorkerLocalReadSlice, WorkerLocalReadView,
63    WorkerLocalSessionBatchView,
64};