Skip to content

Commit 8a51e3e

Browse files
ZhuchkaTriplesixZhuchkaTriplesix
authored
Merge pull request #113 from QueryaHub/issue-96-response-zero-copy
perf(response): reduce copies for str and bytes responses
2 parents 71f2ce6 + 40321d1 commit 8a51e3e

3 files changed

Lines changed: 131 additions & 33 deletions

File tree

‎src/dispatch.rs‎

Lines changed: 91 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -949,14 +949,14 @@ fn map_handler_return(py: Python<'_>, out: &Py<PyAny>) -> PyResult<HandlerMap> {
949949
if let Ok(s) = b.downcast::<PyString>() {
950950
return Ok(HandlerMap::Simple {
951951
status: 200,
952-
body: s.to_string().into_bytes(),
952+
body: SimpleBody::PyString(s.clone().unbind()),
953953
content_type: "text/plain; charset=utf-8".to_string(),
954954
});
955955
}
956956
if let Ok(buf) = b.downcast::<PyBytes>() {
957957
return Ok(HandlerMap::Simple {
958958
status: 200,
959-
body: buf.as_bytes().to_vec(),
959+
body: SimpleBody::PyBytes(buf.clone().unbind()),
960960
content_type: "application/octet-stream".to_string(),
961961
});
962962
}
@@ -974,21 +974,21 @@ fn map_handler_return(py: Python<'_>, out: &Py<PyAny>) -> PyResult<HandlerMap> {
974974
if let Ok(s) = b.extract::<String>() {
975975
return Ok(HandlerMap::Simple {
976976
status: 200,
977-
body: s.into_bytes(),
977+
body: SimpleBody::Owned(s.into_bytes()),
978978
content_type: "text/plain; charset=utf-8".to_string(),
979979
});
980980
}
981981
if let Ok(s) = b.extract::<&str>() {
982982
return Ok(HandlerMap::Simple {
983983
status: 200,
984-
body: s.as_bytes().to_vec(),
984+
body: SimpleBody::Owned(s.as_bytes().to_vec()),
985985
content_type: "text/plain; charset=utf-8".to_string(),
986986
});
987987
}
988988
if let Ok(buf) = b.extract::<Vec<u8>>() {
989989
return Ok(HandlerMap::Simple {
990990
status: 200,
991-
body: buf,
991+
body: SimpleBody::Owned(buf),
992992
content_type: "application/octet-stream".to_string(),
993993
});
994994
}
@@ -1013,7 +1013,7 @@ fn map_handler_return(py: Python<'_>, out: &Py<PyAny>) -> PyResult<HandlerMap> {
10131013
if let (Ok(code), Ok(bstr)) = (sc.extract::<u16>(), body.str()) {
10141014
return Ok(HandlerMap::Simple {
10151015
status: code,
1016-
body: bstr.to_string().into_bytes(),
1016+
body: SimpleBody::Owned(bstr.to_string().into_bytes()),
10171017
content_type: "text/plain; charset=utf-8".to_string(),
10181018
});
10191019
}
@@ -1024,11 +1024,60 @@ fn map_handler_return(py: Python<'_>, out: &Py<PyAny>) -> PyResult<HandlerMap> {
10241024
let s: String = dumped.extract()?;
10251025
Ok(HandlerMap::Simple {
10261026
status: 200,
1027-
body: s.into_bytes(),
1027+
body: SimpleBody::Owned(s.into_bytes()),
10281028
content_type: "application/json; charset=utf-8".to_string(),
10291029
})
10301030
}
10311031

1032+
fn send_simple_body_sync(
1033+
py: Python<'_>,
1034+
protocol: &Py<PyAny>,
1035+
status: u16,
1036+
body: &SimpleBody,
1037+
content_type: &str,
1038+
) -> PyResult<()> {
1039+
match body {
1040+
SimpleBody::Owned(bytes) => {
1041+
if bytes.is_empty() {
1042+
response::send_empty_sync(py, protocol, status, Some(content_type))
1043+
} else {
1044+
response::send_bytes_sync(py, protocol, status, bytes, content_type)
1045+
}
1046+
}
1047+
SimpleBody::PyString(s) => {
1048+
let text = s.bind(py).to_str()?;
1049+
response::send_text_sync(py, protocol, status, text, content_type)
1050+
}
1051+
SimpleBody::PyBytes(b) => {
1052+
response::send_pybytes_sync(py, protocol, status, b.bind(py), content_type)
1053+
}
1054+
}
1055+
}
1056+
1057+
enum SimpleBody {
1058+
Owned(Vec<u8>),
1059+
PyString(Py<PyString>),
1060+
PyBytes(Py<PyBytes>),
1061+
}
1062+
1063+
impl SimpleBody {
1064+
fn byte_len(&self, py: Python<'_>) -> PyResult<usize> {
1065+
match self {
1066+
SimpleBody::Owned(v) => Ok(v.len()),
1067+
SimpleBody::PyString(s) => Ok(s.bind(py).to_str()?.len()),
1068+
SimpleBody::PyBytes(b) => Ok(b.bind(py).as_bytes().len()),
1069+
}
1070+
}
1071+
1072+
fn into_vec(self, py: Python<'_>) -> PyResult<Vec<u8>> {
1073+
match self {
1074+
SimpleBody::Owned(v) => Ok(v),
1075+
SimpleBody::PyString(s) => Ok(s.bind(py).to_str()?.as_bytes().to_vec()),
1076+
SimpleBody::PyBytes(b) => Ok(b.bind(py).as_bytes().to_vec()),
1077+
}
1078+
}
1079+
}
1080+
10321081
enum HandlerMap {
10331082
AlreadySent,
10341083
WithHeaders {
@@ -1038,7 +1087,7 @@ enum HandlerMap {
10381087
},
10391088
Simple {
10401089
status: u16,
1041-
body: Vec<u8>,
1090+
body: SimpleBody,
10421091
content_type: String,
10431092
},
10441093
}
@@ -1066,7 +1115,13 @@ fn send_handler_map_inline(
10661115
status,
10671116
body,
10681117
content_type,
1069-
} => response::send_head_simple_sync(py, protocol, status, body.len(), &content_type),
1118+
} => response::send_head_simple_sync(
1119+
py,
1120+
protocol,
1121+
status,
1122+
body.byte_len(py)?,
1123+
&content_type,
1124+
),
10701125
}
10711126
} else {
10721127
match mapped {
@@ -1080,13 +1135,7 @@ fn send_handler_map_inline(
10801135
status,
10811136
body,
10821137
content_type,
1083-
} => {
1084-
if body.is_empty() {
1085-
response::send_empty_sync(py, protocol, status, Some(&content_type))
1086-
} else {
1087-
response::send_bytes_sync(py, protocol, status, &body, &content_type)
1088-
}
1089-
}
1138+
} => send_simple_body_sync(py, protocol, status, &body, &content_type),
10901139
}
10911140
}
10921141
}
@@ -1110,18 +1159,22 @@ fn merge_config_response_headers(
11101159
return Ok(mapped);
11111160
}
11121161
if if_absent {
1113-
Ok(merge_header_pairs_if_absent(mapped, &pairs))
1162+
merge_header_pairs_if_absent(py, mapped, &pairs)
11141163
} else {
1115-
Ok(merge_header_pairs_replace(mapped, &pairs))
1164+
merge_header_pairs_replace(py, mapped, &pairs)
11161165
}
11171166
}
11181167

