use std::sync::Arc;
use serde_json::Value;
use super::session::{BoundaryId, BoundaryMode, ParkedResponse};
use super::StreamingResponse;
use crate::route_handler::Protocol;
use crate::{FlushWriter, RenderOptions, ResponseWriter, Result, WebUIHandler};
#[derive(Default)]
pub struct BufferSink {
bytes: Vec<u8>,
}
impl ResponseWriter for BufferSink {
fn write(&mut self, content: &str) -> Result<()> {
self.bytes.extend_from_slice(content.as_bytes());
Ok(())
}
fn end(&mut self) -> Result<()> {
Ok(())
}
}
impl FlushWriter for BufferSink {
fn flush(&mut self) -> Result<()> {
Ok(())
}
}
#[derive(Clone, Debug)]
pub struct SessionOptions {
pub entry_id: String,
pub request_path: String,
pub nonce: Option<String>,
pub head_inject: Option<String>,
pub body_inject: Option<String>,
}
impl SessionOptions {
#[must_use]
pub fn new(entry_id: impl Into<String>, request_path: impl Into<String>) -> Self {
Self {
entry_id: entry_id.into(),
request_path: request_path.into(),
nonce: None,
head_inject: None,
body_inject: None,
}
}
}
pub struct StreamingSession {
handler: Arc<WebUIHandler>,
protocol: Arc<Protocol>,
options: SessionOptions,
parked: Option<ParkedResponse>,
sink: BufferSink,
boundary_count: usize,
finished: bool,
}
impl StreamingSession {
pub fn new(
handler: Arc<WebUIHandler>,
protocol: Arc<Protocol>,
options: SessionOptions,
) -> Result<Self> {
let mut sink = BufferSink::default();
let (parked, boundary_count) = {
let render_options = render_options(&options);
let response = handler.stream_response(&protocol, &render_options, &mut sink)?;
let boundary_count = response.boundary_count();
(response.park(), boundary_count)
};
Ok(Self {
handler,
protocol,
options,
parked: Some(parked),
sink,
boundary_count,
finished: false,
})
}
pub fn boundary(&self, name: &str) -> Result<BoundaryId> {
let names = self
.protocol
.streaming_boundary_names(&self.options.entry_id);
names
.iter()
.position(|candidate| candidate == name)
.ok_or_else(|| super::error::unknown_boundary_name_error(name, names))
.and_then(BoundaryId::from_index)
}
#[must_use]
pub fn boundary_count(&self) -> usize {
self.boundary_count
}
#[must_use]
pub fn is_finished(&self) -> bool {
self.finished
}
pub fn write_shell(&mut self, state: &Value) -> Result<Vec<u8>> {
self.run(|response| response.write_shell(state))
}
pub fn write_boundary(
&mut self,
boundary: BoundaryId,
state: &Value,
mode: BoundaryMode,
) -> Result<Vec<u8>> {
self.run(|response| response.write_boundary(boundary, state, mode))
}
pub fn update(&mut self, boundary: BoundaryId, state: &Value) -> Result<Vec<u8>> {
self.run(|response| response.update(boundary, state))
}
pub fn finish(&mut self, state: &Value) -> Result<Vec<u8>> {
let parked = self.take_parked("finish")?;
let result = {
let Self {
handler,
protocol,
options,
sink,
parked: slot,
..
} = self;
let render_options = render_options(options);
match StreamingResponse::unpark(parked, handler, protocol, &render_options, sink) {
Ok(response) => match response.ensure_finishable() {
Ok(()) => response.finish(state),
Err(error) => {
*slot = Some(response.park());
Err(error)
}
},
Err(error) => Err(error),
}
};
match result {
Ok(()) => {
self.finished = true;
Ok(std::mem::take(&mut self.sink.bytes))
}
Err(error) => {
self.sink.bytes = Vec::new();
Err(error)
}
}
}
fn take_parked(&mut self, operation: &str) -> Result<ParkedResponse> {
if let Some(parked) = self.parked.take() {
return Ok(parked);
}
let reason = if self.finished {
"the streaming response has already finished"
} else {
"the streaming response failed while emitting its terminal record and \
cannot continue; bytes may already have been sent"
};
Err(super::error::boundary_order_error(operation, reason))
}
fn run(
&mut self,
operation: impl FnOnce(&mut StreamingResponse<'_, BufferSink>) -> Result<()>,
) -> Result<Vec<u8>> {
let parked = self.take_parked("streaming operation")?;
let Self {
handler,
protocol,
options,
sink,
parked: slot,
..
} = self;
let render_options = render_options(options);
let mut response =
StreamingResponse::unpark(parked, handler, protocol, &render_options, sink)?;
let result = operation(&mut response);
*slot = Some(response.park());
result?;
Ok(std::mem::take(&mut self.sink.bytes))
}
}
fn render_options(options: &SessionOptions) -> RenderOptions<'_> {
RenderOptions {
entry_id: &options.entry_id,
request_path: &options.request_path,
nonce: options.nonce.as_deref(),
head_inject: options.head_inject.as_deref(),
body_inject: options.body_inject.as_deref(),
}
}