Skip to content

Commit 21882fa

Browse files
committed
Add mutex to avoid data race bewteen Start and Shutdown in SSEServer struct
1 parent 71b910b commit 21882fa

File tree

1 file changed

+13
-2
lines changed

1 file changed

+13
-2
lines changed

server/sse.go

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ type SSEServer struct {
6565

6666
keepAlive bool
6767
keepAliveInterval time.Duration
68+
69+
mu sync.RWMutex
6870
}
6971

7072
// SSEOption defines a function type for configuring SSEServer
@@ -189,18 +191,25 @@ func NewTestServer(server *MCPServer, opts ...SSEOption) *httptest.Server {
189191
// Start begins serving SSE connections on the specified address.
190192
// It sets up HTTP handlers for SSE and message endpoints.
191193
func (s *SSEServer) Start(addr string) error {
194+
s.mu.Lock()
192195
s.srv = &http.Server{
193196
Addr: addr,
194197
Handler: s,
195198
}
199+
s.mu.Unlock()
196200

197201
return s.srv.ListenAndServe()
198202
}
199203

200204
// Shutdown gracefully stops the SSE server, closing all active sessions
201205
// and shutting down the HTTP server.
202206
func (s *SSEServer) Shutdown(ctx context.Context) error {
203-
if s.srv != nil {
207+
s.mu.RLock()
208+
srv := s.srv
209+
s.mu.RUnlock()
210+
211+
if srv != nil {
212+
// 关闭所有会话
204213
s.sessions.Range(func(key, value interface{}) bool {
205214
if session, ok := value.(*sseSession); ok {
206215
close(session.done)
@@ -209,7 +218,7 @@ func (s *SSEServer) Shutdown(ctx context.Context) error {
209218
return true
210219
})
211220

212-
return s.srv.Shutdown(ctx)
221+
return srv.Shutdown(ctx)
213222
}
214223
return nil
215224
}
@@ -336,7 +345,9 @@ func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
336345
return
337346
}
338347

348+
s.mu.RLock()
339349
sessionI, ok := s.sessions.Load(sessionID)
350+
s.mu.RUnlock()
340351
if !ok {
341352
s.writeJSONRPCError(w, nil, mcp.INVALID_PARAMS, "Invalid session ID")
342353
return

0 commit comments

Comments
 (0)