portugalfuturista/rust-sdk
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: use Semaphore instead of Notify in OneshotTransport to prevent race condition (#611)

Browse source
This commit is contained in:
Maksim Madžar 2026-01-13 22:15:19 +01:00 committed by GitHub
parent 69dbd5a3ce
commit c4a68295e0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
crates/rmcp/src/transport.rs
View file

@ -177,7 +177,7 @@ where
{
message: Option<RxJsonRpcMessage<R>>,
sender: tokio::sync::mpsc::Sender<TxJsonRpcMessage<R>>,
finished_signal: Arc<tokio::sync::Notify>,
termination: Arc<tokio::sync::Semaphore>,
}
impl<R> OneshotTransport<R>
@ -192,7 +192,7 @@ where
Self {
message: Some(message),
sender,
finished_signal: Arc::new(tokio::sync::Notify::new()),
termination: Arc::new(tokio::sync::Semaphore::new(0)),
},
receiver,
)
@ -212,21 +212,22 @@ where
let sender = self.sender.clone();
let terminate = matches!(item, TxJsonRpcMessage::<R>::Response(_))
|| matches!(item, TxJsonRpcMessage::<R>::Error(_));
let signal = self.finished_signal.clone();
let termination = self.termination.clone();
async move {
sender.send(item).await?;
if terminate {
signal.notify_waiters();
termination.add_permits(1);
}
Ok(())
}
}
async fn receive(&mut self) -> Option<RxJsonRpcMessage<R>> {
if self.message.is_none() {
self.finished_signal.notified().await;
if let Some(msg) = self.message.take() {
return Some(msg);
}
self.message.take()
let _ = self.termination.acquire().await;
None
}
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {