Support consumer update command

References rabbitmq/rabbitmq-server#3753

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/Constants.java

@@ -66,6 +66,7 @@ public final class Constants {
   public static final short COMMAND_HEARTBEAT = 23;
   public static final short COMMAND_ROUTE = 24;
   public static final short COMMAND_PARTITIONS = 25;
+  public static final short COMMAND_CONSUMER_UPDATE = 26;
 
   public static final short VERSION_1 = 1;
 

src/main/java/com/rabbitmq/stream/OffsetSpecification.java

@@ -18,6 +18,7 @@
 /** API to specify an offset in a stream. */
 public class OffsetSpecification {
 
+  private static final short TYPE_NONE = 0;
   private static final short TYPE_FIRST = 1;
   private static final short TYPE_LAST = 2;
   private static final short TYPE_NEXT = 3;
@@ -26,6 +27,7 @@ public class OffsetSpecification {
 
   private static final long UNUSED_OFFSET = -1;
 
+  private static final OffsetSpecification NONE = new OffsetSpecification(TYPE_NONE, UNUSED_OFFSET);
   private static final OffsetSpecification FIRST =
       new OffsetSpecification(TYPE_FIRST, UNUSED_OFFSET);
   private static final OffsetSpecification LAST = new OffsetSpecification(TYPE_LAST, UNUSED_OFFSET);
@@ -39,6 +41,15 @@ private OffsetSpecification(short type, long offset) {
     this.offset = offset;
   }
 
+  /**
+   * When the offset specification is not relevant.
+   *
+   * @return none offset specification
+   */
+  public static OffsetSpecification none() {
+    return NONE;
+  }
+
   /**
    * The first available offset in the stream.
    *

src/main/java/com/rabbitmq/stream/impl/Client.java

@@ -14,6 +14,7 @@
 package com.rabbitmq.stream.impl;
 
 import static com.rabbitmq.stream.Constants.COMMAND_CLOSE;
+import static com.rabbitmq.stream.Constants.COMMAND_CONSUMER_UPDATE;
 import static com.rabbitmq.stream.Constants.COMMAND_CREATE_STREAM;
 import static com.rabbitmq.stream.Constants.COMMAND_CREDIT;
 import static com.rabbitmq.stream.Constants.COMMAND_DECLARE_PUBLISHER;
@@ -39,6 +40,7 @@
 import static com.rabbitmq.stream.Constants.RESPONSE_CODE_SASL_CHALLENGE;
 import static com.rabbitmq.stream.Constants.VERSION_1;
 import static com.rabbitmq.stream.impl.Utils.encodeRequestCode;
+import static com.rabbitmq.stream.impl.Utils.encodeResponseCode;
 import static com.rabbitmq.stream.impl.Utils.extractResponseCode;
 import static com.rabbitmq.stream.impl.Utils.formatConstant;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -159,6 +161,7 @@ public class Client implements AutoCloseable {
   final ChunkListener chunkListener;
   final MessageListener messageListener;
   final CreditNotification creditNotification;
+  final ConsumerUpdateListener consumerUpdateListener;
   final MetadataListener metadataListener;
   final Codec codec;
   final Channel channel;
@@ -216,6 +219,7 @@ public Client(ClientParameters parameters) {
     this.chunkChecksum = parameters.chunkChecksum;
     this.metricsCollector = parameters.metricsCollector;
     this.metadataListener = parameters.metadataListener;
+    this.consumerUpdateListener = parameters.consumerUpdateListener;
     this.compressionCodecFactory =
         parameters.compressionCodecFactory == null
             ? compression -> null
@@ -1279,7 +1283,7 @@ public List<String> partitions(String superStream) {
       throw new IllegalArgumentException("stream must not be null");
     }
     int length =
-        2 + 2 + 4 + +2 + superStream.length(); // API code, version, correlation ID, 1 string
+        2 + 2 + 4 + 2 + superStream.length(); // API code, version, correlation ID, 1 string
     int correlationId = correlationSequence.incrementAndGet();
     try {
       ByteBuf bb = allocate(length + 4);
@@ -1300,6 +1304,29 @@ public List<String> partitions(String superStream) {
     }
   }
 
+  public void consumerUpdateResponse(
+      int correlationId, short responseCode, OffsetSpecification offsetSpecification) {
+    offsetSpecification =
+        offsetSpecification == null ? OffsetSpecification.none() : offsetSpecification;
+    int length = 2 + 2 + 4 + 2 + 2; // API code, version, correlation ID, response code, offset type
+
+    if (offsetSpecification.isOffset() || offsetSpecification.isTimestamp()) {
+      length += 8;
+    }
+
+    ByteBuf bb = allocate(length + 4);
+    bb.writeInt(length);
+    bb.writeShort(encodeResponseCode(COMMAND_CONSUMER_UPDATE));
+    bb.writeShort(VERSION_1);
+    bb.writeInt(correlationId);
+    bb.writeShort(responseCode);
+    bb.writeShort(offsetSpecification.getType());
+    if (offsetSpecification.isOffset() || offsetSpecification.isTimestamp()) {
+      bb.writeLong(offsetSpecification.getOffset());
+    }
+    channel.writeAndFlush(bb);
+  }
+
   void shutdownReason(ShutdownReason reason) {
     this.shutdownReason = reason;
   }
@@ -1397,6 +1424,11 @@ public interface CreditNotification {
     void handle(byte subscriptionId, short responseCode);
   }
 
+  public interface ConsumerUpdateListener {
+
+    OffsetSpecification handle(Client client, byte subscriptionId, boolean active);
+  }
+
   public interface ShutdownListener {
 
     void handle(ShutdownContext shutdownContext);
@@ -1918,6 +1950,8 @@ public static class ClientParameters {
                 "Received notification for subscription {}: {}",
                 subscriptionId,
                 Utils.formatConstant(responseCode));
+    private ConsumerUpdateListener consumerUpdateListener =
+        (client, subscriptionId, active) -> null;
     private ShutdownListener shutdownListener = shutdownContext -> {};
     private SaslConfiguration saslConfiguration = DefaultSaslConfiguration.PLAIN;
     private CredentialsProvider credentialsProvider =
@@ -1965,6 +1999,11 @@ public ClientParameters creditNotification(CreditNotification creditNotification
       return this;
     }
 
+    public ClientParameters consumerUpdateListener(ConsumerUpdateListener consumerUpdateListener) {
+      this.consumerUpdateListener = consumerUpdateListener;
+      return this;
+    }
+
     public ClientParameters codec(Codec codec) {
       this.codec = codec;
       return this;

src/main/java/com/rabbitmq/stream/impl/ServerFrameHandler.java

@@ -14,6 +14,7 @@
 package com.rabbitmq.stream.impl;
 
 import static com.rabbitmq.stream.Constants.COMMAND_CLOSE;
+import static com.rabbitmq.stream.Constants.COMMAND_CONSUMER_UPDATE;
 import static com.rabbitmq.stream.Constants.COMMAND_CREATE_STREAM;
 import static com.rabbitmq.stream.Constants.COMMAND_CREDIT;
 import static com.rabbitmq.stream.Constants.COMMAND_DECLARE_PUBLISHER;
@@ -46,6 +47,7 @@
 import com.rabbitmq.stream.ChunkChecksumValidationException;
 import com.rabbitmq.stream.Codec;
 import com.rabbitmq.stream.Message;
+import com.rabbitmq.stream.OffsetSpecification;
 import com.rabbitmq.stream.StreamException;
 import com.rabbitmq.stream.compression.Compression;
 import com.rabbitmq.stream.compression.CompressionCodec;
@@ -112,6 +114,7 @@ class ServerFrameHandler {
     handlers.put(COMMAND_QUERY_PUBLISHER_SEQUENCE, new QueryPublisherSequenceFrameHandler());
     handlers.put(COMMAND_ROUTE, new RouteFrameHandler());
     handlers.put(COMMAND_PARTITIONS, new PartitionsFrameHandler());
+    handlers.put(COMMAND_CONSUMER_UPDATE, new ConsumerUpdateFrameHandler());
     HANDLERS = new FrameHandler[handlers.size() + 10];
     handlers.entrySet().forEach(entry -> HANDLERS[entry.getKey()] = entry.getValue());
   }
@@ -575,6 +578,27 @@ int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) {
     }
   }
 
+  private static class ConsumerUpdateFrameHandler extends BaseFrameHandler {
+
+    @Override
+    int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) {
+      int correlationId = message.readInt();
+      int read = 4;
+      byte subscriptionId = message.readByte();
+      read += 1;
+      byte activeByte = message.readByte();
+      read += 1;
+
+      OffsetSpecification offsetSpecification =
+          client.consumerUpdateListener.handle(
+              client, subscriptionId, activeByte == 1 ? true : false);
+
+      client.consumerUpdateResponse(correlationId, RESPONSE_CODE_OK, offsetSpecification);
+
+      return read;
+    }
+  }
+
   private static class PeerPropertiesFrameHandler extends BaseFrameHandler {
 
     @Override

src/test/java/com/rabbitmq/stream/impl/SingleActiveConsumerTest.java

@@ -0,0 +1,80 @@
+// Copyright (c) 2021 VMware, Inc. or its affiliates.  All rights reserved.
+//
+// This software, the RabbitMQ Stream Java client library, is dual-licensed under the
+// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
+// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.stream.impl;
+
+import static com.rabbitmq.stream.impl.TestUtils.b;
+import static com.rabbitmq.stream.impl.TestUtils.latchAssert;
+import static com.rabbitmq.stream.impl.TestUtils.waitAtMost;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.rabbitmq.stream.OffsetSpecification;
+import com.rabbitmq.stream.impl.Client.ClientParameters;
+import com.rabbitmq.stream.impl.Client.Response;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@ExtendWith(TestUtils.StreamTestInfrastructureExtension.class)
+public class SingleActiveConsumerTest {
+
+  String stream;
+  TestUtils.ClientFactory cf;
+
+  @Test
+  void subscribe() throws Exception {
+    int messageCount = 10000;
+    Map<Byte, Boolean> consumerStates = new ConcurrentHashMap<>();
+    Map<Byte, AtomicInteger> receivedMessages = new ConcurrentHashMap<>();
+    receivedMessages.put(b(0), new AtomicInteger(0));
+    receivedMessages.put(b(1), new AtomicInteger(0));
+    CountDownLatch consumerUpdateLatch = new CountDownLatch(2);
+    ClientParameters clientParameters =
+        new ClientParameters()
+            .chunkListener(
+                (client, subscriptionId, offset, msgCount, dataSize) ->
+                    client.credit(subscriptionId, 1))
+            .messageListener(
+                (subscriptionId, offset, chunkTimestamp, message) ->
+                    receivedMessages.get(subscriptionId).incrementAndGet())
+            .consumerUpdateListener(
+                (client, subscriptionId, active) -> {
+                  consumerStates.put(subscriptionId, active);
+                  consumerUpdateLatch.countDown();
+                  return null;
+                });
+    Client client = cf.get(clientParameters);
+
+    TestUtils.publishAndWaitForConfirms(cf, messageCount, stream);
+
+    Map<String, String> parameters = new HashMap<>();
+    parameters.put("single-active-consumer", "true");
+    parameters.put("name", "foo");
+    Response response = client.subscribe(b(0), stream, OffsetSpecification.first(), 2, parameters);
+    assertThat(response.isOk()).isTrue();
+    response = client.subscribe(b(1), stream, OffsetSpecification.first(), 2, parameters);
+    assertThat(response.isOk()).isTrue();
+    latchAssert(consumerUpdateLatch).completes();
+    assertThat(consumerStates)
+        .hasSize(2)
+        .containsEntry(b(0), Boolean.TRUE)
+        .containsEntry(b(1), Boolean.FALSE);
+
+    waitAtMost(
+        () -> receivedMessages.getOrDefault(b(0), new AtomicInteger(0)).get() == messageCount);
+  }
+}

src/test/java/com/rabbitmq/stream/impl/TestUtils.java

@@ -170,12 +170,7 @@ static void publishAndWaitForConfirms(
       client.publish(b(1), Collections.singletonList(message));
     }
 
-    try {
-      assertThat(latchConfirm.await(60, SECONDS)).isTrue();
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-      throw new RuntimeException(e);
-    }
+    latchAssert(latchConfirm).completes(Duration.ofSeconds(60));
   }
 
   static Consumer<Object> namedTask(TaskWithException task, String description) {