Expose subscription credit settings

Fixes #262

Author: acogoluegnes

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

@@ -151,6 +151,18 @@ public interface ConsumerBuilder {
    */
   ConsumerBuilder noTrackingStrategy();
 
+  /**
+   * Define credit values for the subscription.
+   *
+   * <p>This is an experimental API, subject to change.
+   *
+   * @param initial the initial number of credits, default is 10
+   * @param onChunkDelivery the number of credits provided on each chunk delivery, default is 1
+   * @return this builder instance
+   * @since 0.10.0
+   */
+  ConsumerBuilder credits(int initial, int onChunkDelivery);
+
   /**
    * Create the configured {@link Consumer}
    *

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

@@ -118,7 +118,9 @@ Runnable subscribe(
       SubscriptionListener subscriptionListener,
       Runnable trackingClosingCallback,
       MessageHandler messageHandler,
-      Map<String, String> subscriptionProperties) {
+      Map<String, String> subscriptionProperties,
+      int initialCredits,
+      int additionalCredits) {
     List<Client.Broker> candidates = findBrokersForStream(stream);
     Client.Broker newNode = pickBroker(candidates);
     if (newNode == null) {
@@ -137,7 +139,9 @@ Runnable subscribe(
             subscriptionListener,
             trackingClosingCallback,
             messageHandler,
-            subscriptionProperties);
+            subscriptionProperties,
+            initialCredits,
+            additionalCredits);
 
     try {
       addToManager(newNode, subscriptionTracker, offsetSpecification, true);
@@ -387,6 +391,8 @@ private static class SubscriptionTracker {
     private volatile ClientSubscriptionsManager manager;
     private volatile AtomicReference<SubscriptionState> state =
         new AtomicReference<>(SubscriptionState.OPENING);
+    private final int initialCredits;
+    private final int additionalCredits;
 
     private SubscriptionTracker(
         long id,
@@ -397,7 +403,9 @@ private SubscriptionTracker(
         SubscriptionListener subscriptionListener,
         Runnable trackingClosingCallback,
         MessageHandler messageHandler,
-        Map<String, String> subscriptionProperties) {
+        Map<String, String> subscriptionProperties,
+        int initialCredits,
+        int additionalCredits) {
       this.id = id;
       this.consumer = consumer;
       this.stream = stream;
@@ -406,6 +414,8 @@ private SubscriptionTracker(
       this.subscriptionListener = subscriptionListener;
       this.trackingClosingCallback = trackingClosingCallback;
       this.messageHandler = messageHandler;
+      this.initialCredits = initialCredits;
+      this.additionalCredits = additionalCredits;
       if (this.offsetTrackingReference == null) {
         this.subscriptionProperties = subscriptionProperties;
       } else {
@@ -548,7 +558,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
             SubscriptionTracker subscriptionTracker =
                 subscriptionTrackers.get(subscriptionId & 0xFF);
             if (subscriptionTracker != null && subscriptionTracker.consumer.isOpen()) {
-              client.credit(subscriptionId, 1);
+              client.credit(subscriptionId, subscriptionTracker.additionalCredits);
             } else {
               LOGGER.debug(
                   "Could not find stream subscription {} or subscription closing, not providing credits",
@@ -943,7 +953,7 @@ synchronized void add(
                         subId,
                         subscriptionTracker.stream,
                         subscriptionContext.offsetSpecification(),
-                        10,
+                        subscriptionTracker.initialCredits,
                         subscriptionTracker.subscriptionProperties),
                 RETRY_ON_TIMEOUT,
                 "Subscribe request for consumer %s on stream '%s'",

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

@@ -80,7 +80,9 @@ class StreamConsumer implements Consumer {
       boolean lazyInit,
       SubscriptionListener subscriptionListener,
       Map<String, String> subscriptionProperties,
-      ConsumerUpdateListener consumerUpdateListener) {
+      ConsumerUpdateListener consumerUpdateListener,
+      int initialCredits,
+      int additionalCredits) {
 
     this.id = ID_SEQUENCE.getAndIncrement();
     Runnable trackingClosingCallback;
@@ -253,7 +255,9 @@ class StreamConsumer implements Consumer {
                     subscriptionListener,
                     trackingClosingCallback,
                     closedAwareMessageHandler,
-                    Collections.unmodifiableMap(subscriptionProperties));
+                    Collections.unmodifiableMap(subscriptionProperties),
+                    initialCredits,
+                    additionalCredits);
 
             this.status = Status.RUNNING;
           };

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

@@ -42,6 +42,8 @@ class StreamConsumerBuilder implements ConsumerBuilder {
   private SubscriptionListener subscriptionListener = subscriptionContext -> {};
   private Map<String, String> subscriptionProperties = new ConcurrentHashMap<>();
   private ConsumerUpdateListener consumerUpdateListener;
+  private int initialCredits = 1;
+  private int additionalCredits = 1;
 
   public StreamConsumerBuilder(StreamEnvironment environment) {
     this.environment = environment;
@@ -130,6 +132,16 @@ public ConsumerBuilder noTrackingStrategy() {
     return this;
   }
 
+  @Override
+  public ConsumerBuilder credits(int initial, int onChunkDelivery) {
+    if (initial <= 0 || onChunkDelivery <= 0) {
+      throw new IllegalArgumentException("Credits must be positive");
+    }
+    this.initialCredits = initial;
+    this.additionalCredits = onChunkDelivery;
+    return this;
+  }
+
   StreamConsumerBuilder lazyInit(boolean lazyInit) {
     this.lazyInit = lazyInit;
     return this;
@@ -192,7 +204,9 @@ public Consumer build() {
               this.lazyInit,
               this.subscriptionListener,
               this.subscriptionProperties,
-              this.consumerUpdateListener);
+              this.consumerUpdateListener,
+              this.initialCredits,
+              this.additionalCredits);
       environment.addConsumer((StreamConsumer) consumer);
     } else {
       if (Utils.isSac(this.subscriptionProperties)) {

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

@@ -659,7 +659,9 @@ Runnable registerConsumer(
       SubscriptionListener subscriptionListener,
       Runnable trackingClosingCallback,
       MessageHandler messageHandler,
-      Map<String, String> subscriptionProperties) {
+      Map<String, String> subscriptionProperties,
+      int initialCredits,
+      int additionalCredits) {
     Runnable closingCallback =
         this.consumersCoordinator.subscribe(
             consumer,
@@ -669,7 +671,9 @@ Runnable registerConsumer(
             subscriptionListener,
             trackingClosingCallback,
             messageHandler,
-            subscriptionProperties);
+            subscriptionProperties,
+            initialCredits,
+            additionalCredits);
     return closingCallback;
   }
 

src/main/java/com/rabbitmq/stream/perf/StreamPerfTest.java

@@ -47,6 +47,7 @@
 import com.rabbitmq.stream.impl.Client;
 import com.rabbitmq.stream.metrics.MetricsCollector;
 import com.rabbitmq.stream.perf.ShutdownService.CloseCallback;
+import com.rabbitmq.stream.perf.Utils.CreditSettings;
 import com.rabbitmq.stream.perf.Utils.NamedThreadFactory;
 import com.rabbitmq.stream.perf.Utils.PerformanceMicrometerMetricsCollector;
 import io.micrometer.core.instrument.Counter;
@@ -425,6 +426,13 @@ public void setMaxSegmentSize(ByteCapacity in) {
       defaultValue = "false")
   private boolean metricsCommandLineArguments;
 
+  @CommandLine.Option(
+      names = {"--credits", "-cr"},
+      description = "initial and additional credits for subscriptions",
+      defaultValue = "10:1",
+      converter = Utils.CreditsTypeConverter.class)
+  private CreditSettings credits;
+
   private MetricsCollector metricsCollector;
   private PerformanceMetrics performanceMetrics;
   private List<Monitoring> monitorings;
@@ -883,8 +891,11 @@ public Integer call() throws Exception {
 
                         AtomicLong messageCount = new AtomicLong(0);
                         String stream = stream(streams, i);
-                        ConsumerBuilder consumerBuilder = environment.consumerBuilder();
-                        consumerBuilder = consumerBuilder.offset(this.offset);
+                        ConsumerBuilder consumerBuilder =
+                            environment
+                                .consumerBuilder()
+                                .offset(this.offset)
+                                .credits(this.credits.initial(), this.credits.additional());
 
                         if (this.superStreams) {
                           consumerBuilder.superStream(stream);

src/main/java/com/rabbitmq/stream/perf/Utils.java

@@ -647,6 +647,69 @@ public Integer convert(String input) {
     }
   }
 
+  static class CreditsTypeConverter implements CommandLine.ITypeConverter<CreditSettings> {
+
+    @Override
+    public CreditSettings convert(String input) {
+      String errorMessage =
+          input + " is not a valid credits setting, " + "valid example values: 20:1, 15";
+      if (input == null || input.trim().isEmpty()) {
+        typeConversionException(errorMessage);
+      }
+      if (input.contains(":") || input.contains("-")) {
+        String separator = input.contains(":") ? ":" : "-";
+        String[] split = input.split(separator);
+        if (split.length != 2) {
+          typeConversionException(errorMessage);
+        } else {
+          int[] credits =
+              Arrays.stream(split)
+                  .mapToInt(Integer::valueOf)
+                  .peek(
+                      c -> {
+                        if (c <= 0) {
+                          typeConversionException("credit values must be positive");
+                        }
+                      })
+                  .toArray();
+          return new CreditSettings(credits[0], credits[1]);
+        }
+      }
+      try {
+        int value = Integer.parseInt(input);
+        if (value <= 0) {
+          typeConversionException("credit values must be positive");
+        }
+        return new CreditSettings(value, 1);
+      } catch (Exception e) {
+        typeConversionException(errorMessage);
+      }
+      return new CreditSettings(10, 1);
+    }
+  }
+
+  static class CreditSettings {
+
+    private final int initial, additional;
+
+    CreditSettings(int initial, int additional) {
+      this.initial = initial;
+      this.additional = additional;
+    }
+
+    int initial() {
+      return this.initial;
+    }
+
+    int additional() {
+      return this.additional;
+    }
+  }
+
+  private static void typeConversionException(String message) {
+    throw new TypeConversionException(message);
+  }
+
   static class CompressionTypeConverter implements CommandLine.ITypeConverter<Compression> {
 
     @Override