Skip to content

Commit dfe88a9

Browse files
committed
feat(response): clear cached data when response is dropped
1 parent 4b4e586 commit dfe88a9

2 files changed

Lines changed: 68 additions & 48 deletions

File tree

src/client/resp/http.rs

Lines changed: 50 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -64,19 +64,19 @@ impl Response {
6464
///
6565
/// This creates a new HTTP response with the same version, status, headers, and extensions
6666
/// as the current response, but with the provided body.
67-
fn build_response(self, body: wreq::Body) -> wreq::Response {
67+
fn build_response(&self, body: wreq::Body) -> wreq::Response {
6868
let mut response = HttpResponse::new(body);
6969
*response.version_mut() = self.parts.version;
7070
*response.status_mut() = self.parts.status;
71-
*response.headers_mut() = self.parts.headers;
72-
*response.extensions_mut() = self.parts.extensions;
71+
*response.headers_mut() = self.parts.headers.clone();
72+
*response.extensions_mut() = self.parts.extensions.clone();
7373
wreq::Response::from(response)
7474
}
7575

7676
/// Creates an empty response with the same metadata but no body content.
7777
///
7878
/// Useful for operations that only need response headers/metadata without consuming the body.
79-
fn empty_response(self) -> wreq::Response {
79+
fn empty_response(&self) -> wreq::Response {
8080
self.build_response(wreq::Body::from(Bytes::new()))
8181
}
8282

@@ -114,14 +114,29 @@ impl Response {
114114
///
115115
/// This method transfers ownership of the streamable body for one-time use.
116116
/// Returns an error if the body has already been consumed or is not streamable.
117-
fn stream_response(self) -> Result<wreq::Response, Error> {
117+
fn stream_response(&self) -> Result<wreq::Response, Error> {
118118
if let Some(arc) = self.body.swap(None) {
119119
if let Ok(Body::Streamable(body)) = Arc::try_unwrap(arc) {
120120
return Ok(self.build_response(body));
121121
}
122122
}
123123
Err(Error::Memory)
124124
}
125+
126+
/// Forcefully destroys the response body, preventing any further reads.
127+
fn destroy(&self) {
128+
if let Some(body) = self.body.swap(None) {
129+
if let Ok(body) = Arc::try_unwrap(body) {
130+
::std::mem::drop(body);
131+
}
132+
}
133+
}
134+
}
135+
136+
impl Drop for Response {
137+
fn drop(&mut self) {
138+
self.destroy();
139+
}
125140
}
126141

127142
#[pymethods]
@@ -159,19 +174,19 @@ impl Response {
159174
/// Get the content length of the response.
160175
#[getter]
161176
pub fn content_length(&self, py: Python) -> Option<u64> {
162-
py.detach(|| self.clone().empty_response().content_length())
177+
py.detach(|| self.empty_response().content_length())
163178
}
164179

165180
/// Get the remote address of the response.
166181
#[getter]
167182
pub fn remote_addr(&self, py: Python) -> Option<SocketAddr> {
168-
py.detach(|| self.clone().empty_response().remote_addr().map(SocketAddr))
183+
py.detach(|| self.empty_response().remote_addr().map(SocketAddr))
169184
}
170185

171186
/// Get the local address of the response.
172187
#[getter]
173188
pub fn local_addr(&self, py: Python) -> Option<SocketAddr> {
174-
py.detach(|| self.clone().empty_response().local_addr().map(SocketAddr))
189+
py.detach(|| self.empty_response().local_addr().map(SocketAddr))
175190
}
176191

177192
/// Get the redirect history of the Response.
@@ -213,8 +228,7 @@ impl Response {
213228

214229
/// Get the response into a `Stream` of `Bytes` from the body.
215230
pub fn stream(&self) -> PyResult<Streamer> {
216-
self.clone()
217-
.stream_response()
231+
self.stream_response()
218232
.map(Streamer::new)
219233
.map_err(Into::into)
220234
}
@@ -259,20 +273,21 @@ impl Response {
259273
///
260274
/// **Current behavior:**
261275
/// - When connection pooling is **disabled**: This method closes the network connection.
262-
/// - When connection pooling is **enabled**: This method closes the response, prevents further body reads,
263-
/// and returns the connection to the pool for reuse.
276+
/// - When connection pooling is **enabled**: This method closes the response, prevents further
277+
/// body reads, and returns the connection to the pool for reuse.
264278
///
265279
/// **Future changes:**
266-
/// In future versions, this method will be changed to always close the network connection regardless of
267-
/// whether connection pooling is enabled or not.
280+
/// In future versions, this method will be changed to always close the network connection
281+
/// regardless of whether connection pooling is enabled or not.
268282
///
269283
/// **Recommendation:**
270-
/// It is **not recommended** to manually call this method at present. Instead, use context managers
271-
/// (async with statement) to properly manage response lifecycle. Wait for the improved implementation
272-
/// in future versions.
273-
pub async fn close(&self) -> PyResult<()> {
274-
self.body.swap(None);
275-
Ok(())
284+
/// It is **not recommended** to manually call this method at present. Instead, use context
285+
/// managers (async with statement) to properly manage response lifecycle. Wait for the
286+
/// improved implementation in future versions.
287+
pub async fn close(&self) {
288+
Python::attach(|py| {
289+
py.detach(|| self.destroy());
290+
});
276291
}
277292
}
278293

@@ -284,12 +299,7 @@ impl Response {
284299
}
285300

286301
#[inline]
287-
async fn __aexit__(
288-
&self,
289-
_exc_type: Py<PyAny>,
290-
_exc_val: Py<PyAny>,
291-
_traceback: Py<PyAny>,
292-
) -> PyResult<()> {
302+
async fn __aexit__(&self, _exc_type: Py<PyAny>, _exc_val: Py<PyAny>, _traceback: Py<PyAny>) {
293303
self.close().await
294304
}
295305
}
@@ -308,6 +318,12 @@ impl Display for Response {
308318

309319
// ===== impl BlockingResponse =====
310320

321+
impl Drop for BlockingResponse {
322+
fn drop(&mut self) {
323+
self.0.destroy();
324+
}
325+
}
326+
311327
#[pymethods]
312328
impl BlockingResponse {
313329
/// Get the URL of the response.
@@ -427,20 +443,20 @@ impl BlockingResponse {
427443
///
428444
/// **Current behavior:**
429445
/// - When connection pooling is **disabled**: This method closes the network connection.
430-
/// - When connection pooling is **enabled**: This method closes the response, prevents further body reads,
431-
/// and returns the connection to the pool for reuse.
446+
/// - When connection pooling is **enabled**: This method closes the response, prevents further
447+
/// body reads, and returns the connection to the pool for reuse.
432448
///
433449
/// **Future changes:**
434-
/// In future versions, this method will be changed to always close the network connection regardless of
435-
/// whether connection pooling is enabled or not.
450+
/// In future versions, this method will be changed to always close the network connection
451+
/// regardless of whether connection pooling is enabled or not.
436452
///
437453
/// **Recommendation:**
438-
/// It is **not recommended** to manually call this method at present. Instead, use context managers
439-
/// (with statement) to properly manage response lifecycle. Wait for the improved implementation
440-
/// in future versions.
454+
/// It is **not recommended** to manually call this method at present. Instead, use context
455+
/// managers (with statement) to properly manage response lifecycle. Wait for the improved
456+
/// implementation in future versions.
441457
#[inline]
442458
pub fn close(&self, py: Python) {
443-
py.detach(|| self.0.body.swap(None));
459+
py.detach(|| self.0.destroy());
444460
}
445461
}
446462

src/lib.rs

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -66,15 +66,17 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
6666
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
6767

6868
mod r#async {
69-
use crate::client::{
70-
Client,
71-
req::Request,
72-
req::WebSocketRequest,
73-
resp::{Response, WebSocket},
74-
};
75-
use crate::http::Method;
7669
use pyo3::{coroutine::CancelHandle, prelude::*, pybacked::PyBackedStr};
7770

71+
use crate::{
72+
client::{
73+
Client,
74+
req::{Request, WebSocketRequest},
75+
resp::{Response, WebSocket},
76+
},
77+
http::Method,
78+
};
79+
7880
/// Make a GET request with the given parameters.
7981
#[inline]
8082
#[pyfunction]
@@ -198,15 +200,17 @@ mod r#async {
198200
}
199201

200202
mod blocking {
201-
use crate::client::{
202-
BlockingClient,
203-
req::Request,
204-
req::WebSocketRequest,
205-
resp::{BlockingResponse, BlockingWebSocket},
206-
};
207-
use crate::http::Method;
208203
use pyo3::{prelude::*, pybacked::PyBackedStr};
209204

205+
use crate::{
206+
client::{
207+
BlockingClient,
208+
req::{Request, WebSocketRequest},
209+
resp::{BlockingResponse, BlockingWebSocket},
210+
},
211+
http::Method,
212+
};
213+
210214
/// Make a GET request with the given parameters (blocking).
211215
#[inline]
212216
#[pyfunction]

0 commit comments

Comments
 (0)