webui_handler/streaming/
mod.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT license.
3
4//! Runtime-discovered progressive streaming hydration.
5
6mod checkpoint;
7mod error;
8mod inventory;
9mod owned;
10mod root;
11mod session;
12mod state;
13mod vm;
14
15use serde_json::Value;
16
17use crate::route_handler::Protocol;
18use crate::{
19    FlushWriter, HandlerError, RenderOptions, ResponseWriter, Result, WebUIHandler,
20    WebUIProcessContext,
21};
22
23pub(crate) use error::streaming_boundary_error;
24pub(crate) use inventory::{record_checkpoint_tag, streaming_template_already_sent};
25pub use owned::{BufferSink, SessionOptions, StreamStep, StreamingSession};
26pub(crate) use root::{
27    consume_streaming_component_root, ensure_no_pending_streaming_root,
28    prepare_generated_streaming_root, validate_pending_streaming_root,
29    validate_streaming_root_opening, ComponentHostOrigin,
30};
31pub use session::{
32    BoundaryDescriptor, BoundaryInstanceId, BoundaryKey, BoundaryMode, SpanInstanceId,
33    StreamStatus, StreamingResponse, StreamingState, MAX_BOUNDARY_OCCURRENCES,
34    MAX_CONTINUATION_DEPTH, MAX_KEYED_INSTANCES, MAX_SPAN_NESTING,
35};
36pub(crate) use state::StreamingRenderState;
37pub(crate) use vm::PreparedContinuationStatePlan;
38
39use root::process_streaming_root_signal;
40use state::{increment_streaming_record_sequence, require_streaming_head_start};
41
42pub(crate) const STREAMING_MARKER: &str = "<meta name=\"webui-streaming\" content=\"1\">";
43
44/// Fixed-capacity scratch used to assemble short streaming markers and record
45/// headers.
46///
47/// Every `ResponseWriter::write` is a dynamic call that reaches the transport,
48/// so composing a marker byte-piece by byte-piece costs several indirections
49/// per boundary. Formatting into this stack buffer collapses each marker or
50/// record header into a single write.
51pub(super) struct MarkerBuffer {
52    bytes: [u8; Self::CAPACITY],
53    len: usize,
54}
55
56impl MarkerBuffer {
57    const CAPACITY: usize = 96;
58
59    pub(super) const fn new() -> Self {
60        Self {
61            bytes: [0; Self::CAPACITY],
62            len: 0,
63        }
64    }
65
66    pub(super) fn push_str(&mut self, text: &str) -> Result<()> {
67        let end = self
68            .len
69            .checked_add(text.len())
70            .filter(|end| *end <= Self::CAPACITY)
71            .ok_or_else(marker_overflow_error)?;
72        self.bytes[self.len..end].copy_from_slice(text.as_bytes());
73        self.len = end;
74        Ok(())
75    }
76
77    pub(super) fn push_usize(&mut self, mut value: usize) -> Result<()> {
78        let mut digits = [0u8; 20];
79        let mut offset = digits.len();
80        loop {
81            offset = offset.checked_sub(1).ok_or_else(marker_overflow_error)?;
82            digits[offset] =
83                b'0' + u8::try_from(value % 10).map_err(|_| marker_overflow_error())?;
84            value /= 10;
85            if value == 0 {
86                break;
87            }
88        }
89        let end = self
90            .len
91            .checked_add(digits.len() - offset)
92            .filter(|end| *end <= Self::CAPACITY)
93            .ok_or_else(marker_overflow_error)?;
94        self.bytes[self.len..end].copy_from_slice(&digits[offset..]);
95        self.len = end;
96        Ok(())
97    }
98
99    pub(super) fn flush_to<W: ResponseWriter + ?Sized>(&mut self, writer: &mut W) -> Result<()> {
100        let text =
101            std::str::from_utf8(&self.bytes[..self.len]).map_err(|_| marker_overflow_error())?;
102        let result = writer.write(text);
103        self.len = 0;
104        result
105    }
106}
107
108/// Write `<prefix><id>-->` with a single writer call.
109pub(super) fn write_range_marker<W: ResponseWriter + ?Sized>(
110    writer: &mut W,
111    prefix: &str,
112    id: u32,
113) -> Result<()> {
114    let mut buffer = MarkerBuffer::new();
115    buffer.push_str(prefix)?;
116    buffer.push_usize(usize::try_from(id).map_err(|_| marker_overflow_error())?)?;
117    buffer.push_str("-->")?;
118    buffer.flush_to(writer)
119}
120
121#[cold]
122#[inline(never)]
123fn marker_overflow_error() -> HandlerError {
124    HandlerError::Invariant("streaming marker exceeded its fixed buffer".to_string())
125}
126
127struct StreamingSink<'w, W: FlushWriter + ?Sized> {
128    transport: &'w mut W,
129    component_opening: Option<ComponentOpening>,
130    written: usize,
131    flushed: usize,
132}
133
134struct ComponentOpening {
135    bytes: String,
136    root_offset: Option<usize>,
137}
138
139impl<W: FlushWriter + ?Sized> ResponseWriter for StreamingSink<'_, W> {
140    fn write(&mut self, content: &str) -> Result<()> {
141        if let Some(opening) = self.component_opening.as_mut() {
142            opening.bytes.push_str(content);
143            return Ok(());
144        }
145        self.written = self.written.wrapping_add(content.len());
146        self.transport.write(content)
147    }
148
149    fn write_attribute(&mut self, name: &str, value: &str) -> Result<()> {
150        if let Some(opening) = self.component_opening.as_mut() {
151            crate::response_writer::append_attribute_to_string(&mut opening.bytes, name, value);
152            return Ok(());
153        }
154        self.written = self
155            .written
156            .wrapping_add(name.len().wrapping_add(value.len()).wrapping_add(4));
157        self.transport.write_attribute(name, value)
158    }
159
160    fn write_boolean_attribute(&mut self, name: &str) -> Result<()> {
161        if let Some(opening) = self.component_opening.as_mut() {
162            crate::response_writer::append_boolean_attribute_to_string(&mut opening.bytes, name);
163            return Ok(());
164        }
165        self.written = self.written.wrapping_add(name.len().wrapping_add(1));
166        self.transport.write_boolean_attribute(name)
167    }
168
169    fn end(&mut self) -> Result<()> {
170        self.transport.end()
171    }
172
173    fn stream_flush(&mut self) -> Result<()> {
174        if self.component_opening.is_some() {
175            return Err(HandlerError::Invariant(
176                "streaming flush attempted while a component opening was buffered".to_string(),
177            ));
178        }
179        // A flush at a position already delivered would cost a syscall without
180        // releasing a byte, so semantic steps that produced no output between
181        // two checkpoints collapse into the earlier flush.
182        if self.written == self.flushed {
183            return Ok(());
184        }
185        self.flushed = self.written;
186        self.transport.flush()
187    }
188
189    fn stream_begin_component(&mut self) -> Result<()> {
190        if self.component_opening.is_some() {
191            return Err(HandlerError::Invariant(
192                "nested component opening buffers are not valid".to_string(),
193            ));
194        }
195        self.component_opening = Some(ComponentOpening {
196            bytes: String::with_capacity(128),
197            root_offset: None,
198        });
199        Ok(())
200    }
201
202    fn stream_mark_component_root(&mut self) -> Result<()> {
203        let Some(opening) = self.component_opening.as_mut() else {
204            return Err(HandlerError::Invariant(
205                "component root marker has no buffered opening".to_string(),
206            ));
207        };
208        if opening.root_offset.replace(opening.bytes.len()).is_some() {
209            return Err(HandlerError::Invariant(
210                "component opening contains duplicate root markers".to_string(),
211            ));
212        }
213        Ok(())
214    }
215
216    fn stream_commit_component(
217        &mut self,
218        span_id: Option<u32>,
219        enclosing_span_id: Option<u32>,
220        deferred: bool,
221    ) -> Result<()> {
222        let Some(opening) = self.component_opening.take() else {
223            return Err(HandlerError::Invariant(
224                "component opening commit has no buffered bytes".to_string(),
225            ));
226        };
227        let Some(root_offset) = opening.root_offset else {
228            return Err(HandlerError::Invariant(
229                "buffered component opening has no compiler root marker".to_string(),
230            ));
231        };
232        if let Some(id) = span_id {
233            write_range_marker(self, "<!--ws:", id)?;
234        }
235        self.write(&opening.bytes[..root_offset])?;
236        if deferred {
237            self.write(" data-ws")?;
238            if let Some(id) = span_id {
239                self.write(" data-ws-span=\"")?;
240                write_u32(self, id)?;
241                self.write("\"")?;
242            }
243            if let Some(id) = enclosing_span_id {
244                self.write(" data-ws-enclosing=\"")?;
245                write_u32(self, id)?;
246                self.write("\"")?;
247            }
248        }
249        self.write(&opening.bytes[root_offset..])
250    }
251}
252
253fn write_u32<W: ResponseWriter + ?Sized>(writer: &mut W, mut value: u32) -> Result<()> {
254    if value == 0 {
255        return writer.write("0");
256    }
257    let mut digits = [0u8; 10];
258    let mut offset = digits.len();
259    while value != 0 {
260        offset -= 1;
261        let digit = u8::try_from(value % 10)
262            .map_err(|_| HandlerError::Invariant("invalid decimal digit".to_string()))?;
263        digits[offset] = b'0' + digit;
264        value /= 10;
265    }
266    let value = std::str::from_utf8(&digits[offset..])
267        .map_err(|_| HandlerError::Invariant("invalid decimal ID bytes".to_string()))?;
268    writer.write(value)
269}
270
271pub(super) fn streaming_state<'a, 'data>(
272    context: &'a mut WebUIProcessContext<'data, '_, '_>,
273) -> Result<&'a mut StreamingRenderState<'data>> {
274    context.streaming.as_deref_mut().ok_or_else(|| {
275        HandlerError::Invariant("streaming signal processed outside a streaming render".to_string())
276    })
277}
278
279pub(super) fn flush_streaming_transport(
280    context: &mut WebUIProcessContext<'_, '_, '_>,
281) -> Result<()> {
282    if context.streaming.is_none() {
283        return Err(HandlerError::Invariant(
284            "streaming flush requested outside a streaming render".to_string(),
285        ));
286    }
287    context.writer.stream_flush()
288}
289
290impl WebUIHandler {
291    pub(crate) fn process_streaming_signal<'data>(
292        &self,
293        value: &'data str,
294        context: &mut WebUIProcessContext<'data, '_, '_>,
295    ) -> Result<bool> {
296        if context
297            .streaming
298            .as_ref()
299            .is_some_and(|streaming| streaming.body_ended)
300        {
301            return Err(streaming_boundary_error(
302                value,
303                "structural signal arrived after body_end",
304            ));
305        }
306
307        if value == "head_start" {
308            if context
309                .streaming
310                .as_ref()
311                .is_some_and(|streaming| streaming.head_marker_emitted)
312            {
313                return Err(HandlerError::DuplicateStreamingHeadStart);
314            }
315            context.writer.write(STREAMING_MARKER)?;
316            streaming_state(context)?.head_marker_emitted = true;
317            return Ok(true);
318        }
319
320        if value.starts_with(root::STREAMING_ROOT_PREFIX) {
321            process_streaming_root_signal(value, context)?;
322            return Ok(true);
323        }
324
325        if value == "body_end" {
326            require_streaming_head_start(context, "body_end")?;
327            if context
328                .streaming
329                .as_ref()
330                .is_some_and(|streaming| streaming.active_boundary.is_some())
331            {
332                return Err(streaming_boundary_error(
333                    value,
334                    "body ended while a boundary body was still rendering",
335                ));
336            }
337            if let Some(html) = context.body_inject {
338                context.writer.write(html)?;
339            }
340            if let Some(html) = context.state_inject.body_end {
341                context.writer.write(html)?;
342            }
343            let sequence = streaming_state(context)?.next_record_sequence;
344            self.emit_streaming_terminal(sequence, context)?;
345            {
346                let streaming = streaming_state(context)?;
347                increment_streaming_record_sequence("terminal", streaming)?;
348                streaming.body_ended = true;
349            }
350            context.body_end_emitted = true;
351            return Ok(true);
352        }
353
354        if value == "head_end" {
355            require_streaming_head_start(context, "head_end")?;
356        } else if value == "body_start" {
357            require_streaming_head_start(context, "body_start")?;
358        }
359        Ok(false)
360    }
361
362    /// Render every runtime boundary in document order as final.
363    ///
364    /// Uses one prepared render context and borrows the caller's state for the
365    /// complete traversal. Every compiler-selected flush still reaches the
366    /// transport, but the synchronous convenience path does not clone the state
367    /// or rebuild the continuation context at each boundary. Hosts that need
368    /// asynchronous work between occurrences should use
369    /// [`Self::stream_response`] or [`StreamingSession`].
370    pub fn render_streaming<'a, W: FlushWriter + ?Sized>(
371        &self,
372        protocol: &'a Protocol,
373        state: &Value,
374        options: &RenderOptions<'a>,
375        writer: &'a mut W,
376    ) -> Result<()> {
377        self.stream_response(protocol, options, writer)?
378            .render_all(state)
379    }
380}
webui_handler/streaming/mod.rs

webui_handler/streaming/
mod.rs
