WIP, will force push later

Author: henry701

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

@@ -5,12 +5,12 @@
 import java.util.Objects;
 import java.util.function.Supplier;
 
-public class AbstractFlowControlStrategy implements ConsumerFlowControlStrategy {
+public abstract class AbstractFlowControlStrategy implements ConsumerFlowControlStrategy {
 
-    private Supplier<Client> clientSupplier;
+    private final Supplier<Client> clientSupplier;
     private volatile Client client;
 
-    public AbstractFlowControlStrategy(Supplier<Client> clientSupplier) {
+    protected AbstractFlowControlStrategy(Supplier<Client> clientSupplier) {
         this.clientSupplier = Objects.requireNonNull(clientSupplier, "clientSupplier");
     }
 

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

@@ -3,7 +3,8 @@
 import com.rabbitmq.stream.impl.Client;
 
 /**
- * Exposes callbacks to handle events from a specific Client, with a leaner interface.
+ * Exposes callbacks to handle events from a particular {@link Client},
+ * with specific names for methods and no {@link Client} parameter.
  */
 public interface ClientDataHandler extends
         Client.PublishConfirmListener,

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

@@ -2,13 +2,41 @@
 
 import com.rabbitmq.stream.impl.Client;
 
+import java.util.Map;
+
 /**
  * A built and configured flow control strategy for consumers.
  * Implementations may freely implement reactions to the various client callbacks.
  * When defined by each implementation, it may internally call {@link Client#credit} to ask for credits.
  */
 public interface ConsumerFlowControlStrategy extends ClientDataHandler, AutoCloseable {
 
+    /**
+     * Callback for handling a new stream subscription.
+     * Called right before the subscription is sent to the actual client.
+     *
+     * @param subscriptionId The subscriptionId as specified by the Stream Protocol
+     * @param stream The name of the stream being subscribed to
+     * @param offsetSpecification The offset specification for this new subscription
+     * @param subscriptionProperties The subscription properties for this new subscription
+     * @return The initial credits that should be granted to this new subscription
+     */
+    int handleSubscribe(
+        byte subscriptionId,
+        String stream,
+        OffsetSpecification offsetSpecification,
+        Map<String, String> subscriptionProperties
+    );
+
+    /**
+     * Callback for handling a stream unsubscription.
+     * @param subscriptionId The subscriptionId as specified by the Stream Protocol
+     */
+    default void handleUnsubscribe(byte subscriptionId) {
+        // No-op by default
+    }
+
+    @Override
     default void handleShutdown(Client.ShutdownContext shutdownContext) {
         this.close();
     }

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

@@ -2,6 +2,7 @@
 
 import com.rabbitmq.stream.impl.Client;
 
+import java.util.Map;
 import java.util.function.Supplier;
 
 /**
@@ -10,13 +11,33 @@
  */
 public class LegacyFlowControlStrategy extends AbstractFlowControlStrategy {
 
+    private final int initialCredits;
     private final int additionalCredits;
 
-    public LegacyFlowControlStrategy(Supplier<Client> clientSupplier, int additionalCredits) {
+    public LegacyFlowControlStrategy(Supplier<Client> clientSupplier) {
+        this(clientSupplier, 1);
+    }
+
+    public LegacyFlowControlStrategy(Supplier<Client> clientSupplier, int initialCredits) {
+        this(clientSupplier, initialCredits, 1);
+    }
+
+    public LegacyFlowControlStrategy(Supplier<Client> clientSupplier, int initialCredits, int additionalCredits) {
         super(clientSupplier);
+        this.initialCredits = initialCredits;
         this.additionalCredits = additionalCredits;
     }
 
+    @Override
+    public int handleSubscribe(
+            byte subscriptionId,
+            String stream,
+            OffsetSpecification offsetSpecification,
+            Map<String, String> subscriptionProperties
+    ) {
+        return this.initialCredits;
+    }
+
     @Override
     public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) {
         mandatoryClient().credit(subscriptionId, this.additionalCredits);

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

@@ -17,12 +17,19 @@ public static class LegacyFlowControlStrategyBuilder implements ConsumerFlowCont
 
         private final ConsumerBuilder consumerBuilder;
 
+        private int initialCredits = 1;
+
         private int additionalCredits = 1;
 
         public LegacyFlowControlStrategyBuilder(ConsumerBuilder consumerBuilder) {
             this.consumerBuilder = consumerBuilder;
         }
 
+        @Override
+        public LegacyFlowControlStrategy build(Supplier<Client> clientSupplier) {
+            return new LegacyFlowControlStrategy(clientSupplier, this.initialCredits, this.additionalCredits);
+        }
+
         @Override
         public ConsumerBuilder builder() {
             return this.consumerBuilder;
@@ -33,11 +40,10 @@ public LegacyFlowControlStrategyBuilder additionalCredits(int additionalCredits)
             return this;
         }
 
-        @Override
-        public LegacyFlowControlStrategy build(Supplier<Client> clientSupplier) {
-            return new LegacyFlowControlStrategy(clientSupplier, this.additionalCredits);
+        public LegacyFlowControlStrategyBuilder initialCredits(int initialCredits) {
+            this.initialCredits = initialCredits;
+            return this;
         }
-
     }
 
 }

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

@@ -543,6 +543,7 @@ private class ClientSubscriptionsManager implements Comparable<ClientSubscriptio
     // the 2 data structures track the subscriptions, they must remain consistent
     private final Map<String, Set<SubscriptionTracker>> streamToStreamSubscriptions =
         new ConcurrentHashMap<>();
+    private final ConsumerFlowControlStrategy consumerFlowControlStrategy;
     // trackers and tracker count must be kept in sync
     private volatile List<SubscriptionTracker> subscriptionTrackers =
         new ArrayList<>(maxConsumersByConnection);
@@ -561,9 +562,10 @@ private ClientSubscriptionsManager(
       IntStream.range(0, maxConsumersByConnection).forEach(i -> subscriptionTrackers.add(null));
       this.trackerCount = 0;
       AtomicReference<Client> clientReference = new AtomicReference<>();
-      ConsumerFlowControlStrategy consumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get);
+      ConsumerFlowControlStrategy localConsumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get);
+      this.consumerFlowControlStrategy = localConsumerFlowControlStrategy;
       ChunkListener chunkListener = (ignoredClient, subscriptionId, offset, messageCount, dataSize) ->
-              consumerFlowControlStrategy.handleChunk(subscriptionId, offset, messageCount, dataSize);
+              localConsumerFlowControlStrategy.handleChunk(subscriptionId, offset, messageCount, dataSize);
       CreditNotification creditNotification =
           (subscriptionId, responseCode) -> {
             SubscriptionTracker subscriptionTracker =
@@ -574,7 +576,7 @@ private ClientSubscriptionsManager(
                 subscriptionId & 0xFF,
                 stream,
                 Utils.formatConstant(responseCode));
-            consumerFlowControlStrategy.handleCreditNotification(subscriptionId, responseCode);
+            localConsumerFlowControlStrategy.handleCreditNotification(subscriptionId, responseCode);
           };
 
       MessageListener messageListener =
@@ -596,7 +598,7 @@ private ClientSubscriptionsManager(
                   this.id,
                   this.name);
             }
-            consumerFlowControlStrategy.handleMessage(subscriptionId, offset, chunkTimestamp, committedOffset, message);
+            localConsumerFlowControlStrategy.handleMessage(subscriptionId, offset, chunkTimestamp, committedOffset, message);
           };
       ShutdownListener shutdownListener =
           shutdownContext -> {
@@ -650,7 +652,7 @@ private ClientSubscriptionsManager(
                           "Consumers re-assignment after disconnection from %s",
                           name));
             }
-            consumerFlowControlStrategy.handleShutdown(shutdownContext);
+            localConsumerFlowControlStrategy.handleShutdown(shutdownContext);
           };
       MetadataListener metadataListener =
           (stream, code) -> {
@@ -702,7 +704,7 @@ private ClientSubscriptionsManager(
                           "Consumers re-assignment after metadata update on stream '%s'",
                           stream));
             }
