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::<B::NativeEvent>()
233                    .await
234                    .map_err(super::codec_error)
235            };
236            let result = match next {
237                Err(error) => Err(error),
238                Ok(Some(NativeFrame::Done)) => self.source_done(),
239                Ok(Some(NativeFrame::Event { value, .. })) => {
240                    self.last_native_event = Some(value.clone());
241                    self.pending_native = Some(value);
242                    Ok(())
243                }
244                Ok(None) => self.finish_source(),
245            };
246            if let Err(error) = result {
247                self.failed = true;
248                self.reader = None;
249                return Err(error);
250            }
251        }
252    }
253    fn image_resource_error(&mut self, error: TransformError) -> TransformError {
254        use crate::transform::TransformErrorKind;
255        if matches!(
256            error.kind(),
257            TransformErrorKind::InvalidInput
258                | TransformErrorKind::InvalidResult
259                | TransformErrorKind::Unsupported
260                | TransformErrorKind::Limit
261        ) {
262            self.failed = true;
263            self.reader = None;
264        }
265        error
266    }
267    fn source_event(
268        &mut self,
269        original: B::NativeEvent,
270        event: B::NativeEvent,
271    ) -> Result<(), TransformError> {
272        self.last_native_event = Some(original.clone());
273        B::NativeEvent::collect(
274            self.source
275                .as_mut()
276                .ok_or_else(|| super::invalid("native collector consumed"))?,
277            original,
278        )?;
279        let bridge = self
280            .bridge
281            .as_mut()
282            .ok_or_else(|| super::invalid("pair stream consumed"))?;
283        let converted = bridge.push(event)?;
284        self.flow = bridge.identities().clone();
285        self.append_report(converted.report)?;
286        self.pending_converted = Some(converted.value);
287        Ok(())
288    }
289    fn source_done(&mut self) -> Result<(), TransformError> {
290        if !B::NativeEvent::DONE {
291            return Err(super::invalid(
292                "DONE is not a terminator for this source dialect",
293            ));
294        }
295        B::NativeEvent::collect_done(
296            self.source
297                .as_mut()
298                .ok_or_else(|| super::invalid("native collector consumed"))?,
299        )?;
300        self.bridge
301            .as_mut()
302            .ok_or_else(|| super::invalid("pair stream consumed"))?
303            .push_done()
304    }
305    fn enqueue(&mut self, events: Vec<B::ClientEvent>) -> Result<(), TransformError> {
306        for event in events {
307            if !self.emit_usage && event.is_usage_only() {
308                continue;
309            }
310            let bytes =
311                codec::encode_json(&event, self.settings.codec).map_err(super::codec_error)?;
312            let size = bytes.len() as u64;
313            self.queued_bytes = self
314                .queued_bytes
315                .checked_add(size)
316                .ok_or_else(|| super::limit("queued client event size overflow"))?;
317            if self.queued_bytes > self.settings.codec.max_buffer_bytes
318                || self.queued_bytes > self.settings.events.max_pending_bytes as u64
319            {
320                return Err(super::limit("pending client events exceed stream budget"));
321            }
322            self.queued.push_back((event, size));
323        }
324        Ok(())
325    }
326    fn prepare_event(&mut self, event: B::ClientEvent) -> Result<(), TransformError> {
327        if !self.eof {
328            B::ClientEvent::collect(
329                self.client
330                    .as_mut()
331                    .ok_or_else(|| super::invalid("client collector consumed"))?,
332                event.clone(),
333            )?;
334        }
335        let bytes = self.encoder.event(&event, self.settings.codec)?;
336        self.ready = Some(ReadyChunk {
337            chunk: StreamChunk {
338                event: Some(event),
339                bytes,
340                finished: false,
341            },
342        });
343        Ok(())
344    }
345    fn finish_source(&mut self) -> Result<(), TransformError> {
346        let native = B::NativeEvent::collected(
347            self.source
348                .take()
349                .ok_or_else(|| super::invalid("native collector consumed"))?,
350        )?;
351        self.native_final = Some(native.value);
352        self.append_report(native.report)?;
353        let end = self
354            .bridge
355            .take()
356            .ok_or_else(|| super::invalid("pair stream consumed"))?
357            .finish()?;
358        self.flow = end.identities;
359        self.append_report(end.report)?;
360        self.queued.extend(std::mem::take(&mut self.held));
361        self.pending_converted = Some(end.chunks);
362        self.finishing_source = true;
363        Ok(())
364    }
365    fn finish_client(&mut self) -> Result<(), TransformError> {
366        let events: Vec<_> = self.queued.iter().map(|(e, _)| e.clone()).collect();
367        for event in events {
368            B::ClientEvent::collect(
369                self.client
370                    .as_mut()
371                    .ok_or_else(|| super::invalid("client collector consumed"))?,
372                event,
373            )?;
374        }
375        if B::ClientEvent::DONE {
376            B::ClientEvent::collect_done(
377                self.client
378                    .as_mut()
379                    .ok_or_else(|| super::invalid("client collector consumed"))?,
380            )?;
381        }
382        let client = B::ClientEvent::collected(
383            self.client
384                .take()
385                .ok_or_else(|| super::invalid("client collector consumed"))?,
386        )?;
387        self.client_final = Some(client.value.value);
388        self.append_report(client.report)?;
389        self.eof = true;
390        Ok(())
391    }
392    fn append_report(&mut self, report: Report) -> Result<(), TransformError> {
393        for diagnostic in report.diagnostics {
394            if !self.report.diagnostics.contains(&diagnostic) {
395                self.report.diagnostics.push(diagnostic);
396            }
397        }
398        Ok(())
399    }
400}
gproxy_protocol/adapt/generate/stream/next.rs

gproxy_protocol/adapt/generate/stream/
next.rs
