[Rust][Benchmark] Preserve UTF-8 across benchmark stream chunks (#50868)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Aug 4, 2026 at 14:15 UTC
72cd5424da80a4a9caa3f42fd65bc0b94e61cbf0
1 file changed
+30
-10
rust/src/bench/src/backends/streaming.rs
+30
-10
@@ -6,7 +6,7 @@
6
/// Accumulates incoming byte chunks and extracts complete SSE messages.
7
/// Mirrors Python's `StreamedResponseHandler` from endpoint_request_func.py:22-60.
8
pub struct StreamedResponseHandler {
9
- buffer: String,
9
+ buffer: Vec<u8>,
10
/// Reusable message buffer — avoids allocating a new Vec per `add_chunk` call.
11
messages: Vec<String>,
12
}
@@ -14,7 +14,7 @@ pub struct StreamedResponseHandler {
14
impl StreamedResponseHandler {
15
pub fn new() -> Self {
16
Self {
17
- buffer: String::with_capacity(4096),
17
+ buffer: Vec::with_capacity(4096),
18
messages: Vec::with_capacity(4),
19
}
20
}
@@ -26,12 +26,11 @@ impl StreamedResponseHandler {
26
pub fn add_chunk(&mut self, chunk_bytes: &[u8]) -> &[String] {
27
self.messages.clear();
28
29
- let chunk_str = String::from_utf8_lossy(chunk_bytes);
30
- self.buffer.push_str(&chunk_str);
29
+ self.buffer.extend_from_slice(chunk_bytes);
30
31
// Split by double newlines (SSE message separator)
33
- while let Some(pos) = self.buffer.find("\n\n") {
34
- let message = self.buffer[..pos].trim().to_string();
32
+ while let Some(pos) = self.buffer.windows(2).position(|window| window == b"\n\n") {
33
+ let message = String::from_utf8_lossy(&self.buffer[..pos]).trim().to_string();
34
// Efficiently remove consumed bytes by shifting remaining data
35
self.buffer.drain(..pos + 2);
36
if !message.is_empty() {
@@ -47,19 +46,20 @@ impl StreamedResponseHandler {
46
//
47
// Also handles multi-field SSE events where the buffer may start with
48
// "event: ...\ndata: ..." (Dynamo frontend).
50
- let data_start = if self.buffer.starts_with("data: ") {
49
+ let buffer = String::from_utf8_lossy(&self.buffer);
50
+ let data_start = if buffer.starts_with("data: ") {
51
Some(0)
52
} else {
53
// Look for a "data: " line in multi-field events
54
- self.buffer.find("\ndata: ").map(|p| p + 1)
54
+ buffer.find("\ndata: ").map(|p| p + 1)
55
};
56
if let Some(offset) = data_start {
57
- let content = self.buffer[offset + 6..].trim();
57
+ let content = buffer[offset + 6..].trim();
58
if content == "[DONE]"
59
|| (!content.is_empty()
60
&& serde_json::from_str::<&serde_json::value::RawValue>(content).is_ok())
61
{
62
- self.messages.push(self.buffer.trim().to_string());
62
+ self.messages.push(buffer.trim().to_string());
63
self.buffer.clear();
64
}
65
}
@@ -102,6 +102,26 @@ mod tests {
102
assert_eq!(msgs2.len(), 1);
103
}
104
105
+ #[test]
106
+ fn test_split_multibyte_utf8() {
107
+ let mut handler = StreamedResponseHandler::new();
108
+ let message = "data: {\"test\":\"中😀\"}\n\n";
109
+ let mut messages = Vec::new();
110
+
111
+ for chunk in message.as_bytes().chunks(1) {
112
+ messages.extend_from_slice(handler.add_chunk(chunk));
113
+ }
114
+
115
+ assert_eq!(messages, &[message.trim().to_string()]);
116
+ }
117
+
118
+ #[test]
119
+ fn test_invalid_utf8_remains_lossy() {
120
+ let mut handler = StreamedResponseHandler::new();
121
+ let msgs = handler.add_chunk(b"data: {\"test\":\"\xFF\"}\n\n");
122
+ assert_eq!(msgs, &["data: {\"test\":\"�\"}".to_string()]);
123
+ }
124
+
125
#[test]
126
fn test_comment_lines() {
127
let mut handler = StreamedResponseHandler::new();