Skip to main content

kevy_embedded/
store_inner.rs

1//! The [`Store`]'s shared internals — per-shard [`Inner`], the
2//! last-clone [`DropGuard`], and the [`WeakStore`] handle (split out
3//! of `store.rs` to keep it under the 500-LOC project ceiling;
4//! behaviour unchanged).
5
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::{Arc, Mutex, RwLock, Weak};
8use std::thread::JoinHandle;
9
10use kevy_persist::Aof;
11
12use crate::config::Config;
13use crate::pubsub::PubsubBus;
14use crate::store::{Shards, Store, lock_write};
15
16/// Weak handle to a `Store` — does not keep the underlying keyspace alive.
17///
18/// Used by the URL-keyed registry in `kevy-client` so that multiple
19/// `Connection::open("mem://name")` calls share the same backing store
20/// without leaking it when all strong handles go away.
21#[derive(Clone)]
22pub struct WeakStore {
23    shards: Weak<Vec<Arc<RwLock<Inner>>>>,
24    guard: Weak<DropGuard>,
25    config: Config,
26    #[cfg(not(target_arch = "wasm32"))]
27    feed_weak: Option<std::sync::Weak<Mutex<kevy_replicate::feed::FeedSource>>>,
28    blocker_weak: Weak<crate::ops_blocking::Blocker>,
29    indexes_weak: Weak<crate::ops_index::IndexReg>,
30    views_weak: Weak<crate::ops_view::ViewReg>,
31}
32
33impl WeakStore {
34    /// Try to upgrade back to a `Store`. Returns `None` if the last strong
35    /// reference has already been dropped.
36    pub fn upgrade(&self) -> Option<Store> {
37        Some(Store {
38            shards: self.shards.upgrade()?,
39            guard: self.guard.upgrade()?,
40            config: self.config.clone(),
41            #[cfg(not(target_arch = "wasm32"))]
42            feed: self.feed_weak.as_ref().and_then(std::sync::Weak::upgrade),
43            blocker: self.blocker_weak.upgrade()?,
44            indexes: self.indexes_weak.upgrade()?,
45            views: self.views_weak.upgrade()?,
46        })
47    }
48}
49
50impl Store {
51    /// Get a weak handle that does not keep the keyspace alive.
52    pub fn downgrade(&self) -> WeakStore {
53        WeakStore {
54            shards: Arc::downgrade(&self.shards),
55            guard: Arc::downgrade(&self.guard),
56            config: self.config.clone(),
57            #[cfg(not(target_arch = "wasm32"))]
58            feed_weak: self.feed.as_ref().map(Arc::downgrade),
59            blocker_weak: Arc::downgrade(&self.blocker),
60            indexes_weak: Arc::downgrade(&self.indexes),
61            views_weak: Arc::downgrade(&self.views),
62        }
63    }
64}
65
66pub(crate) struct Inner {
67    pub(crate) store: kevy_store::Store,
68    pub(crate) aof: Option<Aof>,
69    /// Pub/sub bus. Only shard 0's is ever used (pub/sub is process-wide);
70    /// other shards carry an idle one (cheap).
71    pub(crate) bus: PubsubBus,
72    /// Shared replication source if this store is an embed-as-writer.
73    /// Every shard holds a clone of the same `Arc<Mutex<...>>` so
74    /// `commit_write` can push mutations without reaching back up
75    /// through the `DropGuard`.
76    #[cfg(not(target_arch = "wasm32"))]
77    pub(crate) writer_source:
78        Option<std::sync::Arc<Mutex<kevy_replicate::source::ReplicationSource>>>,
79    /// v2.3 CDC feed (one stream per store); every shard holds a clone
80    /// so `commit_write` pushes effects inline. `None` = feed off.
81    #[cfg(not(target_arch = "wasm32"))]
82    pub(crate) feed: Option<std::sync::Arc<Mutex<kevy_replicate::feed::FeedSource>>>,
83    /// v2.4 blocking-pop wake channel clone (see `Store::blocker`).
84    pub(crate) blocker: Option<Arc<crate::ops_blocking::Blocker>>,
85    /// v2.5: this shard's index segments + the store-level registry
86    /// handle (for the commit_write hook).
87    pub(crate) idx_segs: crate::ops_index::ShardSegs,
88    pub(crate) idx_reg: Option<Arc<crate::ops_index::IndexReg>>,
89    /// v2.6: this shard's view states + registry handle.
90    pub(crate) view_segs: crate::ops_view::ShardViews,
91    pub(crate) view_reg: Option<Arc<crate::ops_view::ViewReg>>,
92}
93
94impl Inner {
95    pub(crate) fn new(store: kevy_store::Store, aof: Option<Aof>) -> Self {
96        Inner {
97            store,
98            aof,
99            bus: PubsubBus::new(),
100            #[cfg(not(target_arch = "wasm32"))]
101            writer_source: None,
102            #[cfg(not(target_arch = "wasm32"))]
103            feed: None,
104            blocker: None,
105            idx_segs: crate::ops_index::ShardSegs::default(),
106            idx_reg: None,
107            view_segs: crate::ops_view::ShardViews::default(),
108            view_reg: None,
109        }
110    }
111}
112
113/// Owns the reaper-thread handle + the shards for the final AOF flush. Lives
114/// in an `Arc<DropGuard>` shared across every `Store` clone; the drop logic
115/// fires only when the last clone goes away.
116pub(crate) struct DropGuard {
117    pub(crate) reaper_stop: Option<Arc<AtomicBool>>,
118    pub(crate) reaper_join: Mutex<Option<JoinHandle<()>>>,
119    pub(crate) shards_for_flush: Shards,
120    /// Replica runner thread + reconnect machinery, present iff this
121    /// store was opened with `Config::replica_upstream = Some(...)`.
122    /// Joined here so the runner stops cleanly when the last `Store`
123    /// clone goes away.
124    #[cfg(not(target_arch = "wasm32"))]
125    pub(crate) replica_runner: Option<crate::replica_runner::ReplicaRunner>,
126    /// v2.3: feed close-marker inputs — the feed handle + data dir,
127    /// present iff feed enabled AND persistent. Written after the AOF
128    /// flush so the marker's cursor describes durable state.
129    #[cfg(not(target_arch = "wasm32"))]
130    pub(crate) feed_close: Option<(
131        std::sync::Arc<Mutex<kevy_replicate::feed::FeedSource>>,
132        std::path::PathBuf,
133    )>,
134    /// Replica-source listener + accepted connection threads, present
135    /// iff this store is an embed-as-writer
136    /// (`Config::embed_writer_listen_addr = Some(...)`). Joined on
137    /// last-clone drop.
138    #[cfg(not(target_arch = "wasm32"))]
139    pub(crate) replica_source: Option<crate::replica_source::ReplicaSource>,
140}
141
142impl Drop for DropGuard {
143    fn drop(&mut self) {
144        // Stop the replica runner FIRST so no more frames arrive while
145        // we're shutting down + flushing the AOF (frames would race
146        // with the shutdown path).
147        #[cfg(not(target_arch = "wasm32"))]
148        if let Some(r) = &self.replica_runner {
149            r.shutdown();
150        }
151        // Stop the writer-source accept + connection threads next, so
152        // no new replica picks up bytes mid-flush.
153        #[cfg(not(target_arch = "wasm32"))]
154        if let Some(rs) = &self.replica_source {
155            rs.shutdown();
156        }
157        // Stop + join the reaper, then flush every shard's AOF so EverySec
158        // users don't lose the last sub-second of writes.
159        if let Some(stop) = &self.reaper_stop {
160            stop.store(true, Ordering::Relaxed);
161        }
162        if let Some(j) = self
163            .reaper_join
164            .lock()
165            .unwrap_or_else(std::sync::PoisonError::into_inner)
166            .take()
167        {
168            let _ = j.join();
169        }
170        for shard in self.shards_for_flush.iter() {
171            let mut g = lock_write(shard);
172            if let Some(aof) = &mut g.aof {
173                let _ = aof.maybe_sync();
174            }
175        }
176        // v2.3: with the AOF durable, record the feed continuity
177        // marker — the cursor now exactly describes on-disk state.
178        #[cfg(not(target_arch = "wasm32"))]
179        if let Some((feed, dir)) = &self.feed_close {
180            Store::feed_write_close_marker(feed, dir);
181        }
182    }
183}