From 552c66feb8d9a267a3fee5fc2e7e5f8e51398d45 Mon Sep 17 00:00:00 2001 From: cryo Date: Mon, 21 Apr 2025 09:35:01 +0000 Subject: [PATCH 1/2] refact: optimize usage of RWMutex in MCPServer for concurrency performance --- server/server.go | 41 ++++++++++++++++++++++++++++++++--------- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/server/server.go b/server/server.go index d121f9b34..91256d11f 100644 --- a/server/server.go +++ b/server/server.go @@ -407,11 +407,18 @@ func (s *MCPServer) AddResource( resource mcp.Resource, handler ResourceHandlerFunc, ) { - s.capabilitiesMu.Lock() + s.capabilitiesMu.RLock() if s.capabilities.resources == nil { - s.capabilities.resources = &resourceCapabilities{} + s.capabilitiesMu.RUnlock() + + s.capabilitiesMu.Lock() + if s.capabilities.resources == nil { + s.capabilities.resources = &resourceCapabilities{} + } + s.capabilitiesMu.Unlock() + } else { + s.capabilitiesMu.RUnlock() } - s.capabilitiesMu.Unlock() s.resourcesMu.Lock() defer s.resourcesMu.Unlock() @@ -438,11 +445,17 @@ func (s *MCPServer) AddResourceTemplate( template mcp.ResourceTemplate, handler ResourceTemplateHandlerFunc, ) { - s.capabilitiesMu.Lock() + s.capabilitiesMu.RLock() if s.capabilities.resources == nil { + s.capabilitiesMu.RUnlock() + + s.capabilitiesMu.Lock() s.capabilities.resources = &resourceCapabilities{} + s.capabilitiesMu.Unlock() + } else { + s.capabilitiesMu.RUnlock() } - s.capabilitiesMu.Unlock() + s.resourcesMu.Lock() defer s.resourcesMu.Unlock() @@ -454,11 +467,16 @@ func (s *MCPServer) AddResourceTemplate( // AddPrompt registers a new prompt handler with the given name func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { - s.capabilitiesMu.Lock() + s.capabilitiesMu.RLock() if s.capabilities.prompts == nil { + s.capabilitiesMu.RUnlock() + + s.capabilitiesMu.Lock() s.capabilities.prompts = &promptCapabilities{} + s.capabilitiesMu.Unlock() + } else { + s.capabilitiesMu.RUnlock() } - s.capabilitiesMu.Unlock() s.promptsMu.Lock() defer s.promptsMu.Unlock() @@ -473,11 +491,16 @@ func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) { // AddTools registers multiple tools at once func (s *MCPServer) AddTools(tools ...ServerTool) { - s.capabilitiesMu.Lock() + s.capabilitiesMu.RLock() if s.capabilities.tools == nil { + s.capabilitiesMu.RUnlock() + + s.capabilitiesMu.Lock() s.capabilities.tools = &toolCapabilities{} + s.capabilitiesMu.Unlock() + } else { + s.capabilitiesMu.RUnlock() } - s.capabilitiesMu.Unlock() s.toolsMu.Lock() for _, entry := range tools { From 0bed1f87b65015acbe8d9a28214b4efd8118d3c3 Mon Sep 17 00:00:00 2001 From: cryo Date: Tue, 22 Apr 2025 23:54:56 +0000 Subject: [PATCH 2/2] fix double-checking for AddTools/AddPromt --- server/server.go | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/server/server.go b/server/server.go index 91256d11f..0e4811841 100644 --- a/server/server.go +++ b/server/server.go @@ -450,7 +450,9 @@ func (s *MCPServer) AddResourceTemplate( s.capabilitiesMu.RUnlock() s.capabilitiesMu.Lock() - s.capabilities.resources = &resourceCapabilities{} + if s.capabilities.resources == nil { + s.capabilities.resources = &resourceCapabilities{} + } s.capabilitiesMu.Unlock() } else { s.capabilitiesMu.RUnlock() @@ -472,7 +474,9 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { s.capabilitiesMu.RUnlock() s.capabilitiesMu.Lock() - s.capabilities.prompts = &promptCapabilities{} + if s.capabilities.prompts == nil { + s.capabilities.prompts = &promptCapabilities{} + } s.capabilitiesMu.Unlock() } else { s.capabilitiesMu.RUnlock() @@ -496,7 +500,9 @@ func (s *MCPServer) AddTools(tools ...ServerTool) { s.capabilitiesMu.RUnlock() s.capabilitiesMu.Lock() - s.capabilities.tools = &toolCapabilities{} + if s.capabilities.tools == nil { + s.capabilities.tools = &toolCapabilities{} + } s.capabilitiesMu.Unlock() } else { s.capabilitiesMu.RUnlock()