Skip to content

Commit 808124b

Browse files
TheSeydiCharyyevTheSeydiCharyyev
authored
fix(streamable-http-server): release the session map lock before waiting on a session (#1323)
* fix(streamable-http-server): release the session map lock before waiting on a session LocalSessionManager kept the read guard on its session map while it waited on a session worker. initialize_session keeps it for the whole ServerHandler::initialize, so one slow initialize blocks create_session for every new client. tokio's RwLock is fair, so has_session calls for other sessions then wait behind that writer, and the whole server waits until the slow initialize returns. Clone the session handle out of the map and drop the guard before waiting, like close_session already does. * fix(streamable-http-server): make the isolation test deterministic and clone the id lazily The test now runs with the tokio clock paused, so each step starts only after all spawned tasks are blocked, instead of after a fixed sleep. It also waits until the worker has passed initialize on before it adds the second call. session_handle now builds the SessionNotFound error only when the session is missing. * test(streamable-http-server): cover close_session during a pending initialize Add a test that close_session returns at once while initialize is still pending and removes the session, and that the pending initialize still gets the handler's response afterwards. The accept_message case now reads the channel capacity from the manager's session config instead of the default constant.
1 parent 07ff92d commit 808124b

2 files changed

Lines changed: 197 additions & 20 deletions

File tree

‎crates/rmcp/src/transport/streamable_http_server/session/local.rs‎

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,20 @@ impl LocalSessionManager {
4242
self.event_store = Some(event_store);
4343
self
4444
}
45+
46+
/// Clone the session handle out of the map, so the lock is released
47+
/// before the caller waits on the session worker.
48+
async fn session_handle(
49+
&self,
50+
id: &SessionId,
51+
) -> Result<LocalSessionHandle, LocalSessionManagerError> {
52+
self.sessions
53+
.read()
54+
.await
55+
.get(id)
56+
.cloned()
57+
.ok_or_else(|| LocalSessionManagerError::SessionNotFound(id.clone()))
58+
}
4559
}
4660

4761
#[derive(Debug, Error)]
@@ -72,10 +86,7 @@ impl SessionManager for LocalSessionManager {
7286
id: &SessionId,
7387
message: ClientJsonRpcMessage,
7488
) -> Result<ServerJsonRpcMessage, Self::Error> {
75-
let sessions = self.sessions.read().await;
76-
let handle = sessions
77-
.get(id)
78-
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
89+
let handle = self.session_handle(id).await?;
7990
let response = handle.initialize(message).await?;
8091
Ok(response)
8192
}
@@ -102,10 +113,7 @@ impl SessionManager for LocalSessionManager {
102113
id: &SessionId,
103114
message: ClientJsonRpcMessage,
104115
) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
105-
let sessions = self.sessions.read().await;
106-
let handle = sessions
107-
.get(id)
108-
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
116+
let handle = self.session_handle(id).await?;
109117
let receiver = handle.establish_request_wise_channel().await?;
110118
let http_request_id = receiver.http_request_id;
111119
handle.push_message(message, http_request_id).await?;
@@ -116,10 +124,7 @@ impl SessionManager for LocalSessionManager {
116124
&self,
117125
id: &SessionId,
118126
) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
119-
let sessions = self.sessions.read().await;
120-
let handle = sessions
121-
.get(id)
122-
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
127+
let handle = self.session_handle(id).await?;
123128
let receiver = handle.establish_common_channel().await?;
124129
Ok(ReceiverStream::new(receiver.inner))
125130
}
@@ -136,10 +141,7 @@ impl SessionManager for LocalSessionManager {
136141
.map_err(SessionError::EventStore)?;
137142
return Ok(stream.left_stream());
138143
}
139-
let sessions = self.sessions.read().await;
140-
let handle = sessions
141-
.get(id)
142-
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
144+
let handle = self.session_handle(id).await?;
143145
let receiver = handle.resume(last_event_id.parse()?).await?;
144146
Ok(ReceiverStream::new(receiver.inner).right_stream())
145147
}
@@ -149,10 +151,7 @@ impl SessionManager for LocalSessionManager {
149151
id: &SessionId,
150152
message: ClientJsonRpcMessage,
151153
) -> Result<(), Self::Error> {
152-
let sessions = self.sessions.read().await;
153-
let handle = sessions
154-
.get(id)
155-
.ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
154+
let handle = self.session_handle(id).await?;
156155
handle.push_message(message, None).await?;
157156
Ok(())
158157
}

