reifydb_sub_server/
dispatch.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4#[cfg(not(reifydb_single_threaded))]
5pub use native::{dispatch, dispatch_subscribe};
6
7#[cfg(not(reifydb_single_threaded))]
8mod native {
9	use std::sync::Arc;
10
11	use reifydb_core::{
12		actors::server::{ServerMessage, ServerResponse, ServerSubscribeResponse, build_server_message},
13		metric::ExecutionMetrics,
14	};
15	use reifydb_runtime::{actor::reply::reply_channel, context::clock::Instant};
16	use reifydb_type::value::{duration::Duration as ReifyDuration, 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", 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(), 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			reply,
104		};
105		let (actor_ref, _handle) = state.spawn_server_actor();
106		actor_ref.send(msg).ok().ok_or(ExecuteError::Disconnected)?;
107		timeout(state.query_timeout(), receiver.recv())
108			.await
109			.map_err(|_| ExecuteError::Timeout)?
110			.map_err(|_| ExecuteError::Disconnected)
111	}
112
113	#[inline]
114	fn finalize_dispatch_metrics(
115		response: ServerResponse,
116		start: Instant,
117	) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
118		let wall_duration = start.elapsed();
119		let (frames, compute_duration, mut metrics) = match response {
120			ServerResponse::Success {
121				frames,
122				duration,
123				metrics,
124			} => (frames, duration, metrics),
125			ServerResponse::EngineError {
126				diagnostic,
127				rql,
128			} => {
129				return Err(ExecuteError::Engine {
130					diagnostic: Arc::from(diagnostic),
131					rql,
132				});
133			}
134		};
135		metrics.total = ReifyDuration::from_nanoseconds(wall_duration.as_nanos() as i64).unwrap_or_default();
136		metrics.compute =
137			ReifyDuration::from_nanoseconds(compute_duration.as_nanos() as i64).unwrap_or_default();
138		Ok((frames, metrics))
139	}
140
141	#[inline]
142	fn finalize_subscribe_metrics(
143		response: ServerSubscribeResponse,
144		start: Instant,
145	) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
146		let wall_duration = start.elapsed();
147		let (frames, compute_duration, mut metrics) = match response {
148			ServerSubscribeResponse::Subscribed {
149				frames,
150				duration,
151				metrics,
152			} => (frames, duration, metrics),
153			ServerSubscribeResponse::EngineError {
154				diagnostic,
155				rql,
156			} => {
157				return Err(ExecuteError::Engine {
158					diagnostic: Arc::from(diagnostic),
159					rql,
160				});
161			}
162		};
163		metrics.total = ReifyDuration::from_nanoseconds(wall_duration.as_nanos() as i64).unwrap_or_default();
164		metrics.compute =
165			ReifyDuration::from_nanoseconds(compute_duration.as_nanos() as i64).unwrap_or_default();
166		Ok((frames, metrics))
167	}
168}
reifydb_sub_server/dispatch.rs

reifydb_sub_server/
dispatch.rs
