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