ankurah_core/
livequery.rs

1use std::{
2    marker::PhantomData,
3    sync::{atomic::AtomicBool, Arc, Weak},
4};
5
6use ankurah_proto::{self as proto, CollectionId};
7
8use ankurah_signals::{
9    broadcast::{BroadcastId, Listener, ListenerGuard},
10    porcelain::subscribe::{IntoSubscribeListener, SubscriptionGuard},
11    Get, Mut, Peek, Read, Signal, Subscribe,
12};
13use tracing::{debug, warn};
14
15use crate::{
16    changes::ChangeSet,
17    entity::Entity,
18    error::RetrievalError,
19    model::View,
20    node::{MatchArgs, TNodeErased},
21    policy::PolicyAgent,
22    reactor::{
23        fetch_gap::{GapFetcher, QueryGapFetcher},
24        ReactorSubscription, ReactorUpdate,
25    },
26    resultset::{EntityResultSet, ResultSet},
27    storage::StorageEngine,
28    Node,
29};
30
31/// A local subscription that handles both reactor subscription and remote cleanup
32/// This is a type-erased version that can be used in the TContext trait
33#[derive(Clone)]
34pub struct EntityLiveQuery(Arc<Inner>);
35struct Inner {
36    pub(crate) query_id: proto::QueryId,
37    pub(crate) node: Box<dyn TNodeErased>,
38    // Store the actual subscription - now non-generic!
39    pub(crate) subscription: ReactorSubscription,
40    pub(crate) resultset: EntityResultSet,
41    pub(crate) error: Mut<Option<RetrievalError>>,
42    pub(crate) initialized: tokio::sync::Notify,
43    pub(crate) initialized_version: std::sync::atomic::AtomicU32,
44    // Version tracking for predicate updates
45    pub(crate) current_version: std::sync::atomic::AtomicU32,
46    // Store selection with its version (starts with version 1, updated on selection changes)
47    // This represents user intent (client-side state), separate from reactor's QueryState.selection (reactor-side state)
48    pub(crate) selection: std::sync::Mutex<(ankql::ast::Selection, u32)>,
49    // Store collection_id for selection updates
50    pub(crate) collection_id: CollectionId,
51    // Gap fetcher for reactor.add_query (type-erased)
52    pub(crate) gap_fetcher: std::sync::Arc<dyn GapFetcher<Entity>>,
53}
54
55/// Weak reference to EntityLiveQuery for breaking circular dependencies
56pub struct WeakEntityLiveQuery(Weak<Inner>);
57
58impl WeakEntityLiveQuery {
59    pub fn upgrade(&self) -> Option<EntityLiveQuery> { self.0.upgrade().map(EntityLiveQuery) }
60}
61
62impl Clone for WeakEntityLiveQuery {
63    fn clone(&self) -> Self { Self(self.0.clone()) }
64}
65
66#[derive(Clone)]
67pub struct LiveQuery<R: View>(EntityLiveQuery, PhantomData<R>);
68
69impl<R: View> std::ops::Deref for LiveQuery<R> {
70    type Target = EntityLiveQuery;
71    fn deref(&self) -> &Self::Target { &self.0 }
72}
73
74impl crate::reactor::PreNotifyHook for &EntityLiveQuery {
75    fn pre_notify(&self, version: u32) {
76        // Mark as initialized before notification is sent
77        self.mark_initialized(version);
78    }
79}
80
81impl EntityLiveQuery {
82    pub fn new<SE, PA>(
83        node: &Node<SE, PA>,
84        collection_id: CollectionId,
85        mut args: MatchArgs,
86        cdata: PA::ContextData,
87    ) -> Result<Self, RetrievalError>
88    where
89        SE: StorageEngine + Send + Sync + 'static,
90        PA: PolicyAgent + Send + Sync + 'static,
91    {
92        node.policy_agent.can_access_collection(&cdata, &collection_id)?;
93        args.selection.predicate = node.policy_agent.filter_predicate(&cdata, &collection_id, args.selection.predicate)?;
94
95        let subscription = node.reactor.subscribe();
96
97        let resultset = EntityResultSet::empty();
98        let query_id = proto::QueryId::new();
99        let gap_fetcher: std::sync::Arc<dyn GapFetcher<Entity>> = std::sync::Arc::new(QueryGapFetcher::new(&node, cdata.clone()));
100
101        let me = Self(Arc::new(Inner {
102            query_id,
103            node: Box::new(node.clone()),
104            subscription,
105            resultset: resultset.clone(),
106            error: Mut::new(None),
107            initialized: tokio::sync::Notify::new(),
108            initialized_version: std::sync::atomic::AtomicU32::new(0), // 0 means uninitialized
109            current_version: std::sync::atomic::AtomicU32::new(1),     // Start at version 1
110            selection: std::sync::Mutex::new((args.selection.clone(), 1)), // Start with version 1
111            collection_id: collection_id.clone(),
112            gap_fetcher,
113        }));
114
115        // Check if this is a durable node (no relay) or ephemeral node (has relay)
116        let has_relay = node.subscription_relay.is_some();
117
118        if args.cached || !has_relay {
119            // Durable node: spawn initialization task directly (no remote subscription needed)
120            let me2 = me.clone();
121
122            debug!("LiveQuery::new() spawning initialization task for durable node predicate {}", query_id);
123            crate::task::spawn(async move {
124                debug!("LiveQuery initialization task starting for predicate {}", query_id);
125                if let Err(e) = me2.activate(1).await {
126                    debug!("LiveQuery initialization failed for predicate {}: {}", query_id, e);
127
128                    me2.0.error.set(Some(e));
129                } else {
130                    debug!("LiveQuery initialization completed for predicate {}", query_id);
131                }
132            });
133        }
134
135        // Ephemeral node: register with relay for remote subscription
136        // Remote will call activate() after applying deltas via subscription_established
137        if has_relay {
138            node.subscribe_remote_query(query_id, collection_id.clone(), args.selection.clone(), cdata.clone(), 1, me.weak());
139        }
140
141        Ok(me)
142    }
143    pub fn map<R: View>(self) -> LiveQuery<R> { LiveQuery(self, PhantomData) }
144
145    /// Wait for the LiveQuery to be fully initialized with initial states
146    pub async fn wait_initialized(&self) {
147        // If already initialized, return immediately
148        if self.0.initialized_version.load(std::sync::atomic::Ordering::Relaxed)
149            >= self.0.current_version.load(std::sync::atomic::Ordering::Relaxed)
150        {
151            return;
152        }
153
154        // Otherwise wait for the notification
155        self.0.initialized.notified().await;
156    }
157
158    pub fn update_selection(
159        &self,
160        new_selection: impl TryInto<ankql::ast::Selection, Error = impl Into<RetrievalError>>,
161    ) -> Result<(), RetrievalError> {
162        let new_selection = new_selection.try_into().map_err(|e| e.into())?;
163
164        // Increment current_version atomically and get the new version number
165        let new_version = self.0.current_version.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
166
167        // Store new selection and version in pending_selection mutex
168        *self.0.selection.lock().unwrap() = (new_selection.clone(), new_version);
169
170        // Check if this node has a relay (ephemeral) or not (durable)
171        let has_relay = self.0.node.has_subscription_relay();
172
173        if has_relay {
174            // Ephemeral node: delegate to relay, which will call update_selection_init after applying deltas
175            self.0.node.update_remote_query(self.0.query_id, new_selection.clone(), new_version)?;
176        } else {
177            // Durable node: spawn task to call update_selection_init directly
178            let me2 = self.clone();
179            let query_id = self.0.query_id;
180
181            crate::task::spawn(async move {
182                if let Err(e) = me2.activate(new_version).await {
183                    tracing::error!("LiveQuery update failed for predicate {}: {}", query_id, e);
184                    me2.0.error.set(Some(e));
185                }
186            });
187        }
188
189        Ok(())
190    }
191
192    pub async fn update_selection_wait(
193        &self,
194        new_selection: impl TryInto<ankql::ast::Selection, Error = impl Into<RetrievalError>>,
195    ) -> Result<(), RetrievalError> {
196        self.update_selection(new_selection)?;
197        self.wait_initialized().await;
198        Ok(())
199    }
200
201    /// Activate the LiveQuery by fetching entities and calling reactor.add_query or reactor.update_query
202    /// Called after deltas have been applied for both initial subscription and selection updates
203    /// Gets all parameters from self.0 (collection_id, query_id, selection)
204    /// Marks initialization as complete regardless of success/failure
205    /// Rejects activation if the version is older than the current selection to prevent regression
206    async fn activate(&self, version: u32) -> Result<(), RetrievalError> {
207        // Get the current selection and its version
208        let (selection, stored_version) = self.0.selection.lock().unwrap().clone();
209
210        // Reject activation if this is an older version than what's currently stored
211        // This prevents out-of-order activations from regressing the state
212        if version < stored_version {
213            warn!("LiveQuery - Dropped stale activation request for version {} (current version is {})", version, stored_version);
214            return Ok(());
215        }
216
217        debug!("LiveQuery.activate() for predicate {} (version {})", self.0.query_id, version);
218
219        let reactor = self.0.node.reactor();
220        let initialized_version = self.0.initialized_version.load(std::sync::atomic::Ordering::Relaxed);
221
222        // Determine if this is the first activation (query not yet in reactor)
223        if initialized_version == 0 {
224            // First activation ever: call reactor.add_query_and_notify which will populate the resultset
225            // Pass self as pre_notify_hook to mark initialized before notification
226            reactor
227                .add_query_and_notify(
228                    self.0.subscription.id(),
229                    self.0.query_id,
230                    self.0.collection_id.clone(),
231                    selection,
232                    &*self.0.node,
233                    self.0.resultset.clone(),
234                    self.0.gap_fetcher.clone(),
235                    self,
236                )
237                .await?
238        } else {
239            // Subsequent activation (including cached re-initialization or selection update): use update_query_and_notify
240            // This handles both: (1) cached queries re-activating after remote deltas, and (2) selection updates
241            reactor
242                .update_query_and_notify(
243                    self.0.subscription.id(),
244                    self.0.query_id,
245                    self.0.collection_id.clone(),
246                    selection,
247                    &*self.0.node,
248                    version,
249                    self,
250                )
251                .await?;
252        };
253
254        Ok(())
255    }
256    pub fn error(&self) -> Read<Option<RetrievalError>> { self.0.error.read() }
257    pub fn query_id(&self) -> proto::QueryId { self.0.query_id }
258
259    /// Create a weak reference to this LiveQuery
260    pub fn weak(&self) -> WeakEntityLiveQuery { WeakEntityLiveQuery(Arc::downgrade(&self.0)) }
261
262    /// Mark initialization as complete for a given version
263    pub fn mark_initialized(&self, version: u32) {
264        // TASK: Serialize or coalesce concurrent activations to prevent version regression https://github.com/ankurah/ankurah/issues/146
265        self.0.initialized_version.store(version, std::sync::atomic::Ordering::Relaxed);
266        self.0.initialized.notify_waiters();
267    }
268}
269
270impl Drop for Inner {
271    fn drop(&mut self) { self.node.unsubscribe_remote_predicate(self.query_id); }
272}
273
274// Implement RemoteQuerySubscriber for WeakEntityLiveQuery to break circular dependencies
275#[async_trait::async_trait]
276impl crate::peer_subscription::RemoteQuerySubscriber for WeakEntityLiveQuery {
277    async fn subscription_established(&self, version: u32) {
278        // Try to upgrade the weak reference
279        if let Some(livequery) = self.upgrade() {
280            // Activate the query (fetch entities, call reactor, and mark initialized)
281            // Handle errors internally by setting last_error
282            if let Err(e) = livequery.activate(version).await {
283                tracing::error!("Failed to activate subscription for query {}: {}", livequery.0.query_id, e);
284                livequery.0.error.set(Some(e));
285            }
286        }
287        // If upgrade fails, the LiveQuery was already dropped - nothing to do
288    }
289
290    fn set_last_error(&self, error: RetrievalError) {
291        // Try to upgrade the weak reference
292        if let Some(livequery) = self.upgrade() {
293            tracing::info!("Setting last error for LiveQuery {}: {}", livequery.0.query_id, error);
294            livequery.0.error.set(Some(error));
295        }
296        // If upgrade fails, the LiveQuery was already dropped - nothing to do
297    }
298}
299
300impl<R: View> LiveQuery<R> {
301    /// Wait for the LiveQuery to be fully initialized with initial states
302    pub async fn wait_initialized(&self) { self.0.wait_initialized().await; }
303
304    pub fn resultset(&self) -> ResultSet<R> { self.0 .0.resultset.wrap::<R>() }
305
306    pub fn loaded(&self) -> bool { self.0 .0.resultset.is_loaded() }
307
308    pub fn ids(&self) -> Vec<proto::EntityId> { self.0 .0.resultset.keys().collect() }
309
310    pub fn ids_sorted(&self) -> Vec<proto::EntityId> {
311        use itertools::Itertools;
312        self.0 .0.resultset.keys().sorted().collect()
313    }
314}
315
316// Implement Signal trait - delegate to the resultset
317impl<R: View> Signal for LiveQuery<R> {
318    fn listen(&self, listener: Listener) -> ListenerGuard { self.0 .0.resultset.listen(listener) }
319
320    fn broadcast_id(&self) -> BroadcastId { self.0 .0.resultset.broadcast_id() }
321}
322
323// Implement Get trait - delegate to ResultSet<R>
324impl<R: View + Clone + 'static> Get<Vec<R>> for LiveQuery<R> {
325    fn get(&self) -> Vec<R> { self.0 .0.resultset.wrap::<R>().get() }
326}
327
328// Implement Peek trait - delegate to ResultSet<R>
329impl<R: View + Clone + 'static> Peek<Vec<R>> for LiveQuery<R> {
330    fn peek(&self) -> Vec<R> { self.0 .0.resultset.wrap().peek() }
331}
332
333// Implement Subscribe trait - convert ReactorUpdate to ChangeSet<R>
334impl<R: View> Subscribe<ChangeSet<R>> for LiveQuery<R>
335where R: Clone + Send + Sync + 'static
336{
337    fn subscribe<L>(&self, listener: L) -> SubscriptionGuard
338    where L: IntoSubscribeListener<ChangeSet<R>> {
339        let listener = listener.into_subscribe_listener();
340
341        let me = self.clone();
342        // Subscribe to the underlying ReactorUpdate stream and convert to ChangeSet<R>
343
344        self.0 .0.subscription.subscribe(move |reactor_update: ReactorUpdate| {
345            // Convert ReactorUpdate to ChangeSet<R>");
346            let changeset: ChangeSet<R> = livequery_change_set_from(me.0 .0.resultset.wrap::<R>(), reactor_update);
347            listener(changeset);
348        })
349    }
350}
351
352/// Notably, this function does not filter by query_id, because it should only be used by LiveQuery, which entails a single-predicate subscription
353fn livequery_change_set_from<R: View>(resultset: ResultSet<R>, reactor_update: ReactorUpdate) -> ChangeSet<R>
354where R: View {
355    use crate::changes::{ChangeSet, ItemChange};
356
357    let mut changes = Vec::new();
358
359    for item in reactor_update.items {
360        let view = R::from_entity(item.entity);
361
362        // Determine the change type based on predicate relevance
363        // ignore the query_id, because it should only be used by LiveQuery, which entails a single-predicate subscription
364        if let Some((_, membership_change)) = item.predicate_relevance.first() {
365            match membership_change {
366                crate::reactor::MembershipChange::Initial => {
367                    changes.push(ItemChange::Initial { item: view });
368                }
369                crate::reactor::MembershipChange::Add => {
370                    changes.push(ItemChange::Add { item: view, events: item.events });
371                }
372                crate::reactor::MembershipChange::Remove => {
373                    changes.push(ItemChange::Remove { item: view, events: item.events });
374                }
375            }
376        } else {
377            // No membership change, just an update
378            changes.push(ItemChange::Update { item: view, events: item.events });
379        }
380    }
381
382    ChangeSet { changes, resultset }
383}
ankurah_core/livequery.rs

ankurah_core/
livequery.rs
