Skip to main content

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
19	use crate::{
20		execute::ExecuteError,
21		interceptor::{RequestContext, ResponseContext},
22		state::AppState,
23	};
24
25	/// Dispatch a query/command/admin operation through the actor with interceptors.
26	///
27	/// This is the single entry point for all native transport handlers.
28	/// The caller is responsible only for:
29	/// 1. Extracting identity from transport-specific auth
30	/// 2. Building `RequestMetadata` from transport-specific headers
31	/// 3. Parsing params from transport-specific wire format
32	/// 4. Converting the result into transport-specific response format
33	pub async fn dispatch(
34		state: &AppState,
35		mut ctx: RequestContext,
36	) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
37		run_pre_execute(state, &mut ctx).await?;
38		let start = state.clock().instant();
39		let response = send_server_message(state, &ctx).await?;
40		let (frames, metrics) = finalize_dispatch_metrics(response, start)?;
41		run_post_execute(state, &ctx, &metrics, frames.len()).await;
42		Ok((frames, metrics))
43	}
44
45	/// Dispatch a subscribe operation through the actor with interceptors.
46	///
47	/// Separate from `dispatch()` because Subscribe uses `Reply<ServerSubscribeResponse>`
48	/// rather than `Reply<ServerResponse>`.
49	pub async fn dispatch_subscribe(
50		state: &AppState,
51		mut ctx: RequestContext,
52	) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
53		run_pre_execute(state, &mut ctx).await?;
54		let start = state.clock().instant();
55		let response = send_subscribe_message(state, &ctx).await?;
56		let (frames, metrics) = finalize_subscribe_metrics(response, start)?;
57		run_post_execute(state, &ctx, &metrics, frames.len()).await;
58		Ok((frames, metrics))
59	}
60
61	#[inline]
62	async fn run_pre_execute(state: &AppState, ctx: &mut RequestContext) -> Result<(), ExecuteError> {
63		if !state.request_interceptors().is_empty() {
64			state.request_interceptors().pre_execute(ctx).await?;
65		}
66		Ok(())
67	}
68
69	#[inline]
70	async fn run_post_execute(
71		state: &AppState,
72		ctx: &RequestContext,
73		metrics: &ExecutionMetrics,
74		frame_count: usize,
75	) {
76		if state.request_interceptors().is_empty() {
77			return;
78		}
79		let response_ctx = ResponseContext {
80			identity: ctx.identity,
81			operation: ctx.operation,
82			rql: ctx.rql.clone(),
83			params: ctx.params.clone(),
84			metadata: ctx.metadata.clone(),
85			metrics: metrics.clone(),
86			result: Ok(frame_count),
87		};
88		state.request_interceptors().post_execute(&response_ctx).await;
89	}
90
91	#[inline]
92	async fn send_server_message(state: &AppState, ctx: &RequestContext) -> Result<ServerResponse, ExecuteError> {
93		let (reply, receiver) = reply_channel();
94		let msg = build_server_message(ctx.operation, ctx.identity, ctx.rql.clone(), ctx.params.clone(), reply);
95		let (actor_ref, _handle) = state.spawn_server_actor();
96		actor_ref.send(msg).ok().ok_or(ExecuteError::Disconnected)?;
97		timeout(state.query_timeout(), receiver.recv())
98			.await
99			.map_err(|_| ExecuteError::Timeout)?
100			.map_err(|_| ExecuteError::Disconnected)
101	}
102
103	#[inline]
104	async fn send_subscribe_message(
105		state: &AppState,
106		ctx: &RequestContext,
107	) -> Result<ServerSubscribeResponse, ExecuteError> {
108		let (reply, receiver) = reply_channel();
109		let msg = ServerMessage::Subscribe {
110			identity: ctx.identity,
111			rql: ctx.rql.clone(),
112			reply,
113		};
114		let (actor_ref, _handle) = state.spawn_server_actor();
115		actor_ref.send(msg).ok().ok_or(ExecuteError::Disconnected)?;
116		timeout(state.query_timeout(), receiver.recv())
117			.await
118			.map_err(|_| ExecuteError::Timeout)?
119			.map_err(|_| ExecuteError::Disconnected)
120	}
121
122	#[inline]
123	fn finalize_dispatch_metrics(
124		response: ServerResponse,
125		start: Instant,
126	) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
127		let wall_duration = start.elapsed();
128		let (frames, compute_duration, mut metrics) = match response {
129			ServerResponse::Success {
130				frames,
131				duration,
132				metrics,
133			} => (frames, duration, metrics),
134			ServerResponse::EngineError {
135				diagnostic,
136				rql,
137			} => {
138				return Err(ExecuteError::Engine {
139					diagnostic: Arc::from(diagnostic),
140					rql,
141				});
142			}
143		};
144		metrics.total = ReifyDuration::from_nanoseconds(wall_duration.as_nanos() as i64).unwrap_or_default();
145		metrics.compute =
146			ReifyDuration::from_nanoseconds(compute_duration.as_nanos() as i64).unwrap_or_default();
147		Ok((frames, metrics))
148	}
149
150	#[inline]
151	fn finalize_subscribe_metrics(
152		response: ServerSubscribeResponse,
153		start: Instant,
154	) -> Result<(Vec<Frame>, ExecutionMetrics), ExecuteError> {
155		let wall_duration = start.elapsed();
156		let (frames, compute_duration, mut metrics) = match response {
157			ServerSubscribeResponse::Subscribed {
158				frames,
159				duration,
160				metrics,
161			} => (frames, duration, metrics),
162			ServerSubscribeResponse::EngineError {
163				diagnostic,
164				rql,
165			} => {
166				return Err(ExecuteError::Engine {
167					diagnostic: Arc::from(diagnostic),
168					rql,
169				});
170			}
171		};
172		metrics.total = ReifyDuration::from_nanoseconds(wall_duration.as_nanos() as i64).unwrap_or_default();
173		metrics.compute =
174			ReifyDuration::from_nanoseconds(compute_duration.as_nanos() as i64).unwrap_or_default();
175		Ok((frames, metrics))
176	}
177}