‎crates/rmcp/tests/test_streamable_http_session_isolation.rs‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
#![cfg(all(feature = "transport-streamable-http-server", not(feature = "local")))]
2+
3+
use std::{sync::Arc, time::Duration};
4+
5+
use rmcp::{
6+
model::{
7+
ClientJsonRpcMessage, ClientRequest, PingRequest, RequestId, ServerJsonRpcMessage,
8+
ServerResult,
9+
},
10+
transport::{
11+
Transport,
12+
streamable_http_server::session::{SessionId, SessionManager, local::LocalSessionManager},
13+
},
14+
};
15+
use rstest::rstest;
16+
use tokio::task::JoinHandle;
17+
18+
/// Virtual time: the tests run with the clock paused.
19+
const PROBE_TIMEOUT: Duration = Duration::from_secs(1);
20+
21+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22+
enum StuckCall {
23+
Initialize,
24+
CreateStream,
25+
CreateStandaloneStream,
26+
Resume,
27+
AcceptMessage,
28+
}
29+
30+
fn ping() -> ClientJsonRpcMessage {
31+
ClientJsonRpcMessage::request(
32+
ClientRequest::PingRequest(PingRequest::default()),
33+
RequestId::Number(1),
34+
)
35+
}
36+
37+
/// With the clock paused, time only moves when every task is blocked, so this
38+
/// returns after all spawned tasks have run as far as they can.
39+
async fn settle() {
40+
tokio::time::sleep(Duration::from_millis(1)).await;
41+
}
42+
43+
fn spawn_call(
44+
manager: &Arc<LocalSessionManager>,
45+
id: &SessionId,
46+
call: StuckCall,
47+
) -> JoinHandle<()> {
48+
let manager = manager.clone();
49+
let id = id.clone();
50+
tokio::spawn(async move {
51+
match call {
52+
StuckCall::Initialize => {
53+
let _ = manager.initialize_session(&id, ping()).await;
54+
}
55+
StuckCall::CreateStream => {
56+
let _ = manager.create_stream(&id, ping()).await;
57+
}
58+
StuckCall::CreateStandaloneStream => {
59+
let _ = manager.create_standalone_stream(&id).await;
60+
}
61+
StuckCall::Resume => {
62+
let _ = manager.resume(&id, "0".to_owned()).await;
63+
}
64+
StuckCall::AcceptMessage => {
65+
// The worker no longer reads its event channel, so the push
66+
// after the channel is full waits.
67+
for _ in 0..=manager.session_config.channel_capacity {
68+
let _ = manager.accept_message(&id, ping()).await;
69+
}
70+
}
71+
}
72+
})
73+
}
74+
75+
/// While a call waits on the worker of one session, the manager must still
76+
/// create new sessions and look up other sessions.
77+
#[rstest]
78+
#[case::initialize(StuckCall::Initialize)]
79+
#[case::create_stream(StuckCall::CreateStream)]
80+
#[case::create_standalone_stream(StuckCall::CreateStandaloneStream)]
81+
#[case::resume(StuckCall::Resume)]
82+
#[case::accept_message(StuckCall::AcceptMessage)]
83+
#[tokio::test(start_paused = true)]
84+
async fn waiting_on_one_session_does_not_block_other_sessions(
85+
#[case] call: StuckCall,
86+
) -> anyhow::Result<()> {
87+
let manager = Arc::new(LocalSessionManager::default());
88+
let (other_id, _other_transport) = manager.create_session().await?;
89+
// Nothing answers on this transport, so the worker never gets the
90+
// initialize response and stops reading session events.
91+
let (stuck_id, mut stuck_transport) = manager.create_session().await?;
92+
93+
let mut stuck = vec![spawn_call(&manager, &stuck_id, StuckCall::Initialize)];
94+
// The worker has passed initialize on and now waits for the response.
95+
assert!(
96+
stuck_transport.receive().await.is_some(),
97+
"the worker did not pass initialize on"
98+
);
99+
if call != StuckCall::Initialize {
100+
stuck.push(spawn_call(&manager, &stuck_id, call));
101+
}
102+
settle().await;
103+
104+
// A new client connects while the call is waiting...
105+
let create = tokio::spawn({
106+
let manager = manager.clone();
107+
async move { manager.create_session().await.map(|(id, _)| id) }
108+
});
109+
settle().await;
110+
111+
// ...and a request arrives on another session.
112+
let lookup = tokio::time::timeout(PROBE_TIMEOUT, manager.has_session(&other_id)).await;
113+
assert!(
114+
matches!(lookup, Ok(Ok(true))),
115+
"has_session blocked: {lookup:?}"
116+
);
117+
let created = tokio::time::timeout(PROBE_TIMEOUT, create).await;
118+
assert!(
119+
matches!(created, Ok(Ok(Ok(_)))),
120+
"create_session blocked: {created:?}"
121+
);
122+
123+
for task in stuck {
124+
assert!(
125+
!task.is_finished(),
126+
"the call on the stuck session should still be waiting"
127+
);
128+
task.abort();
129+
}
130+
Ok(())
131+
}
132+
133+
/// `close_session` must not wait for a pending `initialize`. The worker reads
134+
/// session events only after initialization, so the pending call still waits
135+
/// for the handler and still gets its response.
136+
#[tokio::test(start_paused = true)]
137+
async fn close_session_does_not_wait_for_pending_initialize() -> anyhow::Result<()> {
138+
let manager = Arc::new(LocalSessionManager::default());
139+
let (id, mut transport) = manager.create_session().await?;
140+
let initialize = tokio::spawn({
141+
let manager = manager.clone();
142+
let id = id.clone();
143+
async move { manager.initialize_session(&id, ping()).await }
144+
});
145+
assert!(
146+
transport.receive().await.is_some(),
147+
"the worker did not pass initialize on"
148+
);
149+
150+
let closed = tokio::time::timeout(PROBE_TIMEOUT, manager.close_session(&id)).await;
151+
assert!(
152+
matches!(closed, Ok(Ok(()))),
153+
"close_session blocked: {closed:?}"
154+
);
155+
assert!(
156+
!manager.has_session(&id).await?,
157+
"the closed session is still in the map"
158+
);
159+
160+
settle().await;
161+
assert!(
162+
!initialize.is_finished(),
163+
"initialize should still wait for the handler"
164+
);
165+
166+
transport
167+
.send(ServerJsonRpcMessage::response(
168+
ServerResult::empty(()),
169+
RequestId::Number(1),
170+
))
171+
.await?;
172+
let response = tokio::time::timeout(PROBE_TIMEOUT, initialize).await;
173+
assert!(
174+
matches!(response, Ok(Ok(Ok(ServerJsonRpcMessage::Response(_))))),
175+
"initialize did not get the handler's response: {response:?}"
176+
);
177+
Ok(())
178+
}

0 commit comments

Comments
 (0)