linera_service/
query_subscription.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    collections::HashMap,
6    sync::{Arc, Mutex},
7    time::Duration,
8};
9
10use futures::StreamExt as _;
11use linera_base::identifiers::{ApplicationId, ChainId};
12use linera_client::chain_listener::ClientContext;
13use linera_core::worker::Reason;
14use linera_execution::{Query, QueryResponse};
15use tokio::sync::watch;
16use tokio_util::sync::CancellationToken;
17use tracing::{debug, warn};
18
19/// A named GraphQL query string registered at startup via `--allow-subscription`.
20#[derive(Clone, Debug)]
21pub struct RegisteredQuery {
22    /// The operation name used to refer to the query.
23    pub name: String,
24    /// The full GraphQL query string.
25    pub query: String,
26}
27
28/// Parses a GraphQL query string like `query Name { ... }` and extracts the operation name.
29pub fn parse_allowed_subscription(s: &str) -> anyhow::Result<RegisteredQuery> {
30    let trimmed = s.trim();
31    let rest = trimmed
32        .strip_prefix("query")
33        .ok_or_else(|| anyhow::anyhow!("expected query to start with 'query', got: {s}"))?;
34    // The character right after "query" must be whitespace (not part of a longer word).
35    anyhow::ensure!(
36        rest.starts_with(char::is_whitespace),
37        "expected whitespace after 'query' keyword"
38    );
39    let rest = rest.trim_start();
40    // Extract the operation name: sequence of alphanumeric/underscore chars.
41    let name = rest
42        .split(|c: char| !c.is_alphanumeric() && c != '_')
43        .next()
44        .unwrap_or_default();
45    anyhow::ensure!(
46        !name.is_empty(),
47        "expected an operation name after 'query', e.g. 'query MyQuery {{ ... }}'"
48    );
49    Ok(RegisteredQuery {
50        name: name.to_string(),
51        query: trimmed.to_string(),
52    })
53}
54
55/// Parses a `Name=Secs` string into a query name and TTL in seconds.
56pub fn parse_subscription_ttl(s: &str) -> Result<(String, u64), String> {
57    let (name, secs) = s
58        .split_once('=')
59        .ok_or_else(|| format!("expected format Name=Secs, got: {s}"))?;
60    let secs: u64 = secs
61        .parse()
62        .map_err(|e| format!("invalid seconds value '{secs}': {e}"))?;
63    Ok((name.to_string(), secs))
64}
65
66/// Identifies a unique subscription target: a named query for a specific chain and application.
67#[derive(Clone, Debug, Hash, Eq, PartialEq)]
68pub struct SubscriptionKey {
69    /// The name of the registered query.
70    pub name: String,
71    /// The chain the query runs against.
72    pub chain_id: ChainId,
73    /// The application the query targets.
74    pub application_id: ApplicationId,
75}
76
77/// State for an active watcher: the watch sender for the latest query result.
78/// The channel carries pre-serialized JSON strings so that cloning between
79/// subscribers is a single `memcpy` instead of a deep `serde_json::Value` clone.
80struct WatcherState {
81    sender: watch::Sender<Option<String>>,
82}
83
84/// Manages registered query names and active per-key watchers.
85pub struct QuerySubscriptionManager {
86    /// Registered queries by name.
87    queries: HashMap<String, String>,
88    /// Per-query minimum TTL for cached results.
89    ttls: HashMap<String, Duration>,
90    /// Active watchers keyed by subscription target.
91    watchers: Mutex<HashMap<SubscriptionKey, WatcherState>>,
92}
93
94impl QuerySubscriptionManager {
95    /// Creates a new manager from the registered queries and optional per-query TTLs.
96    pub fn new(registered: Vec<RegisteredQuery>, ttls: HashMap<String, Duration>) -> Self {
97        let queries = registered
98            .into_iter()
99            .map(|rq| (rq.name, rq.query))
100            .collect();
101        Self {
102            queries,
103            ttls,
104            watchers: Mutex::new(HashMap::new()),
105        }
106    }
107
108    /// Returns the GraphQL query string for a given name, if registered.
109    pub fn get_query(&self, name: &str) -> Option<&str> {
110        self.queries.get(name).map(|s| s.as_str())
111    }
112
113    /// Returns a watch receiver for the given key. Lazily spawns a watcher if needed.
114    /// The receiver initially holds `None`; the watcher populates it with `Some(value)`
115    /// after the first query. Callers should filter out `None` values from the stream.
116    pub fn subscribe<C: ClientContext + 'static>(
117        self: &Arc<Self>,
118        key: &SubscriptionKey,
119        context: Arc<futures::lock::Mutex<C>>,
120        token: CancellationToken,
121    ) -> anyhow::Result<watch::Receiver<Option<String>>> {
122        let query_string = self
123            .get_query(&key.name)
124            .ok_or_else(|| {
125                anyhow::anyhow!("no subscription query registered with name '{}'", key.name)
126            })?
127            .to_string();
128
129        let mut watchers = self.watchers.lock().unwrap();
130
131        // If a watcher already exists, reuse it.
132        if let Some(state) = watchers.get(key) {
133            return Ok(state.sender.subscribe());
134        }
135
136        // Create a new watch channel (initial value is None until the first query completes).
137        let (sender, receiver) = watch::channel(None);
138        watchers.insert(
139            key.clone(),
140            WatcherState {
141                sender: sender.clone(),
142            },
143        );
144
145        let ttl = self.ttls.get(&key.name).copied();
146        let manager = Arc::clone(self);
147        let key_clone = key.clone();
148        tokio::spawn(run_query_subscription_watcher(
149            context,
150            manager,
151            key_clone,
152            query_string,
153            sender,
154            token,
155            ttl,
156        ));
157
158        Ok(receiver)
159    }
160}
161
162/// Background task that watches for new blocks on a chain and re-executes the query.
163async fn run_query_subscription_watcher<C: ClientContext + 'static>(
164    context: Arc<futures::lock::Mutex<C>>,
165    manager: Arc<QuerySubscriptionManager>,
166    key: SubscriptionKey,
167    query_string: String,
168    sender: watch::Sender<Option<String>>,
169    token: CancellationToken,
170    ttl: Option<Duration>,
171) {
172    debug!(
173        name = %key.name,
174        chain_id = %key.chain_id,
175        application_id = %key.application_id,
176        ?ttl,
177        "starting query subscription watcher"
178    );
179
180    let notification_stream = {
181        let ctx = context.lock().await;
182        match ctx.make_chain_client(key.chain_id).await {
183            Ok(client) => match client.subscribe() {
184                Ok(stream) => stream,
185                Err(e) => {
186                    warn!("failed to subscribe to chain notifications: {e}");
187                    cleanup_watcher(&manager, &key);
188                    return;
189                }
190            },
191            Err(e) => {
192                warn!("failed to create chain client: {e}");
193                cleanup_watcher(&manager, &key);
194                return;
195            }
196        }
197    };
198
199    let mut notification_stream = Box::pin(notification_stream);
200
201    // Cache the last result as a string to deduplicate.
202    let mut last_result: Option<String> = None;
203
204    // Execute the query once immediately so the first subscriber gets a value.
205    execute_and_maybe_send(&context, &key, &query_string, &sender, &mut last_result).await;
206
207    // Track when the last execution happened, for TTL-based deferral.
208    let mut last_execution = tokio::time::Instant::now();
209    // Whether a deferred re-execution is pending (a NewBlock arrived during the TTL window).
210    let mut pending_invalidation = false;
211
212    loop {
213        // If there's a pending invalidation, compute the remaining TTL sleep.
214        let ttl_sleep = if pending_invalidation {
215            if let Some(ttl) = ttl {
216                let elapsed = last_execution.elapsed();
217                if elapsed < ttl {
218                    tokio::time::sleep(ttl - elapsed)
219                } else {
220                    tokio::time::sleep(Duration::ZERO)
221                }
222            } else {
223                // No TTL configured; should not happen since pending_invalidation
224                // is only set when ttl is Some, but handle gracefully.
225                tokio::time::sleep(Duration::ZERO)
226            }
227        } else {
228            // No pending invalidation: sleep forever (effectively disabled).
229            tokio::time::sleep(Duration::MAX)
230        };
231        tokio::pin!(ttl_sleep);
232
233        tokio::select! {
234            _ = token.cancelled() => {
235                debug!(name = %key.name, "watcher cancelled");
236                break;
237            }
238            () = &mut ttl_sleep, if pending_invalidation => {
239                pending_invalidation = false;
240                execute_and_maybe_send(
241                    &context,
242                    &key,
243                    &query_string,
244                    &sender,
245                    &mut last_result,
246                )
247                .await;
248                last_execution = tokio::time::Instant::now();
249            }
250            notification = notification_stream.next() => {
251                match notification {
252                    Some(n) => {
253                        if matches!(n.reason, Reason::NewBlock { .. }) {
254                            if ttl.is_some() {
255                                // Defer re-execution until the TTL expires.
256                                if !pending_invalidation {
257                                    debug!(name = %key.name, "deferring invalidation until TTL expires");
258                                }
259                                pending_invalidation = true;
260                            } else {
261                                execute_and_maybe_send(
262                                    &context,
263                                    &key,
264                                    &query_string,
265                                    &sender,
266                                    &mut last_result,
267                                )
268                                .await;
269                                last_execution = tokio::time::Instant::now();
270                            }
271                        }
272                    }
273                    None => {
274                        debug!(name = %key.name, "notification stream ended");
275                        break;
276                    }
277                }
278
279                // If no receivers remain, stop the watcher.
280                if sender.is_closed() {
281                    debug!(name = %key.name, "no more subscribers, stopping watcher");
282                    break;
283                }
284            }
285        }
286    }
287
288    cleanup_watcher(&manager, &key);
289}
290
291/// Executes the query against the application and updates the watch channel if the result changed.
292async fn execute_and_maybe_send<C: ClientContext + 'static>(
293    context: &Arc<futures::lock::Mutex<C>>,
294    key: &SubscriptionKey,
295    query_string: &str,
296    sender: &watch::Sender<Option<String>>,
297    last_result: &mut Option<String>,
298) {
299    // The application service expects a JSON-encoded GraphQL request.
300    let json_request = serde_json::json!({ "query": query_string });
301    let request_bytes = serde_json::to_vec(&json_request).unwrap();
302    let query = Query::User {
303        application_id: key.application_id,
304        bytes: request_bytes,
305    };
306
307    let result = {
308        let ctx = context.lock().await;
309        match ctx.make_chain_client(key.chain_id).await {
310            Ok(client) => client.query_application(query, None).await,
311            Err(e) => {
312                warn!(name = %key.name, "failed to create chain client: {e}");
313                return;
314            }
315        }
316    };
317
318    match result {
319        Ok((outcome, _height)) => {
320            let response_bytes = match outcome.response {
321                QueryResponse::User(bytes) => bytes,
322                QueryResponse::System(_) => {
323                    warn!(name = %key.name, "unexpected system response for user query");
324                    return;
325                }
326            };
327
328            // Convert to string. The bytes are already valid UTF-8 JSON
329            // from the application service.
330            let json_string = match String::from_utf8(response_bytes) {
331                Ok(s) => s,
332                Err(error) => {
333                    warn!(%error, name = %key.name, "response bytes are not valid UTF-8");
334                    return;
335                }
336            };
337
338            // Deduplicate: only send if the result changed.
339            if last_result.as_ref() == Some(&json_string) {
340                return;
341            }
342
343            if let Err(e) = sender.send(Some(json_string.clone())) {
344                debug!(name = %key.name, "Failed to send graphql response: {e}");
345            }
346            *last_result = Some(json_string);
347        }
348        Err(e) => {
349            warn!(name = %key.name, "query execution failed: {e}");
350        }
351    }
352}
353
354/// Removes the watcher entry from the manager.
355fn cleanup_watcher(manager: &QuerySubscriptionManager, key: &SubscriptionKey) {
356    let mut watchers = manager.watchers.lock().unwrap();
357    watchers.remove(key);
358    debug!(name = %key.name, "watcher cleaned up");
359}
linera_service/query_subscription.rs

linera_service/
query_subscription.rs
