actix/actix-extras
1
0
Fork 0
mirror of https://github.com/actix/actix-extras.git synced 2024-11-23 23:51:06 +01:00
Code Issues Releases Wiki Activity

actix-ws: use BigBytes to avoid copying ws messages between various buffers

Browse Source
This commit is contained in:
asonix 2024-11-04 17:42:58 -06:00
parent 77406cbb71
commit a47d61c8d6
1 changed files with 13 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
actix-ws/src/stream.rs
View File

@ -1,13 +1,14 @@
use std::{
collections::VecDeque,
future::poll_fn,
io, mem,
io,
pin::Pin,
task::{Context, Poll},
};
use actix_codec::{Decoder, Encoder};
use actix_codec::Decoder;
use actix_http::{
big_bytes::BigBytes,
ws::{Codec, Frame, Message, ProtocolError},
Payload,
};
@ -24,8 +25,7 @@ use crate::AggregatedMessageStream;
/// Response body for a WebSocket.
pub struct StreamingBody {
session_rx: Receiver<Message>,
messages: VecDeque<Message>,
buf: BytesMut,
buf: BigBytes,
codec: Codec,
closing: bool,
}
@ -34,8 +34,7 @@ impl StreamingBody {
pub(super) fn new(session_rx: Receiver<Message>) -> Self {
StreamingBody {
session_rx,
messages: VecDeque::new(),
buf: BytesMut::new(),
buf: BigBytes::with_capacity(0),
codec: Codec::new(),
closing: false,
}
@ -118,14 +117,12 @@ impl Stream for StreamingBody {
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
if this.closing {
return Poll::Ready(None);
}
loop {
while !this.closing {
match Pin::new(&mut this.session_rx).poll_recv(cx) {
Poll::Ready(Some(msg)) => {
this.messages.push_back(msg);
if let Err(err) = this.codec.encode_bigbytes(msg, &mut this.buf) {
return Poll::Ready(Some(Err(err.into())));
}
}
Poll::Ready(None) => {
this.closing = true;
@ -135,14 +132,12 @@ impl Stream for StreamingBody {
}
}
while let Some(msg) = this.messages.pop_front() {
if let Err(err) = this.codec.encode(msg, &mut this.buf) {
return Poll::Ready(Some(Err(err.into())));
}
if let Some(bytes) = this.buf.pop_front() {
return Poll::Ready(Some(Ok(bytes)));
}
if !this.buf.is_empty() {
return Poll::Ready(Some(Ok(mem::take(&mut this.buf).freeze())));
if this.closing {
return Poll::Ready(None);
}
Poll::Pending