1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
// SPDX-License-Identifier: BUSL-1.1
//! Streaming gateway entry point: [`Gateway::execute_stream`].
//!
//! Mirrors [`Gateway::execute`](super::core::Gateway::execute)'s routing but
//! produces a merged [`ResultStream`] of row batches instead of a collected
//! `Vec<Vec<u8>>`, so rows flow to the client as they arrive.
use std::sync::Arc;
use std::sync::atomic::Ordering;
use crate::Error;
use crate::control::server::result_stream::ResultStream;
use nodedb_physical::physical_plan::PhysicalPlan;
use super::core::{Gateway, QueryContext};
use super::dispatcher::{DispatchRouteStreamParams, default_deadline_ms, dispatch_route_stream};
use super::retry::retry_not_leader;
use super::route::TaskRoute;
use super::router::resolve_decision;
impl Gateway {
/// Streaming sibling of [`execute`](Gateway::execute).
///
/// Routes `plan` the same way as `execute` (via [`Gateway::compute_routes`]),
/// but each route produces a [`ResultStream`] of row
/// batches instead of a collected `Vec<Vec<u8>>`:
///
/// - Local route → `gather_all_cores_stream` over the route's plan.
/// - Remote route → `dispatch_remote_stream` (eager first-frame + typed
/// NotLeader retry before the first row; terminal-after-first thereafter).
///
/// All per-route streams are merged with `futures::stream::select_all`, so
/// rows interleave as they arrive. The not-leader retry wraps ONLY the eager
/// pre-stream phase (opening the stream + first frame), matching the
/// retry-vs-stream contract — once a route's stream is live, its errors are
/// terminal.
///
/// The streaming path does not go through `execute` /
/// `execute_with_version_set` (which collect a materialized `Vec<Vec<u8>>`);
/// it frames rows incrementally instead.
pub async fn execute_stream(
&self,
ctx: &QueryContext,
plan: PhysicalPlan,
) -> Result<ResultStream, Error> {
let shared = self.shared()?;
let version_set =
self.collect_version_set(&plan, ctx.tenant_id.as_u64(), ctx.database_id)?;
let routes = self.compute_routes(plan, ctx)?;
let deadline_ms = default_deadline_ms(&shared);
let mut per_route: Vec<ResultStream> = Vec::with_capacity(routes.len());
for route in routes {
let vshard_id_u32 = route.vshard_id;
let plan_for_retry = route.plan.clone();
let routing_ref = shared.cluster_routing.as_deref();
let retry_counter = Arc::clone(&self.not_leader_retry_count);
let version_set_for_route = version_set.clone();
let shared_for_route = Arc::clone(&shared);
// The not-leader retry wraps only the eager pre-stream phase. Each
// attempt re-resolves the routing decision and re-opens the stream;
// a pre-row NotLeader / DescriptorMismatch is retryable, anything
// after the first frame is terminal (handled inside the stream).
let stream = retry_not_leader(routing_ref, move |attempt| {
if attempt > 0 {
retry_counter.fetch_add(1, Ordering::Relaxed);
}
let plan = plan_for_retry.clone();
let shared = Arc::clone(&shared_for_route);
let tenant_id = ctx.tenant_id;
let database_id = ctx.database_id;
let trace_id = ctx.trace_id;
let version_set = version_set_for_route.clone();
async move {
let decision = {
let routing_guard = shared
.cluster_routing
.as_ref()
.map(|rw| rw.read().unwrap_or_else(|p| p.into_inner()));
let raft_snapshot: Vec<nodedb_cluster::GroupStatus> =
shared.raft_status_fn.get().map(|f| f()).unwrap_or_default();
let live_leader = move |group_id: u64| -> u64 {
raft_snapshot
.iter()
.find(|gs| gs.group_id == group_id)
.map(|gs| gs.leader_id)
.unwrap_or(0)
};
let live_lookup: Option<&dyn Fn(u64) -> u64> =
if shared.raft_status_fn.get().is_some() {
Some(&live_leader)
} else {
None
};
resolve_decision(
vshard_id_u32,
shared.node_id,
routing_guard.as_deref(),
live_lookup,
)
};
let route = TaskRoute {
plan,
decision,
vshard_id: vshard_id_u32,
};
dispatch_route_stream(DispatchRouteStreamParams {
route,
shared: &shared,
tenant_id,
database_id,
trace_id,
deadline_ms,
version_set: &version_set,
})
.await
}
})
.await?;
per_route.push(stream);
}
Ok(Box::pin(futures::stream::select_all(per_route)))
}
}