ankurah_core/
livequery.rs

1use std::{
2    marker::PhantomData,
3    sync::{Arc, Weak},
4};
5
6use ankurah_proto::{self as proto, CollectionId};
7
8use ankurah_signals::{
9    broadcast::BroadcastId,
10    porcelain::subscribe::{IntoSubscribeListener, SubscriptionGuard},
11    signal::{Listener, ListenerGuard},
12    Get, Mut, Peek, Read, Signal, Subscribe,
13};
14use tracing::{debug, warn};
15
16use crate::{
17    changes::ChangeSet,
18    entity::Entity,
19    error::RetrievalError,
20    model::View,
21    node::{MatchArgs, NodeInner, TNodeErased},
22    policy::PolicyAgent,
23    reactor::{
24        fetch_gap::{GapFetcher, QueryGapFetcher},
25        ReactorSubscription, ReactorUpdate,
26    },
27    resultset::{EntityResultSet, ResultSet},
28    storage::StorageEngine,
29    Node,
30};
31
32/// A local subscription that handles both reactor subscription and remote cleanup
33/// This is a type-erased version that can be used in the TContext trait
34///
35/// Whether the query keeps its node alive is a construction-time choice:
36/// [`EntityLiveQuery::new`] holds the node strongly, [`EntityLiveQuery::new_weak_node`] does not.
37#[derive(Clone)]
38pub struct EntityLiveQuery(Arc<Inner>);
39
40/// Type-erased reference to a node. Strong variants keep the node alive; weak variants do not.
41trait NodeRef: Send + Sync {
42    fn upgrade(&self) -> Option<Box<dyn TNodeErased>>;
43}
44
45/// Strong node reference — keeps the node alive as long as Inner exists.
46struct StrongNodeRef<SE, PA: PolicyAgent>(Arc<NodeInner<SE, PA>>);
47
48impl<SE, PA> NodeRef for StrongNodeRef<SE, PA>
49where
50    SE: StorageEngine + Send + Sync + 'static,
51    PA: PolicyAgent + Send + Sync + 'static,
52{
53    fn upgrade(&self) -> Option<Box<dyn TNodeErased>> { Some(Box::new(Node(self.0.clone()))) }
54}
55
56/// Weak node reference — does NOT keep the node alive.
57struct WeakNodeRefImpl<SE, PA: PolicyAgent>(Weak<NodeInner<SE, PA>>);
58
59impl<SE, PA> NodeRef for WeakNodeRefImpl<SE, PA>
60where
61    SE: StorageEngine + Send + Sync + 'static,
62    PA: PolicyAgent + Send + Sync + 'static,
63{
64    fn upgrade(&self) -> Option<Box<dyn TNodeErased>> { self.0.upgrade().map(|inner| Box::new(Node(inner)) as Box<dyn TNodeErased>) }
65}
66
67struct Inner {
68    pub(crate) query_id: proto::QueryId,
69    // subscription must be declared before node so it drops first —
70    // dropping node (StrongNodeRef) deallocates the reactor, and
71    // subscription's Drop needs the reactor to unsubscribe.
72    pub(crate) subscription: ReactorSubscription,
73    pub(crate) node: Box<dyn NodeRef>,
74    pub(crate) resultset: EntityResultSet,
75    pub(crate) error: Mut<Option<RetrievalError>>,
76    pub(crate) initialized: tokio::sync::Notify,
77    pub(crate) initialized_version: std::sync::atomic::AtomicU32,
78    // Version tracking for predicate updates
79    pub(crate) current_version: std::sync::atomic::AtomicU32,
80    // Store selection with its version (starts with version 1, updated on selection changes)
81    // This represents user intent (client-side state), separate from reactor's QueryState.selection (reactor-side state)
82    // Using Mut for reactive updates that can be observed in WASM
83    pub(crate) selection: Mut<(ankql::ast::Selection, u32)>,
84    // Store collection_id for selection updates
85    pub(crate) collection_id: CollectionId,
86    // Gap fetcher for reactor.add_query (type-erased)
87    pub(crate) gap_fetcher: std::sync::Arc<dyn GapFetcher<Entity>>,
88}
89
90/// Weak reference to EntityLiveQuery for breaking circular dependencies
91pub struct WeakEntityLiveQuery(Weak<Inner>);
92
93impl WeakEntityLiveQuery {
94    pub fn upgrade(&self) -> Option<EntityLiveQuery> { self.0.upgrade().map(EntityLiveQuery) }
95}
96
97impl Clone for WeakEntityLiveQuery {
98    fn clone(&self) -> Self { Self(self.0.clone()) }
99}
100
101#[derive(Clone)]
102pub struct LiveQuery<R: View>(EntityLiveQuery, PhantomData<R>);
103
104impl<R: View> std::ops::Deref for LiveQuery<R> {
105    type Target = EntityLiveQuery;
106    fn deref(&self) -> &Self::Target { &self.0 }
107}
108
109impl Inner {
110    fn node(&self) -> Option<Box<dyn TNodeErased>> { self.node.upgrade() }
111
112    async fn wait_initialized(&self) {
113        // If already initialized, return immediately
114        if self.initialized_version.load(std::sync::atomic::Ordering::Relaxed)
115            >= self.current_version.load(std::sync::atomic::Ordering::Relaxed)
116        {
117            return;
118        }
119
120        // FIXME - this should be waiting for the correct version, not any version
121        // Otherwise wait for the notification
122        self.initialized.notified().await;
123    }
124
125    /// Activate the LiveQuery by fetching entities and calling reactor.add_query or reactor.update_query
126    /// Called after deltas have been applied for both initial subscription and selection updates
127    /// Gets all parameters from self (collection_id, query_id, selection)
128    /// Marks initialization as complete regardless of success/failure
129    /// Rejects activation if the version is older than the current selection to prevent regression
130    async fn activate(&self, version: u32) -> Result<(), RetrievalError> {
131        // Get the current selection and its version
132        let (selection, stored_version) = self.selection.value();
133
134        // Reject activation if this is an older version than what's currently stored
135        // This prevents out-of-order activations from regressing the state
136        if version < stored_version {
137            warn!("LiveQuery - Dropped stale activation request for version {} (current version is {})", version, stored_version);
138            return Ok(());
139        }
140
141        debug!("LiveQuery.activate() for predicate {} (version {})", self.query_id, version);
142
143        let node = self.node().ok_or_else(|| RetrievalError::Other("Node has been dropped".into()))?;
144        let reactor = node.reactor();
145        let initialized_version = self.initialized_version.load(std::sync::atomic::Ordering::Relaxed);
146
147        let hook = InnerPreNotifyHook(self);
148        // Determine if this is the first activation (query not yet in reactor)
149        if initialized_version == 0 {
150            // First activation ever: call reactor.add_query_and_notify which will populate the resultset
151            // Pass the hook as pre_notify_hook to mark initialized before notification
152            reactor
153                .add_query_and_notify(
154                    self.subscription.id(),
155                    self.query_id,
156                    self.collection_id.clone(),
157                    selection,
158                    &*node,
159                    self.resultset.clone(),
160                    self.gap_fetcher.clone(),
161                    &hook,
162                )
163                .await?
164        } else {
165            // Subsequent activation (including cached re-initialization or selection update): use update_query_and_notify
166            // This handles both: (1) cached queries re-activating after remote deltas, and (2) selection updates
167            reactor
168                .update_query_and_notify(
169                    self.subscription.id(),
170                    self.query_id,
171                    self.collection_id.clone(),
172                    selection,
173                    &*node,
174                    version,
175                    &hook,
176                )
177                .await?;
178        };
179
180        Ok(())
181    }
182
183    /// Mark initialization as complete for a given version
184    fn mark_initialized(&self, version: u32) {
185        // TASK: Serialize or coalesce concurrent activations to prevent version regression https://github.com/ankurah/ankurah/issues/146
186        self.initialized_version.store(version, std::sync::atomic::Ordering::Relaxed);
187        self.initialized.notify_waiters();
188    }
189}
190
191/// Adapts a borrowed Inner to the reactor's PreNotifyHook (previously implemented on &EntityLiveQuery,
192/// but activation now lives on Inner so both LiveQuery variants share it)
193struct InnerPreNotifyHook<'a>(&'a Inner);
194impl crate::reactor::PreNotifyHook for &InnerPreNotifyHook<'_> {
195    fn pre_notify(&self, version: u32) {
196        // Mark as initialized before notification is sent
197        self.0.mark_initialized(version);
198    }
199}
200
201/// Helper: create the Inner and set up initialization (shared by strong- and weak-node constructors)
202fn create_inner<SE, PA>(
203    node: &Node<SE, PA>,
204    node_ref: Box<dyn NodeRef>,
205    collection_id: CollectionId,
206    mut args: MatchArgs,
207    cdata: PA::ContextData,
208) -> Result<(Arc<Inner>, proto::QueryId), RetrievalError>
209where
210    SE: StorageEngine + Send + Sync + 'static,
211    PA: PolicyAgent + Send + Sync + 'static,
212{
213    node.policy_agent.can_access_collection(&cdata, &collection_id)?;
214    args.selection.predicate = node.policy_agent.filter_predicate(&cdata, &collection_id, args.selection.predicate)?;
215
216    // Resolve types in the AST (converts literals for JSON path comparisons)
217    args.selection = node.type_resolver.resolve_selection_types(args.selection);
218
219    let subscription = node.reactor.subscribe();
220
221    let resultset = EntityResultSet::empty();
222    let query_id = proto::QueryId::new();
223    let gap_fetcher: std::sync::Arc<dyn GapFetcher<Entity>> = std::sync::Arc::new(QueryGapFetcher::new(&node, cdata.clone()));
224
225    let inner = Arc::new(Inner {
226        query_id,
227        node: node_ref,
228        subscription,
229        resultset: resultset.clone(),
230        error: Mut::new(None),
231        initialized: tokio::sync::Notify::new(),
232        initialized_version: std::sync::atomic::AtomicU32::new(0), // 0 means uninitialized
233        current_version: std::sync::atomic::AtomicU32::new(1),     // Start at version 1
234        selection: Mut::new((args.selection.clone(), 1)),          // Start with version 1
235        collection_id: collection_id.clone(),
236        gap_fetcher,
237    });
238
239    // Check if this is a durable node (no relay) or ephemeral node (has relay)
240    let has_relay = node.subscription_relay.is_some();
241
242    if args.cached || !has_relay {
243        // Durable node: spawn initialization task directly (no remote subscription needed)
244        let inner2 = inner.clone();
245
246        debug!("LiveQuery::new() spawning initialization task for durable node predicate {}", query_id);
247        crate::task::spawn(async move {
248            debug!("LiveQuery initialization task starting for predicate {}", query_id);
249            if let Err(e) = inner2.activate(1).await {
250                debug!("LiveQuery initialization failed for predicate {}: {}", query_id, e);
251                inner2.error.set(Some(e));
252            } else {
253                debug!("LiveQuery initialization completed for predicate {}", query_id);
254            }
255        });
256    }
257
258    Ok((inner, query_id))
259}
260
261impl EntityLiveQuery {
262    pub fn new<SE, PA>(
263        node: &Node<SE, PA>,
264        collection_id: CollectionId,
265        args: MatchArgs,
266        cdata: PA::ContextData,
267    ) -> Result<Self, RetrievalError>
268    where
269        SE: StorageEngine + Send + Sync + 'static,
270        PA: PolicyAgent + Send + Sync + 'static,
271    {
272        let node_ref: Box<dyn NodeRef> = Box::new(StrongNodeRef(Arc::clone(&node.0)));
273        Self::new_with_node_ref(node, node_ref, collection_id, args, cdata)
274    }
275
276    /// Create a LiveQuery that does NOT keep the node alive.
277    ///
278    /// Used by PolicyAgent and other internal subscribers that should not create
279    /// reference cycles (node → agent → livequery → node). Operations that need
280    /// the node (activation, selection updates) fail with "Node has been dropped"
281    /// once the node is gone.
282    pub fn new_weak_node<SE, PA>(
283        node: &Node<SE, PA>,
284        collection_id: CollectionId,
285        args: MatchArgs,
286        cdata: PA::ContextData,
287    ) -> Result<Self, RetrievalError>
288    where
289        SE: StorageEngine + Send + Sync + 'static,
290        PA: PolicyAgent + Send + Sync + 'static,
291    {
292        let node_ref: Box<dyn NodeRef> = Box::new(WeakNodeRefImpl(Arc::downgrade(&node.0)));
293        Self::new_with_node_ref(node, node_ref, collection_id, args, cdata)
294    }
295
296    fn new_with_node_ref<SE, PA>(
297        node: &Node<SE, PA>,
298        node_ref: Box<dyn NodeRef>,
299        collection_id: CollectionId,
300        args: MatchArgs,
301        cdata: PA::ContextData,
302    ) -> Result<Self, RetrievalError>
303    where
304        SE: StorageEngine + Send + Sync + 'static,
305        PA: PolicyAgent + Send + Sync + 'static,
306    {
307        let has_relay = node.subscription_relay.is_some();
308        let (inner, query_id) = create_inner(node, node_ref, collection_id.clone(), args, cdata.clone())?;
309
310        let me = Self(inner.clone());
311
312        // Ephemeral node: register with relay for remote subscription
313        // Remote will call activate() after applying deltas via subscription_established
314        if has_relay {
315            node.subscribe_remote_query(query_id, collection_id, inner.selection.value().0, cdata, 1, me.weak());
316        }
317
318        Ok(me)
319    }
320    pub fn map<R: View>(self) -> LiveQuery<R> { LiveQuery(self, PhantomData) }
321
322    /// Wait for the LiveQuery to be fully initialized with initial states
323    pub async fn wait_initialized(&self) { self.0.wait_initialized().await; }
324
325    pub fn update_selection(
326        &self,
327        new_selection: impl TryInto<ankql::ast::Selection, Error = impl Into<RetrievalError>>,
328    ) -> Result<(), RetrievalError> {
329        let new_selection = new_selection.try_into().map_err(|e| e.into())?;
330        let node = self.0.node().ok_or_else(|| RetrievalError::Other("Node has been dropped".into()))?;
331
332        // Increment current_version atomically and get the new version number
333        let new_version = self.0.current_version.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
334
335        // Mark resultset as not loaded since we're changing the selection
336        self.0.resultset.set_loaded(false);
337
338        // Store new selection and version
339        self.0.selection.set((new_selection.clone(), new_version));
340
341        // Check if this node has a relay (ephemeral) or not (durable)
342        let has_relay = node.has_subscription_relay();
343
344        if has_relay {
345            // Ephemeral node: delegate to relay, which will call update_selection_init after applying deltas
346            node.update_remote_query(self.0.query_id, new_selection.clone(), new_version)?;
347        } else {
348            // Durable node: spawn task to call update_selection_init directly
349            let inner = self.0.clone();
350            let query_id = self.0.query_id;
351
352            crate::task::spawn(async move {
353                if let Err(e) = inner.activate(new_version).await {
354                    tracing::error!("LiveQuery update failed for predicate {}: {}", query_id, e);
355                    inner.error.set(Some(e));
356                }
357            });
358        }
359
360        Ok(())
361    }
362
363    pub async fn update_selection_wait(
364        &self,
365        new_selection: impl TryInto<ankql::ast::Selection, Error = impl Into<RetrievalError>>,
366    ) -> Result<(), RetrievalError> {
367        self.update_selection(new_selection)?;
368        self.0.wait_initialized().await;
369        Ok(())
370    }
371
372    pub fn error(&self) -> Read<Option<RetrievalError>> { self.0.error.read() }
373    pub fn query_id(&self) -> proto::QueryId { self.0.query_id }
374    pub fn selection(&self) -> Read<(ankql::ast::Selection, u32)> { self.0.selection.read() }
375    pub fn resultset(&self) -> EntityResultSet { self.0.resultset.clone() }
376
377    /// Create a weak reference to this LiveQuery
378    pub fn weak(&self) -> WeakEntityLiveQuery { WeakEntityLiveQuery(Arc::downgrade(&self.0)) }
379}
380
381impl Drop for Inner {
382    fn drop(&mut self) {
383        if let Some(node) = self.node.upgrade() {
384            node.unsubscribe_remote_predicate(self.query_id);
385        }
386    }
387}
388
389// Implement RemoteQuerySubscriber for WeakEntityLiveQuery to break circular dependencies
390#[async_trait::async_trait]
391impl crate::peer_subscription::RemoteQuerySubscriber for WeakEntityLiveQuery {
392    async fn subscription_established(&self, version: u32) {
393        // Try to upgrade the weak reference
394        if let Some(inner) = self.0.upgrade() {
395            // Activate the query (fetch entities, call reactor, and mark initialized)
396            // Handle errors internally by setting last_error
397            tracing::debug!("Subscription established for query {}: {}", inner.query_id, version);
398            if let Err(e) = inner.activate(version).await {
399                tracing::error!("Failed to activate subscription for query {}: {}", inner.query_id, e);
400                inner.error.set(Some(e));
401            }
402        }
403        // If upgrade fails, the LiveQuery was already dropped - nothing to do
404    }
405
406    fn set_last_error(&self, error: RetrievalError) {
407        // Try to upgrade the weak reference
408        if let Some(inner) = self.0.upgrade() {
409            tracing::info!("Setting last error for LiveQuery {}: {}", inner.query_id, error);
410            inner.error.set(Some(error));
411        }
412        // If upgrade fails, the LiveQuery was already dropped - nothing to do
413    }
414}
415
416impl<R: View> LiveQuery<R> {
417    /// Wait for the LiveQuery to be fully initialized with initial states
418    pub async fn wait_initialized(&self) { self.0.wait_initialized().await; }
419
420    pub fn resultset(&self) -> ResultSet<R> { self.0 .0.resultset.wrap::<R>() }
421
422    pub fn loaded(&self) -> bool { self.0 .0.resultset.is_loaded() }
423
424    pub fn ids(&self) -> Vec<proto::EntityId> { self.0 .0.resultset.keys().collect() }
425
426    pub fn ids_sorted(&self) -> Vec<proto::EntityId> {
427        use itertools::Itertools;
428        self.0 .0.resultset.keys().sorted().collect()
429    }
430}
431
432// Implement Signal trait - delegate to the subscription (not resultset)
433// This ensures that LiveQuery tracking fires on ALL entity changes, not just membership changes
434impl<R: View> Signal for LiveQuery<R> {
435    fn listen(&self, listener: Listener) -> ListenerGuard { self.0 .0.subscription.listen(listener) }
436
437    fn broadcast_id(&self) -> BroadcastId { self.0 .0.subscription.broadcast_id() }
438}
439
440// Implement Get trait - delegate to ResultSet<R>
441impl<R: View + Clone + 'static> Get<Vec<R>> for LiveQuery<R> {
442    fn get(&self) -> Vec<R> {
443        use ankurah_signals::CurrentObserver;
444        CurrentObserver::track(&self);
445        self.0 .0.resultset.wrap::<R>().peek()
446    }
447}
448
449// Implement Peek trait - delegate to ResultSet<R>
450impl<R: View + Clone + 'static> Peek<Vec<R>> for LiveQuery<R> {
451    fn peek(&self) -> Vec<R> { self.0 .0.resultset.wrap().peek() }
452}
453
454// Implement Subscribe trait - convert ReactorUpdate to ChangeSet<R>
455impl<R: View> Subscribe<ChangeSet<R>> for LiveQuery<R>
456where R: Clone + Send + Sync + 'static
457{
458    fn subscribe<L>(&self, listener: L) -> SubscriptionGuard
459    where L: IntoSubscribeListener<ChangeSet<R>> {
460        let listener = listener.into_subscribe_listener();
461
462        let me = self.clone();
463        // Subscribe to the underlying ReactorUpdate stream and convert to ChangeSet<R>
464        self.0 .0.subscription.subscribe(move |reactor_update: ReactorUpdate| {
465            let changeset: ChangeSet<R> = livequery_change_set_from(me.0 .0.resultset.wrap::<R>(), reactor_update);
466            listener(changeset);
467        })
468    }
469}
470
471/// Notably, this function does not filter by query_id, because it should only be used by LiveQuery, which entails a single-predicate subscription
472fn livequery_change_set_from<R: View>(resultset: ResultSet<R>, reactor_update: ReactorUpdate) -> ChangeSet<R>
473where R: View {
474    use crate::changes::{ChangeSet, ItemChange};
475
476    let mut changes = Vec::new();
477
478    for item in reactor_update.items {
479        let view = R::from_entity(item.entity);
480
481        // Determine the change type based on predicate relevance
482        // ignore the query_id, because it should only be used by LiveQuery, which entails a single-predicate subscription
483        if let Some((_, membership_change)) = item.predicate_relevance.first() {
484            match membership_change {
485                crate::reactor::MembershipChange::Initial => {
486                    changes.push(ItemChange::Initial { item: view });
487                }
488                crate::reactor::MembershipChange::Add => {
489                    changes.push(ItemChange::Add { item: view, events: item.events });
490                }
491                crate::reactor::MembershipChange::Remove => {
492                    changes.push(ItemChange::Remove { item: view, events: item.events });
493                }
494            }
495        } else {
496            // No membership change, just an update
497            changes.push(ItemChange::Update { item: view, events: item.events });
498        }
499    }
500
501    ChangeSet { changes, resultset }
502}
ankurah_core/livequery.rs

ankurah_core/
livequery.rs
