Handle WebSocket end and exceptions

Signed-off-by: Julien Ponge <jponge@redhat.com>

Author: jponge

src/main/java/kafka/vertx/demo/WebSocketServer.java

@@ -11,6 +11,7 @@
 import io.vertx.core.AbstractVerticle;
 import io.vertx.core.Promise;
 import io.vertx.core.eventbus.EventBus;
+import io.vertx.core.eventbus.MessageConsumer;
 import io.vertx.core.http.ServerWebSocket;
 import io.vertx.core.json.JsonObject;
 import io.vertx.ext.web.Router;
@@ -90,8 +91,17 @@ private void handleProduceSocket(ServerWebSocket webSocket) {
       }
     });
 
-    eventBus.<JsonObject>consumer(Main.PERIODIC_PRODUCER_BROADCAST, message ->
+    MessageConsumer<JsonObject> consumer = eventBus.<JsonObject>consumer(Main.PERIODIC_PRODUCER_BROADCAST, message ->
       eventBus.send(webSocket.textHandlerID(), message.body().encode()));
+
+    webSocket.endHandler(ended -> {
+      logger.info("WebSocket closed from {}", webSocket.remoteAddress().host());
+      consumer.unregister();
+    });
+    webSocket.exceptionHandler(err -> {
+      logger.error("WebSocket error", err);
+      consumer.unregister();
+    });
   }
 
   private void handleConsumeSocket(ServerWebSocket webSocket) {
@@ -126,5 +136,18 @@ private void handleConsumeSocket(ServerWebSocket webSocket) {
       }
     });
 
+    webSocket.endHandler(done -> {
+      logger.info("WebSocket closed from {}", webSocket.remoteAddress().host());
+      kafkaConsumer.close().onFailure(err -> {
+        logger.error("Closing Kafka consumer failed", err);
+      });
+    });
+
+    webSocket.exceptionHandler(err -> {
+      logger.error("WebSocket error", err);
+      kafkaConsumer.close().onFailure(kerr -> {
+        logger.error("Closing Kafka consumer failed", kerr);
+      });
+    });
   }
 }