-            consumerFlowControlStrategy.handleMetadata(stream, code);
+            localConsumerFlowControlStrategy.handleMetadata(stream, code);
           };
       ConsumerUpdateListener consumerUpdateListener =
           (client, subscriptionId, active) -> {
@@ -721,7 +723,7 @@ private ClientSubscriptionsManager(
               LOGGER.debug(
                   "Could not find stream subscription {} for consumer update", subscriptionId);
             }
-            consumerFlowControlStrategy.handleConsumerUpdate(subscriptionId, active);
+            localConsumerFlowControlStrategy.handleConsumerUpdate(subscriptionId, active);
             return result;
           };
       String connectionName = connectionNamingStrategy.apply(ClientConnectionType.CONSUMER);
@@ -954,14 +956,20 @@ synchronized void add(
 
         checkNotClosed();
         byte subId = subscriptionId;
+        int initialCredits = this.consumerFlowControlStrategy.handleSubscribe(
+                subId,
+                subscriptionTracker.stream,
+                subscriptionContext.offsetSpecification(),
+                subscriptionTracker.subscriptionProperties
+        );
         Client.Response subscribeResponse =
             Utils.callAndMaybeRetry(
                 () ->
                     client.subscribe(
                         subId,
                         subscriptionTracker.stream,
                         subscriptionContext.offsetSpecification(),
-                        subscriptionTracker.initialCredits,
+                        initialCredits,
                         subscriptionTracker.subscriptionProperties),
                 RETRY_ON_TIMEOUT,
                 "Subscribe request for consumer %s on stream '%s'",
@@ -1025,7 +1033,7 @@ synchronized void remove(SubscriptionTracker subscriptionTracker) {
             subscriptionTracker.consumer.id(),
             subscriptionTracker.stream);
       }
-
+      this.consumerFlowControlStrategy.handleUnsubscribe(subscriptionIdInClient);
       this.setSubscriptionTrackers(update(this.subscriptionTrackers, subscriptionIdInClient, null));
       streamToStreamSubscriptions.compute(
           subscriptionTracker.stream,
@@ -1090,6 +1098,7 @@ synchronized void close() {
                 if (this.client != null && this.client.isOpen() && tracker.consumer.isOpen()) {
                   this.client.unsubscribe(tracker.subscriptionIdInClient);
                 }
+                this.consumerFlowControlStrategy.handleUnsubscribe(tracker.subscriptionIdInClient);
               } catch (Exception e) {
                 // OK, moving on
                 LOGGER.debug(

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

@@ -135,12 +135,23 @@ public ConsumerBuilder noTrackingStrategy() {
     return this;
   }
 
-  public ConsumerBuilder credits(int initial, int onChunkDelivery) {
+  /**
+   *
+   * @param initial Credits to ask for with each new subscription
+   * @param onChunkDelivery Credits to ask for after a chunk is delivered
+   * @return this {@link StreamConsumerBuilder}
+   * @deprecated Prefer using {@link ConsumerBuilder#flowControlStrategy(ConsumerFlowControlStrategyBuilderFactory)}
+   *             to define flow control strategies instead.
+   */
+  @Deprecated
+  public StreamConsumerBuilder credits(int initial, int onChunkDelivery) {
     if (initial <= 0 || onChunkDelivery <= 0) {
       throw new IllegalArgumentException("Credits must be positive");
     }
-    this.initialCredits = initial;
-    this.additionalCredits = onChunkDelivery;
+    this.consumerFlowControlStrategyBuilder = LegacyFlowControlStrategyBuilderFactory.INSTANCE
+            .builder(this)
+            .initialCredits(initial)
+            .additionalCredits(additionalCredits);
     return this;
   }
 

src/test/java/com/rabbitmq/stream/Host.java

@@ -43,7 +43,6 @@ public static String capture(InputStream is) throws IOException {
 
   private static Process executeCommand(String command) throws IOException {
     Process pr = executeCommandProcess(command);
-
     int ev = waitForExitValue(pr);
     if (ev != 0) {
       String stdout = capture(pr.getInputStream());
@@ -83,10 +82,10 @@ private static Process executeCommandProcess(String command) throws IOException
     String[] finalCommand;
     if (System.getProperty("os.name").toLowerCase().contains("windows")) {
       finalCommand = new String[4];
-      finalCommand[0] = "C:\\winnt\\system32\\cmd.exe";
+      finalCommand[0] = "C:\\Windows\\system32\\cmd.exe";
       finalCommand[1] = "/y";
       finalCommand[2] = "/c";
-      finalCommand[3] = command;
+      finalCommand[3] = command.replaceAll("\"", "\"\"\"").replaceAll("'", "\"");
     } else {
       finalCommand = new String[3];
       finalCommand[0] = "/bin/sh";