1119-
fn merge_header_pairs_replace(mapped: HandlerMap, extra: &[(String, String)]) -> HandlerMap {
1168+
fn merge_header_pairs_replace(
1169+
py: Python<'_>,
1170+
mapped: HandlerMap,
1171+
extra: &[(String, String)],
1172+
) -> PyResult<HandlerMap> {
11201173
if extra.is_empty() {
1121-
return mapped;
1174+
return Ok(mapped);
11221175
}
11231176
match mapped {
1124-
HandlerMap::AlreadySent => HandlerMap::AlreadySent,
1177+
HandlerMap::AlreadySent => Ok(HandlerMap::AlreadySent),
11251178
HandlerMap::WithHeaders {
11261179
status,
11271180
body,
@@ -1131,37 +1184,42 @@ fn merge_header_pairs_replace(mapped: HandlerMap, extra: &[(String, String)]) ->
11311184
headers.retain(|(k, _)| !k.eq_ignore_ascii_case(a));
11321185
headers.push((a.clone(), b.clone()));
11331186
}
1134-
HandlerMap::WithHeaders {
1187+
Ok(HandlerMap::WithHeaders {
11351188
status,
11361189
body,
11371190
headers,
1138-
}
1191+
})
11391192
}
11401193
HandlerMap::Simple {
11411194
status,
11421195
body,
11431196
content_type,
11441197
} => {
1198+
let body = body.into_vec(py)?;
11451199
let mut headers = vec![("content-type".to_string(), content_type)];
11461200
for (a, b) in extra {
11471201
headers.retain(|(k, _)| !k.eq_ignore_ascii_case(a));
11481202
headers.push((a.clone(), b.clone()));
11491203
}
1150-
HandlerMap::WithHeaders {
1204+
Ok(HandlerMap::WithHeaders {
11511205
status,
11521206
body,
11531207
headers,
1154-
}
1208+
})
11551209
}
11561210
}
11571211
}
11581212

