[websocket] send messages with a queue each after each (#638)

Fixes #632 Co-authored-by: Jonah Graham <[email protected]>
eclipse-lsp4j · Jun 28, 2022 · 19500c5 · 19500c5
1 parent 286d022
commit 19500c5
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 2 deletions.
diff --git a/...t.jakarta/src/main/java/org/eclipse/lsp4j/websocket/jakarta/WebSocketMessageConsumer.java b/...t.jakarta/src/main/java/org/eclipse/lsp4j/websocket/jakarta/WebSocketMessageConsumer.java
@@ -12,9 +12,13 @@
 package org.eclipse.lsp4j.websocket.jakarta;
 
 import java.io.IOException;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import jakarta.websocket.SendHandler;
+import jakarta.websocket.SendResult;
 import jakarta.websocket.Session;
 
 import org.eclipse.lsp4j.jsonrpc.JsonRpcException;
@@ -31,6 +35,8 @@ public class WebSocketMessageConsumer implements MessageConsumer {
 
 	private final Session session;
 	private final MessageJsonHandler jsonHandler;
+	private final ConcurrentLinkedQueue<String> messageQueue = new ConcurrentLinkedQueue<>();
+	private final WebSocketSendHandler handler = new WebSocketSendHandler();
 
 	public WebSocketMessageConsumer(Session session, MessageJsonHandler jsonHandler) {
 		this.session = session;
@@ -55,7 +61,8 @@ protected void sendMessage(String message) throws IOException {
 		if (session.isOpen()) {
 			int length = message.length();
 			if (length <= session.getMaxTextMessageBufferSize()) {
-				session.getAsyncRemote().sendText(message);
+				messageQueue.add(message);
+				handler.handleNextMessage();
 			} else {
 				int currentOffset = 0;
 				while (currentOffset < length) {
@@ -69,4 +76,20 @@ protected void sendMessage(String message) throws IOException {
 		}
 	}
 
+	private class WebSocketSendHandler implements SendHandler {
+		private final AtomicBoolean isSending = new AtomicBoolean();
+
+		@Override
+		public void onResult(SendResult result) {
+			isSending.set(false);
+			handleNextMessage();
+		}
+
+		void handleNextMessage() {
+			if (session.isOpen() && !messageQueue.isEmpty() && isSending.compareAndSet(false, true)) {
+				session.getAsyncRemote().sendText(messageQueue.poll(), this);
+			}
+		}
+	}
+
 }
diff --git a/...et.jakarta/src/test/java/org/eclipse/lsp4j/websocket/jakarta/test/MockConnectionTest.java b/...et.jakarta/src/test/java/org/eclipse/lsp4j/websocket/jakarta/test/MockConnectionTest.java
@@ -68,6 +68,23 @@ public void testNotifications() throws Exception {
 		Assert.assertEquals("123456", client.result);
 	}
 
+	@Test
+	public void testManyConcurrentNotifications() throws Exception {
+		String expectedResult = "";
+		for (char i = 'a'; i <= 'z'; i ++) {
+			String x = Character.toString(i);
+			client.server.notify(x);
+			expectedResult += x;
+		}
+		int expectedResultLenght = expectedResult.length();
+		try {
+			await(() -> server.result.length() == expectedResultLenght);
+		} catch (Error e) {
+			// discard this error so that the nice error displays in the assertEquals
+		}
+		Assert.assertEquals(expectedResult, server.result);
+	}
+
 	@Test
 	public void testChunkedNotification() throws Exception {
 		StringBuilder messageBuilder = new StringBuilder();

diff --git a/...e.lsp4j.websocket/src/main/java/org/eclipse/lsp4j/websocket/WebSocketMessageConsumer.java b/...e.lsp4j.websocket/src/main/java/org/eclipse/lsp4j/websocket/WebSocketMessageConsumer.java
@@ -12,8 +12,12 @@
 package org.eclipse.lsp4j.websocket;
 
 import java.io.IOException;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
 
+import javax.websocket.SendHandler;
+import javax.websocket.SendResult;
 import javax.websocket.Session;
 
 import org.eclipse.lsp4j.jsonrpc.JsonRpcException;
@@ -30,6 +34,8 @@ public class WebSocketMessageConsumer implements MessageConsumer {
 
 	private final Session session;
 	private final MessageJsonHandler jsonHandler;
+	private final ConcurrentLinkedQueue<String> messageQueue = new ConcurrentLinkedQueue<>();
+	private final WebSocketSendHandler handler = new WebSocketSendHandler();
 
 	public WebSocketMessageConsumer(Session session, MessageJsonHandler jsonHandler) {
 		this.session = session;
@@ -54,7 +60,8 @@ protected void sendMessage(String message) throws IOException {
 		if (session.isOpen()) {
 			int length = message.length();
 			if (length <= session.getMaxTextMessageBufferSize()) {
-				session.getAsyncRemote().sendText(message);
+				messageQueue.add(message);
+				handler.handleNextMessage();
 			} else {
 				int currentOffset = 0;
 				while (currentOffset < length) {
@@ -68,4 +75,20 @@ protected void sendMessage(String message) throws IOException {
 		}
 	}
 
+	private class WebSocketSendHandler implements SendHandler {
+		private final AtomicBoolean isSending = new AtomicBoolean();
+
+		@Override
+		public void onResult(SendResult result) {
+			isSending.set(false);
+			handleNextMessage();
+		}
+
+		void handleNextMessage() {
+			if (session.isOpen() && !messageQueue.isEmpty() && isSending.compareAndSet(false, true)) {
+				session.getAsyncRemote().sendText(messageQueue.poll(), this);
+			}
+		}
+	}
+
 }
diff --git a/...se.lsp4j.websocket/src/test/java/org/eclipse/lsp4j/websocket/test/MockConnectionTest.java b/...se.lsp4j.websocket/src/test/java/org/eclipse/lsp4j/websocket/test/MockConnectionTest.java
@@ -68,6 +68,23 @@ public void testNotifications() throws Exception {
 		Assert.assertEquals("123456", client.result);
 	}
 
+	@Test
+	public void testManyConcurrentNotifications() throws Exception {
+		String expectedResult = "";
+		for (char i = 'a'; i <= 'z'; i ++) {
+			String x = Character.toString(i);
+			client.server.notify(x);
+			expectedResult += x;
+		}
+		int expectedResultLenght = expectedResult.length();
+		try {
+			await(() -> server.result.length() == expectedResultLenght);
+		} catch (Error e) {
+			// discard this error so that the nice error displays in the assertEquals
+		}
+		Assert.assertEquals(expectedResult, server.result);
+	}
+
 	@Test
 	public void testChunkedNotification() throws Exception {
 		StringBuilder messageBuilder = new StringBuilder();