Skip to content

Commit 14e9968

Browse files
committed
fix: rid unsoundness of integration test util
1 parent 680fd76 commit 14e9968

File tree

1 file changed

+12
-6
lines changed

1 file changed

+12
-6
lines changed

crates/base/src/utils/integration_test_helper.rs

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
#![allow(dead_code)]
22

33
use std::{
4-
mem::transmute,
4+
marker::PhantomPinned,
55
sync::Arc,
66
task::{ready, Poll},
77
};
@@ -99,9 +99,10 @@ impl RequestScope {
9999
RequestScopeGuard {
100100
cancelled: false,
101101
req_end_tx: self.req_end_tx.clone(),
102-
termination_token: self.termination_token.clone(),
102+
termination_token: Some(self.termination_token.clone()),
103103
conn_tx: self.conn.0.take().unwrap(),
104104
inner: None,
105+
_pinned: PhantomPinned,
105106
}
106107
}
107108
}
@@ -110,9 +111,10 @@ impl RequestScope {
110111
pub struct RequestScopeGuard {
111112
cancelled: bool,
112113
req_end_tx: mpsc::UnboundedSender<()>,
113-
termination_token: TerminationToken,
114+
termination_token: Option<TerminationToken>,
114115
conn_tx: watch::Sender<ConnSync>,
115116
inner: Option<BoxFuture<'static, ()>>,
117+
_pinned: PhantomPinned,
116118
}
117119

118120
impl Future for RequestScopeGuard {
@@ -127,11 +129,11 @@ impl Future for RequestScopeGuard {
127129
if !(*this.cancelled) {
128130
*this.cancelled = true;
129131
this.req_end_tx.send(()).unwrap();
130-
this.termination_token.inbound.cancel();
132+
this.termination_token.as_ref().unwrap().inbound.cancel();
131133
}
132134

133-
let inner = this.inner.get_or_insert_with(|| unsafe {
134-
transmute(this.termination_token.outbound.cancelled().boxed())
135+
let inner = this.inner.get_or_insert_with(|| {
136+
wait_termination(this.termination_token.take().unwrap()).boxed()
135137
});
136138

137139
ready!(inner.as_mut().poll_unpin(cx));
@@ -165,3 +167,7 @@ pub fn test_user_runtime_opts() -> UserWorkerRuntimeOpts {
165167
..Default::default()
166168
}
167169
}
170+
171+
async fn wait_termination(token: TerminationToken) {
172+
token.outbound.cancelled().await;
173+
}

0 commit comments

Comments
 (0)