1159-
fn merge_header_pairs_if_absent(mapped: HandlerMap, extra: &[(String, String)]) -> HandlerMap {
1213+
fn merge_header_pairs_if_absent(
1214+
py: Python<'_>,
1215+
mapped: HandlerMap,
1216+
extra: &[(String, String)],
1217+
) -> PyResult<HandlerMap> {
11601218
if extra.is_empty() {
1161-
return mapped;
1219+
return Ok(mapped);
11621220
}
11631221
match mapped {
1164-
HandlerMap::AlreadySent => HandlerMap::AlreadySent,
1222+
HandlerMap::AlreadySent => Ok(HandlerMap::AlreadySent),
11651223
HandlerMap::WithHeaders {
11661224
status,
11671225
body,
@@ -1172,28 +1230,29 @@ fn merge_header_pairs_if_absent(mapped: HandlerMap, extra: &[(String, String)])
11721230
headers.push((a.clone(), b.clone()));
11731231
}
11741232
}
1175-
HandlerMap::WithHeaders {
1233+
Ok(HandlerMap::WithHeaders {
11761234
status,
11771235
body,
11781236
headers,
1179-
}
1237+
})
11801238
}
11811239
HandlerMap::Simple {
11821240
status,
11831241
body,
11841242
content_type,
11851243
} => {
1244+
let body = body.into_vec(py)?;
11861245
let mut headers = vec![("content-type".to_string(), content_type)];
11871246
for (a, b) in extra {
11881247
if !headers.iter().any(|(k, _)| k.eq_ignore_ascii_case(a)) {
11891248
headers.push((a.clone(), b.clone()));
11901249
}
11911250
}
1192-
HandlerMap::WithHeaders {
1251+
Ok(HandlerMap::WithHeaders {
11931252
status,
11941253
body,
11951254
headers,
1196-
}
1255+
})
11971256
}
11981257
}
11991258
}

‎src/response.rs‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! RSGI `response_bytes` / `response_str` / `response_empty` helpers (Granian RSGI spec).
22
33
use pyo3::prelude::*;
4-
use pyo3::types::{PyList, PyString, PyTuple};
4+
use pyo3::types::{PyBytes, PyList, PyString, PyTuple};
55

66
pub async fn send_text(
77
protocol: &Py<PyAny>,
@@ -162,6 +162,23 @@ pub fn send_text_sync(
162162
send_str_sync(py, protocol, status, text, content_type)
163163
}
164164

165+
/// Sync `protocol.response_bytes(...)` with a ``PyBytes`` buffer (no ``Vec`` copy).
166+
pub fn send_pybytes_sync(
167+
py: Python<'_>,
168+
protocol: &Py<PyAny>,
169+
status: u16,
170+
body: &Bound<'_, PyBytes>,
171+
content_type: &str,
172+
) -> PyResult<()> {
173+
if body.as_bytes().is_empty() {
174+
return send_empty_sync(py, protocol, status, Some(content_type));
175+
}
176+
let p = protocol.bind(py);
177+
let h = build_headers_ct(py, Some(content_type))?;
178+
p.getattr("response_bytes")?.call1((status, h, body))?;
179+
Ok(())
180+
}
181+
165182
/// Sync `protocol.response_bytes(...)`; falls back to `send_empty_sync` on empty body.
166183
pub fn send_bytes_sync(
167184
py: Python<'_>,

‎tests/test_sync_rsgi_short_circuit.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,28 @@ def hello() -> str:
8080
assert proto.sent and proto.sent[0][0] == 200 and proto.sent[0][1] == "hello"
8181

8282

83+
def test_handle_rsgi_trivial_bytes_response() -> None:
84+
app = App()
85+
86+
@app.get("/bin")
87+
def bin_route() -> bytes:
88+
return b"\x00\x01"
89+
90+
scope = SimpleNamespace(
91+
proto="http",
92+
method="GET",
93+
path="/bin",
94+
query_string="",
95+
headers={},
96+
)
97+
proto = _ProtoText()
98+
r = app.handle_rsgi(scope, proto)
99+
assert r is None
100+
assert not inspect.isawaitable(r)
101+
assert proto.sent[0][0] == 200
102+
assert proto.sent[0][1] == "\x00\x01"
103+
104+
83105
def test_handle_rsgi_route_with_deps_still_awaitable() -> None:
84106
app = App()
85107

0 commit comments

Comments
 (0)