Skip to main content

reifydb_sub_server/
dispatch.rs

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