From d21565a5598638d3bb1dbb647950e4eab53e9d35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dariusz=20J=C4=99drzejczyk?= Date: Fri, 6 Dec 2024 23:31:13 +0100 Subject: [PATCH] This commit is messy but it reworks the transport to work with Sinks and be non-blocking on the caller side --- .../client/stdio/StdioServerTransport.java | 37 ++++------ .../ai/mcp/spec/DefaultMcpTransport.java | 71 ++++++++----------- 2 files changed, 45 insertions(+), 63 deletions(-) diff --git a/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java b/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java index a522cd0..2907fee 100644 --- a/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java +++ b/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java @@ -8,7 +8,6 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; import spring.ai.mcp.spec.DefaultMcpTransport; import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; @@ -107,7 +106,8 @@ public class StdioServerTransport extends DefaultMcpTransport { while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) { try { System.out.println("Received error line: " + line); - this.getErrorReadQueue().offer(line, 100, TimeUnit.MILLISECONDS); + // TODO: handle errors, etc. + this.getErrorSink().tryEmitNext(line); } catch (Exception e) { throw new RuntimeException(e); @@ -150,7 +150,10 @@ public class StdioServerTransport extends DefaultMcpTransport { while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) { try { JSONRPCMessage message = deserializeJsonRpcMessage(line); - this.getDataReadQueue().offer(message, 100, TimeUnit.MILLISECONDS); + if (!this.getInboundSink().tryEmitNext(message).isSuccess()) { + // TODO: Back off, reschedule -> same comment as for writing + throw new RuntimeException("Failed to enqueue message"); + } } catch (Exception e) { throw new RuntimeException(e); @@ -170,28 +173,18 @@ public class StdioServerTransport extends DefaultMcpTransport { private void startWriterThread() { this.executorService.execute(() -> { - try { - while (this.isRunning && this.processWriter != null && !Thread.currentThread().isInterrupted()) { + this.getOutboundSink().asFlux().handle((message, s) -> { + if (message != null) { try { - JSONRPCMessage message = this.getDataWriteQueue().poll(100, TimeUnit.MILLISECONDS); - if (message != null) { - this.processWriter.write(ModelOptionsUtils.toJsonString(message)); - this.processWriter.newLine(); - this.processWriter.flush(); - } - } - catch (InterruptedException e) { - Thread.currentThread().interrupt(); - break; + this.processWriter.write(ModelOptionsUtils.toJsonString(message)); + this.processWriter.newLine(); + this.processWriter.flush(); + s.next(message); + } catch (IOException e) { + s.error(new RuntimeException(e)); } } - } - catch (IOException e) { - throw new RuntimeException(e); - } - finally { - this.isRunning = false; - } + }).subscribe(); }); } diff --git a/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java b/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java index c2dd107..3105368 100644 --- a/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java +++ b/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java @@ -16,13 +16,11 @@ package spring.ai.mcp.spec; import java.time.Duration; -import java.util.concurrent.BlockingQueue; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.LinkedBlockingQueue; -import java.util.concurrent.TimeUnit; import java.util.function.Consumer; +import reactor.core.publisher.Sinks; import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; /** @@ -31,11 +29,11 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCMessage; */ public class DefaultMcpTransport implements McpTransport { - private final BlockingQueue errorReadQueue; + private final Sinks.Many errorSink; - private final BlockingQueue dataReadQueue; + private final Sinks.Many inboundSink; - private final BlockingQueue dataWriteQueue; + private final Sinks.Many outboundSink; private final ExecutorService executorService; @@ -47,9 +45,15 @@ public class DefaultMcpTransport implements McpTransport { public DefaultMcpTransport(Duration readTimeout) { - this.errorReadQueue = new LinkedBlockingQueue<>(); - this.dataReadQueue = new LinkedBlockingQueue<>(); - this.dataWriteQueue = new LinkedBlockingQueue<>(); + // TODO: consider the effects of buffering here -> the inter-process pipes are + // independent and the notifications can flood the client/server. + // Potentially, the interest in reading could be communicated from one party + // to the other so the Blocking IO Threads can pause consuming the stream + // buffers when there is no expectation for reading. + + this.errorSink = Sinks.many().unicast().onBackpressureBuffer(); + this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); + this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.writeTimeout = readTimeout; this.executorService = Executors.newFixedThreadPool(2); @@ -58,43 +62,30 @@ public class DefaultMcpTransport implements McpTransport { private void handleIncomingMessages() { this.executorService.execute(() -> { - while (!Thread.currentThread().isInterrupted()) { - try { - JSONRPCMessage message = this.dataReadQueue.take(); - this.messageReader.accept(message); - } - catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - } + this.inboundSink.asFlux() + .subscribe(message -> this.messageReader.accept(message)); }); } private void handleIncomingErrors() { this.executorService.execute(() -> { - while (!Thread.currentThread().isInterrupted()) { - try { - String error = this.errorReadQueue.take(); - this.errorReader.accept(error); - System.err.println(error); - } - catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - } + this.errorSink.asFlux().subscribe(e -> { + this.errorReader.accept(e); + System.err.println(e); + }); }); } - protected BlockingQueue getDataReadQueue() { - return dataReadQueue; + protected Sinks.Many getInboundSink() { + return inboundSink; } - protected BlockingQueue getDataWriteQueue() { - return dataWriteQueue; + protected Sinks.Many getOutboundSink() { + return outboundSink; } - protected BlockingQueue getErrorReadQueue() { - return errorReadQueue; + protected Sinks.Many getErrorSink() { + return errorSink; } public void setMessageHandler(Consumer messageReader) { @@ -119,13 +110,11 @@ public class DefaultMcpTransport implements McpTransport { } public void sendMessage(JSONRPCMessage message) { - // Use offer with timeout to prevent blocking indefinitely - try { - if (!this.dataWriteQueue.offer(message, writeTimeout.toMillis(), TimeUnit.MILLISECONDS)) { - throw new RuntimeException("Failed to enqueue message"); - } - } catch (InterruptedException e) { - throw new RuntimeException(e); + if (!this.outboundSink.tryEmitNext(message).isSuccess()) { + // TODO: essentially we could reschedule ourselves in some time and make + // another attempt with the already read data but pause reading until + // success + throw new RuntimeException("Failed to enqueue message"); } } }