gproxy_protocol/adapt/generate/stream/
next.rs

1use super::super::GenerationStateAccess;
2use super::{
3    StreamChunk, StreamInvocation, bridge::StreamBridge, event::NativeEvent, invoke::ReadyChunk,
4    reader::NativeFrame,
5};
6use crate::{
7    capability::StateStore,
8    codec,
9    transform::{Report, TransformError},
10};
11
12impl<B: StreamBridge> StreamInvocation<B> {
13    /// Collect the mapped stream into its concrete complete client wire response.
14    /// The host calls `start` once first. This does not perform another POST.
15    pub async fn collect<S: StateStore>(
16        &mut self,
17        state: &GenerationStateAccess<'_, S>,
18    ) -> Result<super::super::GenerationOutcome<super::invoke::ClientFull<B>>, TransformError> {
19        self.binding.check(state)?;
20        if let Some(raw) = &self.rejected {
21            return Ok(super::super::GenerationOutcome::Rejected(
22                crate::WireResponse {
23                    status: raw.status,
24                    headers: raw.headers.clone(),
25                    body: raw.body.clone(),
26                },
27            ));
28        }
29        while self.next(state).await?.is_some() {}
30        let body = self
31            .client_result()
32            .ok_or_else(|| super::missing("client final response is not durably committed"))?
33            .clone();
34        codec::encode_json(&body, self.settings.codec).map_err(super::codec_error)?;
35        let metadata = self
36            .metadata
37            .as_ref()
38            .ok_or_else(|| super::missing("native HTTP metadata unavailable"))?;
39        Ok(super::super::GenerationOutcome::Success {
40            response: crate::WireResponse {
41                status: metadata.status,
42                headers: super::output::headers(
43                    metadata.headers.clone(),
44                    super::reader::SourceFraming::JsonArray,
45                ),
46                body,
47            },
48            report: self.report.clone(),
49        })
50    }
51
52    /// Returns the next encoded client event. No alias is persisted before it
53    /// is yielded: an emitted ID is either the upstream's own or a
54    /// per-request alias, and the only state writes happen once, at the end of
55    /// the stream. Cancellation retains a prepared event.
56    pub async fn next<S: StateStore>(
57        &mut self,
58        state: &GenerationStateAccess<'_, S>,
59    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
60        if self.websocket_terminal {
61            self.next_inner(
62                state,
63                Some(&mut futures_util::stream::poll_fn(|_| {
64                    std::task::Poll::Ready(None)
65                })),
66            )
67            .await
68        } else {
69            self.next_inner(state, None).await
70        }
71    }
72    pub(super) async fn next_inner<S: StateStore>(
73        &mut self,
74        state: &GenerationStateAccess<'_, S>,
75        external: Option<&mut super::resource_map::ExternalSource<'_, B::NativeEvent>>,
76    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
77        self.next_mapped(state, external, None).await
78    }
79    pub(super) async fn next_mapped<S: StateStore>(
80        &mut self,
81        state: &GenerationStateAccess<'_, S>,
82        mut external: Option<&mut super::resource_map::ExternalSource<'_, B::NativeEvent>>,
83        mut mapping: Option<&mut dyn super::resource_map::ResourceMapping<B, S>>,
84    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
85        use futures_util::StreamExt;
86        if self.image_resources_required && mapping.is_none() {
87            return Err(super::missing(
88                "continue this invocation with its image resource progress",
89            ));
90        }
91
92        self.binding.check(state)?;
93        if mapping
94            .as_ref()
95            .is_some_and(|mapper| mapper.revision() != self.resource_revision)
96        {
97            return Err(super::conflict(
98                "image resource progress changed during invocation",
99            ));
100        }
101        if self.failed {
102            return Err(super::invalid("stream invocation failed"));
103        }
104        if self.finished {
105            return Ok(None);
106        }
107        if self.reader.is_none() && external.is_none() {
108            return Err(super::missing("successful stream start required"));
109        }
110        loop {
111            if let Some(original) = self.pending_native.as_ref() {
112                let view = if let Some(mapper) = mapping.as_deref_mut() {
113                    self.resource_revision = mapper.begin_step()?;
114                    let value = match mapper.source(original).await {
115                        Ok(value) => value,
116                        Err(error) => {
117                            if mapper.failed() {
118                                self.failed = true;
119                                self.reader = None;
120                            }
121                            return Err(self.image_resource_error(error));
122                        }
123                    };
124                    self.resource_revision = mapper.revision();
125                    value
126                } else {
127                    original.clone()
128                };
129                let original = self.pending_native.take().expect("retained native event");
130                let terminal = external.is_some() && original.is_terminal();
131                if let Err(error) = self.source_event(original, view) {
132                    self.failed = true;
133                    self.reader = None;
134                    return Err(error);
135                }
136                if terminal {
137                    self.websocket_terminal = true;
138                }
139            }
140            if let Some(events) = self.pending_converted.as_ref() {
141                let mapped = if let Some(mapper) = mapping.as_deref_mut() {
142                    self.resource_revision = mapper.begin_step()?;
143                    let value = match mapper.client(events).await {
144                        Ok(value) => value,
145                        Err(error) => {
146                            if mapper.failed() {
147                                self.failed = true;
148                                self.reader = None;
149                            }
150                            return Err(self.image_resource_error(error));
151                        }
152                    };
153                    self.resource_revision = mapper.revision();
154                    value
155                } else {
156                    events.clone()
157                };
158                self.pending_converted = None;
159                if let Err(error) = self.enqueue(mapped) {
160                    self.failed = true;
161                    self.reader = None;
162                    return Err(error);
163                }
164            }
165            if self.finishing_source {
166                if let Err(error) = self.finish_client() {
167                    self.failed = true;
168                    self.reader = None;
169                    return Err(error);
170                }
171                self.finishing_source = false;
172            }
173            if self.eof && !self.final_saved {
174                if let Some(history) = &mut self.history {
175                    let response = B::ClientEvent::responses_history(
176                        self.client_final.as_ref().expect("final response checked"),
177                    )
178                    .ok_or_else(|| {
179                        super::invalid("Responses history attached to another dialect")
180                    })?;
181                    history.save(response, state, self.settings.codec).await?;
182                }
183                self.final_saved = true;
184            }
185            if self.ready.is_some() {
186                let ready = self.ready.take().expect("pending encoded event retained");
187                self.finished = ready.chunk.finished;
188                return Ok(Some(ready.chunk));
189            }
190            if let Some((event, size)) = self.queued.pop_front() {
191                if !self.eof && (self.holding || event.is_terminal()) {
192                    self.holding = true;
193                    self.held.push((event, size));
194                    continue;
195                }
196                self.queued_bytes = self.queued_bytes.saturating_sub(size);
197                if let Err(error) = self.prepare_event(event) {
198                    self.failed = true;
199                    self.reader = None;
200                    return Err(error);
201                }
202                continue;
203            }
204            if self.eof {
205                let bytes = match self.encoder.finish::<B::ClientEvent>() {
206                    Ok(value) => value,
207                    Err(error) => {
208                        self.failed = true;
209                        self.reader = None;
210                        return Err(error);
211                    }
212                };
213                self.ready = Some(ReadyChunk {
214                    chunk: StreamChunk {
215                        event: None,
216                        bytes,
217                        finished: true,
218                    },
219                });
220                continue;
221            }
222            let next = if let Some(source) = external.as_mut() {
223                source
224                    .next()
225                    .await
226                    .transpose()
227                    .map(|event| event.map(|value| NativeFrame::Event { name: None, value }))
228            } else {
229                self.reader
230                    .as_mut()
231                    .expect("successful start")
232                    .next::<serde_json::Value>()
233                    .await
234                    .map_err(super::codec_error)
235                    .and_then(|frame| match frame {
236                        Some(NativeFrame::Event { name, value }) => B::NativeEvent::decode(value)
237                            .map(|value| Some(NativeFrame::Event { name, value })),
238                        Some(NativeFrame::Done) => Ok(Some(NativeFrame::Done)),
239                        None => Ok(None),
240                    })
241            };
242            let result = match next {
243                Err(error) => Err(error),
244                Ok(Some(NativeFrame::Done)) => self.source_done(),
245                Ok(Some(NativeFrame::Event { value, .. })) => {
246                    self.last_native_event = Some(value.clone());
247                    self.pending_native = Some(value);
248                    Ok(())
249                }
250                Ok(None) => self.finish_source(),
251            };
252            if let Err(error) = result {
253                self.failed = true;
254                self.reader = None;
255                return Err(error);
256            }
257        }
258    }
259    fn image_resource_error(&mut self, error: TransformError) -> TransformError {
260        use crate::transform::TransformErrorKind;
261        if matches!(
262            error.kind(),
263            TransformErrorKind::InvalidInput
264                | TransformErrorKind::InvalidResult
265                | TransformErrorKind::Unsupported
266                | TransformErrorKind::Limit
267        ) {
268            self.failed = true;
269            self.reader = None;
270        }
271        error
272    }
273    fn source_event(
274        &mut self,
275        original: B::NativeEvent,
276        event: B::NativeEvent,
277    ) -> Result<(), TransformError> {
278        self.last_native_event = Some(original.clone());
279        B::NativeEvent::collect(
280            self.source
281                .as_mut()
282                .ok_or_else(|| super::invalid("native collector consumed"))?,
283            original,
284        )?;
285        let bridge = self
286            .bridge
287            .as_mut()
288            .ok_or_else(|| super::invalid("pair stream consumed"))?;
289        let converted = bridge.push(event)?;
290        self.flow = bridge.identities().clone();
291        self.append_report(converted.report)?;
292        self.pending_converted = Some(converted.value);
293        Ok(())
294    }
295    fn source_done(&mut self) -> Result<(), TransformError> {
296        if !B::NativeEvent::DONE {
297            return Err(super::invalid(
298                "DONE is not a terminator for this source dialect",
299            ));
300        }
301        B::NativeEvent::collect_done(
302            self.source
303                .as_mut()
304                .ok_or_else(|| super::invalid("native collector consumed"))?,
305        )?;
306        self.bridge
307            .as_mut()
308            .ok_or_else(|| super::invalid("pair stream consumed"))?
309            .push_done()
310    }
311    fn enqueue(&mut self, events: Vec<B::ClientEvent>) -> Result<(), TransformError> {
312        for event in events {
313            if !self.emit_usage && event.is_usage_only() {
314                continue;
315            }
316            let bytes =
317                codec::encode_json(&event, self.settings.codec).map_err(super::codec_error)?;
318            let size = bytes.len() as u64;
319            self.queued_bytes = self
320                .queued_bytes
321                .checked_add(size)
322                .ok_or_else(|| super::limit("queued client event size overflow"))?;
323            if self.queued_bytes > self.settings.codec.max_buffer_bytes
324                || self.queued_bytes > self.settings.events.max_pending_bytes as u64
325            {
326                return Err(super::limit("pending client events exceed stream budget"));
327            }
328            self.queued.push_back((event, size));
329        }
330        Ok(())
331    }
332    fn prepare_event(&mut self, event: B::ClientEvent) -> Result<(), TransformError> {
333        if !self.eof {
334            B::ClientEvent::collect(
335                self.client
336                    .as_mut()
337                    .ok_or_else(|| super::invalid("client collector consumed"))?,
338                event.clone(),
339            )?;
340        }
341        let bytes = self.encoder.event(&event, self.settings.codec)?;
342        self.ready = Some(ReadyChunk {
343            chunk: StreamChunk {
344                event: Some(event),
345                bytes,
346                finished: false,
347            },
348        });
349        Ok(())
350    }
351    fn finish_source(&mut self) -> Result<(), TransformError> {
352        let native = B::NativeEvent::collected(
353            *self
354                .source
355                .take()
356                .ok_or_else(|| super::invalid("native collector consumed"))?,
357        )?;
358        self.native_final = Some(native.value);
359        self.append_report(native.report)?;
360        let end = self
361            .bridge
362            .take()
363            .ok_or_else(|| super::invalid("pair stream consumed"))?
364            .finish()?;
365        self.flow = end.identities;
366        self.append_report(end.report)?;
367        self.queued.extend(std::mem::take(&mut self.held));
368        self.pending_converted = Some(end.chunks);
369        self.finishing_source = true;
370        Ok(())
371    }
372    fn finish_client(&mut self) -> Result<(), TransformError> {
373        let events: Vec<_> = self.queued.iter().map(|(e, _)| e.clone()).collect();
374        for event in events {
375            B::ClientEvent::collect(
376                self.client
377                    .as_mut()
378                    .ok_or_else(|| super::invalid("client collector consumed"))?,
379                event,
380            )?;
381        }
382        if B::ClientEvent::DONE {
383            B::ClientEvent::collect_done(
384                self.client
385                    .as_mut()
386                    .ok_or_else(|| super::invalid("client collector consumed"))?,
387            )?;
388        }
389        let client = B::ClientEvent::collected(
390            *self
391                .client
392                .take()
393                .ok_or_else(|| super::invalid("client collector consumed"))?,
394        )?;
395        self.client_final = Some(client.value.value);
396        self.append_report(client.report)?;
397        self.eof = true;
398        Ok(())
399    }
400    fn append_report(&mut self, report: Report) -> Result<(), TransformError> {
401        for diagnostic in report.diagnostics {
402            if !self.report.diagnostics.contains(&diagnostic) {
403                self.report.diagnostics.push(diagnostic);
404            }
405        }
406        Ok(())
407    }
408}
gproxy_protocol/adapt/generate/stream/next.rs

gproxy_protocol/adapt/generate/stream/
next.rs
