vgi_rpc/
stream.rs

1//! Streaming primitives: OutputCollector, ProducerState, ExchangeState, StreamResult.
2
3use std::sync::Arc;
4
5use arrow_array::RecordBatch;
6use arrow_schema::{Schema, SchemaRef};
7
8use crate::errors::{Result, RpcError};
9use crate::log::{LogLevel, LogMessage};
10use crate::wire::Metadata;
11
12/// An entry in the output collector — either a data batch or a pending log.
13pub(crate) enum Emitted {
14    Batch {
15        batch: RecordBatch,
16        metadata: Option<Metadata>,
17    },
18    Log(LogMessage),
19}
20
21/// Accumulates batches and log messages for one streaming iteration.
22pub struct OutputCollector {
23    schema: SchemaRef,
24    pub(crate) items: Vec<Emitted>,
25    data_emitted: bool,
26    finished: bool,
27    is_producer: bool,
28}
29
30impl OutputCollector {
31    pub(crate) fn new(schema: SchemaRef, is_producer: bool) -> Self {
32        Self {
33            schema,
34            items: Vec::new(),
35            data_emitted: false,
36            finished: false,
37            is_producer,
38        }
39    }
40
41    /// The stream's output schema.
42    pub fn schema(&self) -> SchemaRef {
43        self.schema.clone()
44    }
45
46    /// Emit a data batch. Schema must match `self.schema()` exactly.
47    pub fn emit(&mut self, batch: RecordBatch) -> Result<()> {
48        self.ensure_data_slot()?;
49        if batch.schema() != self.schema {
50            return Err(RpcError::runtime_error(format!(
51                "emit(): schema mismatch — expected {:?}, got {:?}",
52                self.schema.fields(),
53                batch.schema().fields()
54            )));
55        }
56        self.items.push(Emitted::Batch {
57            batch,
58            metadata: None,
59        });
60        self.data_emitted = true;
61        Ok(())
62    }
63
64    /// Emit a data batch with per-batch custom metadata (e.g. VGI's
65    /// `vgi_batch_index` / `vgi_partition_values#b64` ordering tags).
66    pub fn emit_with_metadata(&mut self, batch: RecordBatch, metadata: Metadata) -> Result<()> {
67        self.ensure_data_slot()?;
68        if batch.schema() != self.schema {
69            return Err(RpcError::runtime_error(format!(
70                "emit_with_metadata(): schema mismatch — expected {:?}, got {:?}",
71                self.schema.fields(),
72                batch.schema().fields()
73            )));
74        }
75        self.items.push(Emitted::Batch {
76            batch,
77            metadata: Some(metadata),
78        });
79        self.data_emitted = true;
80        Ok(())
81    }
82
83    fn ensure_data_slot(&self) -> Result<()> {
84        if self.data_emitted {
85            return Err(RpcError::protocol_error(
86                "only one data batch may be emitted per stream turn",
87            ));
88        }
89        Ok(())
90    }
91
92    /// Mark the stream as finished (producer only).
93    pub fn finish(&mut self) {
94        self.finished = true;
95    }
96
97    pub fn finished(&self) -> bool {
98        self.finished
99    }
100
101    /// Append a client-directed log message.
102    pub fn client_log(&mut self, level: LogLevel, message: impl Into<String>) {
103        self.items
104            .push(Emitted::Log(LogMessage::new(level, message)));
105    }
106
107    /// Append a client-directed log message with extras.
108    pub fn client_log_with(&mut self, msg: LogMessage) {
109        self.items.push(Emitted::Log(msg));
110    }
111
112    pub fn is_producer(&self) -> bool {
113        self.is_producer
114    }
115}
116
117/// Server-driven producer state — called once per tick to emit at most one data batch.
118pub trait ProducerState: Send {
119    fn produce(&mut self, out: &mut OutputCollector, ctx: &CallContext) -> Result<()>;
120
121    /// Optional cancel hook — invoked when the client signals cancellation.
122    fn on_cancel(&mut self, _ctx: &CallContext) {}
123
124    /// Serialize this state for stateless HTTP continuation. The default
125    /// returns an error; override via [`crate::stream_codec::StreamStateCodec`]
126    /// for any state type that will be served over HTTP. Pipe/unix
127    /// transports never call this.
128    fn encode_state(&self) -> Result<Vec<u8>> {
129        Err(RpcError::runtime_error(
130            "producer state does not implement encode_state(); \
131             override this method or register the method via MethodInfo::stream_with_codec",
132        ))
133    }
134}
135
136/// Bidirectional exchange state — called once per client input batch.
137pub trait ExchangeState: Send {
138    fn exchange(
139        &mut self,
140        input: &RecordBatch,
141        out: &mut OutputCollector,
142        ctx: &CallContext,
143    ) -> Result<()>;
144
145    fn on_cancel(&mut self, _ctx: &CallContext) {}
146
147    /// Serialize this state for stateless HTTP continuation. See
148    /// [`ProducerState::encode_state`].
149    fn encode_state(&self) -> Result<Vec<u8>> {
150        Err(RpcError::runtime_error(
151            "exchange state does not implement encode_state(); \
152             override this method or register the method via MethodInfo::stream_with_codec",
153        ))
154    }
155}
156
157/// What a streaming method returns after init: its output/input schemas,
158/// an optional header, and the state object.
159pub struct StreamResult {
160    pub output_schema: SchemaRef,
161    /// `None` for producer streams, or a schema for exchange streams.
162    pub input_schema: Option<SchemaRef>,
163    pub state: StreamStateKind,
164    /// Optional 1-row header batch produced at stream start.
165    pub header: Option<RecordBatch>,
166    /// Arbitrary metadata to attach to the header batch.
167    pub header_metadata: Option<Metadata>,
168}
169
170pub enum StreamStateKind {
171    Producer(Box<dyn ProducerState>),
172    Exchange(Box<dyn ExchangeState>),
173}
174
175impl StreamResult {
176    pub fn producer(schema: SchemaRef, state: Box<dyn ProducerState>) -> Self {
177        Self {
178            output_schema: schema,
179            input_schema: None,
180            state: StreamStateKind::Producer(state),
181            header: None,
182            header_metadata: None,
183        }
184    }
185
186    pub fn exchange(
187        output_schema: SchemaRef,
188        input_schema: SchemaRef,
189        state: Box<dyn ExchangeState>,
190    ) -> Self {
191        Self {
192            output_schema,
193            input_schema: Some(input_schema),
194            state: StreamStateKind::Exchange(state),
195            header: None,
196            header_metadata: None,
197        }
198    }
199
200    pub fn with_header(mut self, header: RecordBatch) -> Self {
201        self.header = Some(header);
202        self
203    }
204}
205
206/// Build a [`crate::server::StateDecoder`] for a `ProducerState` that
207/// also implements [`crate::stream_codec::StreamStateCodec`].
208///
209/// **Internal:** invoked by the `#[producer]` macro expansion; user
210/// code should not call this directly.
211// Only needs `StreamStateCodec` (the `stream-codec` feature), not the http stack.
212// Gating on `stream-codec` lets producer/exchange workers build for wasm (no tokio).
213#[cfg(feature = "stream-codec")]
214#[doc(hidden)]
215pub fn producer_decoder<S>() -> crate::server::StateDecoder
216where
217    S: ProducerState + crate::stream_codec::StreamStateCodec + 'static,
218{
219    Arc::new(|bytes: &[u8]| Ok(StreamStateKind::Producer(Box::new(S::decode(bytes)?))))
220}
221
222/// Build a [`crate::server::StateDecoder`] for an `ExchangeState`. See
223/// [`producer_decoder`].
224///
225/// **Internal:** invoked by the `#[exchange]` macro expansion.
226#[cfg(feature = "stream-codec")]
227#[doc(hidden)]
228pub fn exchange_decoder<S>() -> crate::server::StateDecoder
229where
230    S: ExchangeState + crate::stream_codec::StreamStateCodec + 'static,
231{
232    Arc::new(|bytes: &[u8]| Ok(StreamStateKind::Exchange(Box::new(S::decode(bytes)?))))
233}
234
235pub(crate) fn empty_schema() -> SchemaRef {
236    Arc::new(Schema::empty())
237}
238
239// Re-export for trait bounds below.
240pub use crate::server::CallContext;
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245    use arrow_array::Int64Array;
246    use arrow_schema::{DataType, Field};
247
248    fn batch(schema: SchemaRef, value: i64) -> RecordBatch {
249        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![value]))]).unwrap()
250    }
251
252    #[test]
253    fn collector_rejects_a_second_data_batch_as_protocol_error() {
254        let schema = Arc::new(Schema::new(vec![Field::new(
255            "value",
256            DataType::Int64,
257            false,
258        )]));
259        let mut out = OutputCollector::new(schema.clone(), true);
260
261        out.emit(batch(schema.clone(), 1)).unwrap();
262        let err = out
263            .emit_with_metadata(batch(schema, 2), Metadata::default())
264            .unwrap_err();
265
266        assert_eq!(err.error_type, "ProtocolError");
267        assert_eq!(out.items.len(), 1);
268    }
269
270    #[test]
271    fn collector_allows_logs_after_data() {
272        let schema = Arc::new(Schema::new(vec![Field::new(
273            "value",
274            DataType::Int64,
275            false,
276        )]));
277        let mut out = OutputCollector::new(schema.clone(), true);
278
279        out.emit(batch(schema, 1)).unwrap();
280        out.client_log(LogLevel::Info, "still allowed");
281
282        assert_eq!(out.items.len(), 2);
283    }
284}
vgi_rpc/stream.rs

vgi_rpc/
stream.rs
