feat: Refactor flow control of consumer + producer (#19)

Refactor the flow of the consumer and producer.
Stop the producer when the websocket disconnects.
Only subscribe to the consumer when start is sent.

Signed-off-by: Katherine Stanley <katheris@uk.ibm.com>

Author: katheris

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

@@ -48,12 +48,12 @@ private void setup(HashMap<String, String> props) {
   }
 
   private void handleCommand(TimeoutStream timerStream, Message<JsonObject> message) {
-    String command = message.body().getString("action", "none");
-    if ("start".equals(command)) {
+    String command = message.body().getString(WebSocketServer.ACTION, "none");
+    if (WebSocketServer.START_ACTION.equals(command)) {
       logger.info("Producing Kafka records");
       customMessage = message.body().getString("custom", "Hello World");
       timerStream.resume();
-    } else if ("stop".equals(command)) {
+    } else if (WebSocketServer.STOP_ACTION.equals(command)) {
       logger.info("Stopping producing Kafka records");
       timerStream.pause();
     }
@@ -62,6 +62,7 @@ private void handleCommand(TimeoutStream timerStream, Message<JsonObject> messag
   private void produceKafkaRecord(KafkaProducer<String, String> kafkaProducer, String topic) {
     String payload = customMessage;
     KafkaProducerRecord<String, String> record = KafkaProducerRecord.create(topic, payload);
+    logger.debug("Producing record to topic {} with payload {}", topic, payload);
 
     kafkaProducer
       .send(record)

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

@@ -21,9 +21,14 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.HashMap;
+import java.util.Set;
 
 public class WebSocketServer extends AbstractVerticle {
 
+  public static final String ACTION = "action";
+  public static final String START_ACTION = "start";
+  public static final String STOP_ACTION = "stop";
+
   private static final String PRODUCE_PATH = "/demoproduce";
   private static final String CONSUME_PATH = "/democonsume";
 
@@ -73,8 +78,8 @@ private void handleProduceSocket(ServerWebSocket webSocket) {
 
     webSocket.handler(buffer -> {
       JsonObject message = buffer.toJsonObject();
-      String action = message.getString("action", "none");
-      if ("start".equals(action) || "stop".equals(action)) {
+      String action = message.getString(ACTION, "none");
+      if (START_ACTION.equals(action) || STOP_ACTION.equals(action)) {
         eventBus.send(Main.PERIODIC_PRODUCER_ADDRESS, message);
       }
     });
@@ -83,11 +88,13 @@ private void handleProduceSocket(ServerWebSocket webSocket) {
       eventBus.send(webSocket.textHandlerID(), message.body().encode()));
 
     webSocket.endHandler(ended -> {
-      logger.info("WebSocket closed from {}", webSocket.remoteAddress().host());
+      logger.info("Producer WebSocket closed from {}", webSocket.remoteAddress().host());
+      eventBus.send(Main.PERIODIC_PRODUCER_ADDRESS, new JsonObject().put(ACTION, STOP_ACTION));
       consumer.unregister();
     });
     webSocket.exceptionHandler(err -> {
-      logger.error("WebSocket error", err);
+      logger.error("Producer WebSocket error", err);
+      eventBus.send(Main.PERIODIC_PRODUCER_ADDRESS, new JsonObject().put(ACTION, STOP_ACTION));
       consumer.unregister();
     });
   }
@@ -97,47 +104,43 @@ private void handleConsumeSocket(ServerWebSocket webSocket) {
 
     kafkaConsumer.exceptionHandler(err -> logger.error("Kafka error", err));
 
+    String topic = kafkaConfig.get(Main.TOPIC_KEY);
+    TopicPartition topicPartition = new TopicPartition().setTopic(topic);
+
     kafkaConsumer.handler(record -> {
       JsonObject payload = new JsonObject()
         .put("topic", record.topic())
         .put("partition", record.partition())
         .put("offset", record.offset())
         .put("timestamp", record.timestamp())
         .put("value", record.value());
+      logger.debug("Received record {}", payload);
       vertx.eventBus().send(webSocket.textHandlerID(), payload.encode());
     });
 
-    String topic = kafkaConfig.get(Main.TOPIC_KEY);
-
-    kafkaConsumer.subscribe(topic)
-      .onSuccess(ok -> logger.info("Subscribed to {}", topic))
-      .onFailure(err -> logger.error("Could not subscribe to {}", topic, err));
-
-    TopicPartition partition = new TopicPartition().setTopic(topic);
-
     webSocket.handler(buffer -> {
-      String action = buffer.toJsonObject().getString("action", "none");
-      if ("start".equals(action)) {
-        kafkaConsumer.resume(partition)
-          .onFailure(err -> logger.error("Cannot resume consumer", err));
-      } else if ("stop".equals(action)) {
-        kafkaConsumer.pause(partition)
+      String action = buffer.toJsonObject().getString(ACTION, "none");
+      if (START_ACTION.equals(action)) {
+        kafkaConsumer.subscription()
+          .compose(sub -> (sub.size() > 0) ? kafkaConsumer.resume(topicPartition) : kafkaConsumer.subscribe(topic))
+            .onSuccess(ok -> logger.info("Subscribed to {}", topic))
+            .onFailure(err -> logger.error("Could not subscribe to {}", topic, err));
+      } else if (STOP_ACTION.equals(action)) {
+        kafkaConsumer.pause(topicPartition)
           .onFailure(err -> logger.error("Cannot pause consumer", err));
       }
     });
 
     webSocket.endHandler(done -> {
-      logger.info("WebSocket closed from {}", webSocket.remoteAddress().host());
-      kafkaConsumer.close().onFailure(err -> {
-        logger.error("Closing Kafka consumer failed", err);
-      });
+      logger.info("Consumer 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);
-      });
+      logger.error("Consumer WebSocket error", err);
+      kafkaConsumer.close()
+        .onFailure(kerr -> logger.error("Closing Kafka consumer failed", kerr));
     });
   }
 }

src/main/resources/logback.xml

@@ -10,7 +10,7 @@
   <logger name="io.vertx" level="info"/>
   <logger name="org.apache.kafka" level="error"/>
 
-  <root level="debug">
+  <root level="info">
     <appender-ref ref="STDOUT"/>
   </root>