1use std::{collections::HashMap, sync::Arc};
5
6use reifydb_core::{
7 execution::ExecutionResult,
8 interface::catalog::{id::QueueId, token::Token},
9 retry::RetryStrategy,
10};
11use reifydb_runtime::{context::clock::Instant, sync::waiter::WaiterHandle};
12use reifydb_value::{
13 params::Params,
14 value::{Value, duration::Duration, identity::IdentityId},
15};
16use tracing::instrument;
17
18use crate::{
19 engine::StandardEngine,
20 queue::{lookup::find_queue_id, wake::QueueWakeRegistry},
21};
22
23const CLAIM_RQL: &str = "CALL queue::claim($worker, $queue, $max_n, $lease_ttl)";
24
25pub struct Session {
26 engine: StandardEngine,
27 identity: IdentityId,
28 authenticated: bool,
29 token: Option<String>,
30 retry: RetryStrategy,
31}
32
33impl Session {
34 pub fn from_token(engine: StandardEngine, info: &Token) -> Self {
35 Self {
36 engine,
37 identity: info.identity,
38 authenticated: true,
39 token: None,
40 retry: RetryStrategy::default(),
41 }
42 }
43
44 pub fn from_token_with_value(engine: StandardEngine, info: &Token) -> Self {
45 Self {
46 engine,
47 identity: info.identity,
48 authenticated: true,
49 token: Some(info.token.clone()),
50 retry: RetryStrategy::default(),
51 }
52 }
53
54 pub fn trusted(engine: StandardEngine, identity: IdentityId) -> Self {
55 Self {
56 engine,
57 identity,
58 authenticated: false,
59 token: None,
60 retry: RetryStrategy::default(),
61 }
62 }
63
64 pub fn anonymous(engine: StandardEngine) -> Self {
65 Self::trusted(engine, IdentityId::anonymous())
66 }
67
68 pub fn with_retry(mut self, strategy: RetryStrategy) -> Self {
69 self.retry = strategy;
70 self
71 }
72
73 #[inline]
74 pub fn identity(&self) -> IdentityId {
75 self.identity
76 }
77
78 #[inline]
79 pub fn token(&self) -> Option<&str> {
80 self.token.as_deref()
81 }
82
83 #[inline]
84 pub fn is_authenticated(&self) -> bool {
85 self.authenticated
86 }
87
88 #[instrument(name = "session::query", level = "debug", skip(self, params), fields(rql = %rql))]
89 pub fn query(&self, rql: &str, params: impl Into<Params>) -> ExecutionResult {
90 self.engine.query_as(self.identity, rql, params.into())
91 }
92
93 #[instrument(name = "session::command", level = "debug", skip(self, params), fields(rql = %rql))]
94 pub fn command(&self, rql: &str, params: impl Into<Params>) -> ExecutionResult {
95 let params = params.into();
96 self.retry
97 .execute(self.engine.rng(), rql, || self.engine.command_as(self.identity, rql, params.clone()))
98 }
99
100 #[instrument(name = "session::admin", level = "debug", skip(self, params), fields(rql = %rql))]
101 pub fn admin(&self, rql: &str, params: impl Into<Params>) -> ExecutionResult {
102 let params = params.into();
103 self.retry.execute(self.engine.rng(), rql, || self.engine.admin_as(self.identity, rql, params.clone()))
104 }
105
106 #[instrument(name = "queue::claim_wait", level = "debug", skip(self), fields(queue = %queue, worker = %worker))]
107 pub fn claim_wait(
108 &self,
109 queue: &str,
110 worker: &str,
111 max_n: u32,
112 lease_ttl: Duration,
113 wait_for: Duration,
114 ) -> ExecutionResult {
115 let params = claim_params(queue, worker, max_n, lease_ttl);
116
117 let mut result = self.command(CLAIM_RQL, params.clone());
118 if !wait_for.is_positive() || result.error.is_some() || claimed_any(&result) {
119 return result;
120 }
121
122 let Some(queue_id) = find_queue_id(&self.engine, self.identity, queue) else {
123 return result;
124 };
125 let registry = self.engine.queue_wake();
126 let clock = self.engine.clock();
127 let deadline = clock.instant() + wait_for;
128
129 loop {
130 let waiter = Arc::new(WaiterHandle::on_clock(clock.clone()));
131 let guard = ParkGuard::park(®istry, queue_id, waiter);
132
133 result = self.command(CLAIM_RQL, params.clone());
134 if result.error.is_some() || claimed_any(&result) {
135 guard.forward_if_consumed();
136 return result;
137 }
138
139 let Some(remaining) = remaining_budget(&clock.instant(), &deadline) else {
140 return result;
141 };
142 guard.wait(remaining);
143 }
144 }
145}
146
147struct ParkGuard<'a> {
148 registry: &'a QueueWakeRegistry,
149 queue: QueueId,
150 waiter: Arc<WaiterHandle>,
151}
152
153impl<'a> ParkGuard<'a> {
154 fn park(registry: &'a QueueWakeRegistry, queue: QueueId, waiter: Arc<WaiterHandle>) -> Self {
155 registry.register(queue, waiter.clone());
156 Self {
157 registry,
158 queue,
159 waiter,
160 }
161 }
162
163 fn wait(&self, timeout: Duration) {
164 self.waiter.wait_timeout(timeout);
165 }
166
167 fn forward_if_consumed(&self) {
168 if self.waiter.wait_timeout(Duration::zero()) {
169 self.registry.nudge(self.queue, 1);
170 }
171 }
172}
173
174impl Drop for ParkGuard<'_> {
175 fn drop(&mut self) {
176 self.registry.deregister(self.queue, &self.waiter);
177 }
178}
179
180fn claim_params(queue: &str, worker: &str, max_n: u32, lease_ttl: Duration) -> Params {
181 Params::Named(Arc::new(HashMap::from_iter([
182 ("worker".to_string(), Value::Utf8(worker.to_string())),
183 ("queue".to_string(), Value::Utf8(queue.to_string())),
184 ("max_n".to_string(), Value::Uint4(max_n)),
185 ("lease_ttl".to_string(), Value::Duration(lease_ttl)),
186 ])))
187}
188
189fn claimed_any(result: &ExecutionResult) -> bool {
190 result.frames.iter().any(|frame| frame.row_count() > 0)
191}
192
193fn remaining_budget(now: &Instant, deadline: &Instant) -> Option<Duration> {
194 if now >= deadline {
195 return None;
196 }
197 Duration::from_nanoseconds(deadline.duration_since(now).as_nanos().min(i64::MAX as u128) as i64).ok()
198}