1#[cfg(not(reifydb_single_threaded))]
5pub use host::{dispatch, dispatch_subscribe};
6
7#[cfg(not(reifydb_single_threaded))]
8mod host {
9 use std::sync::Arc;
10
11 use reifydb_core::{
12 actors::server::{ServerMessage, ServerResponse, ServerSubscribeResponse, build_server_message},
13 interface::catalog::subscription::{SubscribeOptions, SubscribeOutcome},
14 metrics::execution::ExecutionMetrics,
15 };
16 use reifydb_runtime::{actor::reply::reply_channel, context::clock::Instant};
17 use reifydb_value::{
18 error::{IntoDiagnostic, TypeError},
19 value::{duration::Duration, frame::frame::Frame},
20 };
21 use tokio::time::timeout;
22 use tracing::instrument;
23
24 use crate::{
25 execute::ExecuteError,
26 interceptor::{RequestContext, ResponseContext},
27 state::AppState,
28 };
29
30 #[instrument(name = "dispatch::request", level = "debug", skip_all, fields(op = ?ctx.operation))]
31 pub async fn dispatch(
32 state: &AppState,
33 mut ctx: RequestContext,
34 ) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
35 run_pre_execute(state, &mut ctx).await?;
36 let start = state.clock().instant();
37 let response = send_server_message(state, &ctx).await?;
38 let (frames, metrics) = finalize_dispatch_metrics(response, start)?;
39 run_post_execute(state, &ctx, &metrics, frames.len()).await;
40 Ok((frames, metrics))
41 }
42
43 #[instrument(name = "dispatch::subscribe", level = "debug", skip_all, fields(op = ?ctx.operation))]
44 pub async fn dispatch_subscribe(
45 state: &AppState,
46 mut ctx: RequestContext,
47 options: SubscribeOptions,
48 ) -> Result<SubscribeOutcome, ExecuteError> {
49 run_pre_execute(state, &mut ctx).await?;
50 let start = state.clock().instant();
51 let response = send_subscribe_message(state, &ctx, options).await?;
52 let (outcome, metrics) = finalize_subscribe_metrics(response, start)?;
53 run_post_execute(state, &ctx, &metrics, 0).await;
54 Ok(outcome)
55 }
56
57 #[inline]
58 async fn run_pre_execute(state: &AppState, ctx: &mut RequestContext) -> Result<(), ExecuteError> {
59 if !state.request_interceptors().is_empty() {
60 state.request_interceptors().pre_execute(ctx).await?;
61 }
62 Ok(())
63 }
64
65 #[inline]
66 async fn run_post_execute(
67 state: &AppState,
68 ctx: &RequestContext,
69 metrics: &ExecutionMetrics,
70 frame_count: usize,
71 ) {
72 if state.request_interceptors().is_empty() {
73 return;
74 }
75 let response_ctx = ResponseContext {
76 identity: ctx.identity,
77 operation: ctx.operation,
78 rql: ctx.rql.clone(),
79 params: ctx.params.clone(),
80 metadata: ctx.metadata.clone(),
81 metrics: metrics.clone(),
82 result: Ok(frame_count),
83 };
84 state.request_interceptors().post_execute(&response_ctx).await;
85 }
86
87 #[instrument(name = "dispatch::send_server_message", level = "debug", skip_all)]
88 async fn send_server_message(state: &AppState, ctx: &RequestContext) -> Result<ServerResponse, ExecuteError> {
89 let (reply, receiver) = reply_channel();
90 let msg = build_server_message(ctx.operation, ctx.identity, ctx.rql.clone(), ctx.params.clone(), reply);
91 let (actor_ref, _handle) = state.spawn_server_actor();
92 actor_ref.send(msg).ok().ok_or(ExecuteError::Disconnected)?;
93 timeout(state.query_timeout().to_std(), receiver.recv())
94 .await
95 .map_err(|_| ExecuteError::Timeout)?
96 .map_err(|_| ExecuteError::Disconnected)
97 }
98
99 #[instrument(name = "dispatch::send_subscribe_message", level = "debug", skip_all)]
100 async fn send_subscribe_message(
101 state: &AppState,
102 ctx: &RequestContext,
103 options: SubscribeOptions,
104 ) -> Result<ServerSubscribeResponse, ExecuteError> {
105 let (reply, receiver) = reply_channel();
106 let msg = ServerMessage::Subscribe {
107 identity: ctx.identity,
108 rql: ctx.rql.clone(),
109 params: ctx.params.clone(),
110 options,
111 reply,
112 };
113 let (actor_ref, _handle) = state.spawn_server_actor();
114 actor_ref.send(msg).ok().ok_or(ExecuteError::Disconnected)?;
115 timeout(state.query_timeout().to_std(), receiver.recv())
116 .await
117 .map_err(|_| ExecuteError::Timeout)?
118 .map_err(|_| ExecuteError::Disconnected)
119 }
120
121 #[inline]
122 fn finalize_dispatch_metrics(
123 response: ServerResponse,
124 start: Instant,
125 ) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
126 let wall_duration = start.elapsed();
127 let (frames, compute_duration, mut metrics) = match response {
128 ServerResponse::Success {
129 frames,
130 duration,
131 metrics,
132 } => (frames, duration, metrics),
133 ServerResponse::EngineError {
134 diagnostic,
135 rql,
136 } => {
137 return Err(ExecuteError::Engine {
138 diagnostic: Arc::from(diagnostic),
139 rql,
140 });
141 }
142 };
143 metrics.total =
144 Duration::from_nanoseconds(wall_duration.as_nanos() as i64).map_err(|e| duration_error(*e))?;
145 metrics.compute = Duration::from_nanoseconds(compute_duration.to_std().as_nanos() as i64)
146 .map_err(|e| duration_error(*e))?;
147 Ok((frames, metrics))
148 }
149
150 #[inline]
151 fn finalize_subscribe_metrics(
152 response: ServerSubscribeResponse,
153 start: Instant,
154 ) -> Result<(SubscribeOutcome, ExecutionMetrics), ExecuteError> {
155 let wall_duration = start.elapsed();
156 let (outcome, compute_duration) = match response {
157 ServerSubscribeResponse::Subscribed {
158 outcome,
159 duration,
160 } => (outcome, duration),
161 ServerSubscribeResponse::EngineError {
162 diagnostic,
163 rql,
164 } => {
165 return Err(ExecuteError::Engine {
166 diagnostic: Arc::from(diagnostic),
167 rql,
168 });
169 }
170 };
171 let metrics = ExecutionMetrics {
172 total: Duration::from_nanoseconds(wall_duration.as_nanos() as i64)
173 .map_err(|e| duration_error(*e))?,
174 compute: Duration::from_nanoseconds(compute_duration.to_std().as_nanos() as i64)
175 .map_err(|e| duration_error(*e))?,
176 ..Default::default()
177 };
178 Ok((outcome, metrics))
179 }
180
181 fn duration_error(error: TypeError) -> ExecuteError {
182 ExecuteError::Engine {
183 diagnostic: Arc::from(error.into_diagnostic()),
184 rql: String::new(),
185 }
186 }
187}