Remove finished streams

niloc132 · Jul 11, 2022 · b7ddeac · b7ddeac
1 parent 7ad1dfb
commit b7ddeac
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/...jakarta/src/main/java/io/grpc/servlet/web/websocket/MultiplexedWebSocketServerStream.java b/...jakarta/src/main/java/io/grpc/servlet/web/websocket/MultiplexedWebSocketServerStream.java
@@ -93,6 +93,7 @@ public void onMessage(String message) {
             // End the stream first.
             stream.transportReportStatus(Status.fromCode(Status.Code.UNKNOWN));
         }
+        streams.clear();
         try {
             websocketSession
                     .close(new CloseReason(CloseReason.CloseCodes.PROTOCOL_ERROR, "Can't read string payloads"));
@@ -120,6 +121,7 @@ public void onMessage(ByteBuffer message) throws IOException {
             // message is empty (no control flow, no data), error
             if (stream != null) {
                 stream.transportReportStatus(Status.fromCode(Status.Code.UNKNOWN));
+                streams.remove(streamId);
             }
             websocketSession.close(new CloseReason(CloseReason.CloseCodes.PROTOCOL_ERROR, "Unexpected empty message"));
             return;
@@ -139,11 +141,13 @@ public void onMessage(ByteBuffer message) throws IOException {
             // if first byte is 1, the client is finished sending
             if (message.remaining() != 0) {
                 stream.transportReportStatus(Status.fromCode(Status.Code.UNKNOWN));
+                streams.remove(streamId);
                 websocketSession.close(
                         new CloseReason(CloseReason.CloseCodes.PROTOCOL_ERROR, "Unexpected bytes in close message"));
                 return;
             }
             stream.inboundDataReceived(ReadableBuffers.empty(), true);
+            streams.remove(streamId);
             return;
         }
         assert !closed;
@@ -161,6 +165,7 @@ public void onError(Throwable error) {
         for (MultiplexedWebsocketStreamImpl stream : streams.values()) {
             stream.transportReportStatus(Status.UNKNOWN);// transport failure of some kind
         }
+        streams.clear();
         // onClose will be called automatically
         if (error instanceof ClosedChannelException) {
             // ignore this for now