From 7add6068d81abab3e40d4eba348eeb74bb0d2626 Mon Sep 17 00:00:00 2001 From: henry701 Date: Sun, 21 May 2023 14:10:01 -0300 Subject: [PATCH 01/14] Add flow control strategy mechanism Added also a single default flow control strategy, named legacy, which is set by default. This change should be transparent by all API consumers, but a new method to configure flow control strategies has been added. --- .../stream/AbstractFlowControlStrategy.java | 30 ++++ .../rabbitmq/stream/ClientDataHandler.java | 92 ++++++++++++ .../com/rabbitmq/stream/ConsumerBuilder.java | 22 +-- .../stream/ConsumerFlowControlStrategy.java | 49 +++++++ .../ConsumerFlowControlStrategyBuilder.java | 21 +++ ...umerFlowControlStrategyBuilderFactory.java | 14 ++ .../stream/LegacyFlowControlStrategy.java | 46 ++++++ ...gacyFlowControlStrategyBuilderFactory.java | 49 +++++++ .../stream/impl/ConsumersCoordinator.java | 75 +++++----- .../rabbitmq/stream/impl/StreamConsumer.java | 14 +- .../stream/impl/StreamConsumerBuilder.java | 35 +++-- .../stream/impl/StreamEnvironment.java | 43 ++---- src/test/java/com/rabbitmq/stream/Host.java | 5 +- .../stream/impl/ConsumersCoordinatorTest.java | 134 +++++++++--------- .../stream/impl/RetentionClientTest.java | 2 +- .../stream/impl/StreamEnvironmentTest.java | 3 +- 16 files changed, 463 insertions(+), 171 deletions(-) create mode 100644 src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java create mode 100644 src/main/java/com/rabbitmq/stream/ClientDataHandler.java create mode 100644 src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java create mode 100644 src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java create mode 100644 src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java create mode 100644 src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java create mode 100644 src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java diff --git a/src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java new file mode 100644 index 0000000000..08ed1980ab --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java @@ -0,0 +1,30 @@ +package com.rabbitmq.stream; + +import com.rabbitmq.stream.impl.Client; + +import java.util.Objects; +import java.util.function.Supplier; + +public abstract class AbstractFlowControlStrategy implements ConsumerFlowControlStrategy { + + private final Supplier clientSupplier; + private volatile Client client; + + protected AbstractFlowControlStrategy(Supplier clientSupplier) { + this.clientSupplier = Objects.requireNonNull(clientSupplier, "clientSupplier"); + } + + protected Client mandatoryClient() { + Client localClient = this.client; + if(localClient != null) { + return localClient; + } + localClient = clientSupplier.get(); + if(localClient == null) { + throw new IllegalStateException("Requested client, but client is not yet available! Supplier: " + this.clientSupplier); + } + this.client = localClient; + return localClient; + } + +} diff --git a/src/main/java/com/rabbitmq/stream/ClientDataHandler.java b/src/main/java/com/rabbitmq/stream/ClientDataHandler.java new file mode 100644 index 0000000000..7a116a0f68 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/ClientDataHandler.java @@ -0,0 +1,92 @@ +package com.rabbitmq.stream; + +import com.rabbitmq.stream.impl.Client; + +/** + * 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, + Client.PublishErrorListener, + Client.ChunkListener, + Client.MessageListener, + Client.CreditNotification, + Client.ConsumerUpdateListener, + Client.ShutdownListener, + Client.MetadataListener { + + @Override + default void handle(byte publisherId, long publishingId) { + this.handlePublishConfirm(publisherId, publishingId); + } + + default void handlePublishConfirm(byte publisherId, long publishingId) { + // No-op by default + } + + @Override + default void handle(byte publisherId, long publishingId, short errorCode) { + this.handlePublishError(publisherId, publishingId, errorCode); + } + + default void handlePublishError(byte publisherId, long publishingId, short errorCode) { + // No-op by default + } + + @Override + default void handle(Client client, byte subscriptionId, long offset, long messageCount, long dataSize) { + this.handleChunk(subscriptionId, offset, messageCount, dataSize); + } + + default void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + // No-op by default + } + + @Override + default void handle(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { + this.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); + } + + default void handleMessage(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { + // No-op by default + } + + @Override + default void handle(byte subscriptionId, short responseCode) { + this.handleCreditNotification(subscriptionId, responseCode); + } + + default void handleCreditNotification(byte subscriptionId, short responseCode) { + // No-op by default + } + + @Override + default OffsetSpecification handle(Client client, byte subscriptionId, boolean active) { + this.handleConsumerUpdate(subscriptionId, active); + return null; + } + + default void handleConsumerUpdate(byte subscriptionId, boolean active) { + // No-op by default + } + + @Override + default void handle(Client.ShutdownContext shutdownContext) { + this.handleShutdown(shutdownContext); + } + + default void handleShutdown(Client.ShutdownContext shutdownContext) { + // No-op by default + } + + @Override + default void handle(String stream, short code) { + this.handleMetadata(stream, code); + } + + default void handleMetadata(String stream, short code) { + // No-op by default + } + +} diff --git a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java index c4d1e3f1f7..c5f0b9c511 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java @@ -59,6 +59,14 @@ public interface ConsumerBuilder { */ ConsumerBuilder messageHandler(MessageHandler messageHandler); + /** + * Factory for the flow control strategy to be used when consuming messages. + * @param consumerFlowControlStrategyBuilderFactory the factory + * @return a fluent configurable builder for the flow control strategy + * @param + */ + > T flowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory); + /** * The logical name of the {@link Consumer}. * @@ -159,7 +167,7 @@ public interface ConsumerBuilder { Consumer build(); /** Manual tracking strategy. */ - interface ManualTrackingStrategy { + interface ManualTrackingStrategy extends ConsumerBuilderAccessor { /** * Interval to check if the last requested stored offset has been actually stored. @@ -170,17 +178,10 @@ interface ManualTrackingStrategy { * @return the manual tracking strategy */ ManualTrackingStrategy checkInterval(Duration checkInterval); - - /** - * Go back to the builder. - * - * @return the consumer builder - */ - ConsumerBuilder builder(); } /** Auto-tracking strategy. */ - interface AutoTrackingStrategy { + interface AutoTrackingStrategy extends ConsumerBuilderAccessor { /** * Number of messages before storing. @@ -201,7 +202,9 @@ interface AutoTrackingStrategy { * @return the auto-tracking strategy */ AutoTrackingStrategy flushInterval(Duration flushInterval); + } + interface ConsumerBuilderAccessor { /** * Go back to the builder. * @@ -209,4 +212,5 @@ interface AutoTrackingStrategy { */ ConsumerBuilder builder(); } + } diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..7681f1e118 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java @@ -0,0 +1,49 @@ +package com.rabbitmq.stream; + +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 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(); + } + + @Override + default void close() { + // Override with cleanup logic, if applicable + } + +} diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java new file mode 100644 index 0000000000..b9f2db30d3 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java @@ -0,0 +1,21 @@ +package com.rabbitmq.stream; + +import com.rabbitmq.stream.impl.Client; + +import java.util.function.Supplier; + +/** + * Fluent builder for a {@link ConsumerFlowControlStrategyBuilderFactory}. + * + * @param the type of {@link ConsumerFlowControlStrategy} to be built + */ +public interface ConsumerFlowControlStrategyBuilder extends ConsumerBuilder.ConsumerBuilderAccessor { + /** + * Builds the actual FlowControlStrategy instance, for the Client with which it interoperates + * + * @param clientSupplier {@link Supplier } for retrieving the {@link Client}. + * Is not a {@link Client} instance because the {@link Client} may be lazily initialized. + * @return the FlowControlStrategy + */ + T build(Supplier clientSupplier); +} diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java new file mode 100644 index 0000000000..3f52ab2d1c --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java @@ -0,0 +1,14 @@ +package com.rabbitmq.stream; + +/** + * A strategy for regulating consumer flow when consuming from a RabbitMQ Stream. + * @param the type of {@link ConsumerFlowControlStrategy} to be built + * @param the type of fluent builder exposed by this factory. Must subclass {@link ConsumerFlowControlStrategyBuilder} + */ +public interface ConsumerFlowControlStrategyBuilderFactory> { + /** + * Accessor for configuration builder with settings specific to each implementing strategy + * @return {@link C} the specific consumer flow control strategy configuration builder + */ + C builder(ConsumerBuilder consumerBuilder); +} diff --git a/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java new file mode 100644 index 0000000000..27a5966061 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java @@ -0,0 +1,46 @@ +package com.rabbitmq.stream; + +import com.rabbitmq.stream.impl.Client; + +import java.util.Map; +import java.util.function.Supplier; + +/** + * The flow control strategy that was always applied before the flow control strategy mechanism existed in the codebase. + * Requests a set amount of credits after each chunk arrives. + */ +public class LegacyFlowControlStrategy extends AbstractFlowControlStrategy { + + private final int initialCredits; + private final int additionalCredits; + + public LegacyFlowControlStrategy(Supplier clientSupplier) { + this(clientSupplier, 1); + } + + public LegacyFlowControlStrategy(Supplier clientSupplier, int initialCredits) { + this(clientSupplier, initialCredits, 1); + } + + public LegacyFlowControlStrategy(Supplier clientSupplier, int initialCredits, int additionalCredits) { + super(clientSupplier); + this.initialCredits = initialCredits; + this.additionalCredits = additionalCredits; + } + + @Override + public int handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ) { + return this.initialCredits; + } + + @Override + public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + mandatoryClient().credit(subscriptionId, this.additionalCredits); + } + +} diff --git a/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java new file mode 100644 index 0000000000..1dd4226522 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java @@ -0,0 +1,49 @@ +package com.rabbitmq.stream; + +import com.rabbitmq.stream.impl.Client; + +import java.util.function.Supplier; + +public class LegacyFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { + + public static final LegacyFlowControlStrategyBuilderFactory INSTANCE = new LegacyFlowControlStrategyBuilderFactory(); + + @Override + public LegacyFlowControlStrategyBuilder builder(ConsumerBuilder consumerBuilder) { + return new LegacyFlowControlStrategyBuilder(consumerBuilder); + } + + public static class LegacyFlowControlStrategyBuilder implements ConsumerFlowControlStrategyBuilder { + + private final ConsumerBuilder consumerBuilder; + + private int initialCredits = 1; + + private int additionalCredits = 1; + + public LegacyFlowControlStrategyBuilder(ConsumerBuilder consumerBuilder) { + this.consumerBuilder = consumerBuilder; + } + + @Override + public LegacyFlowControlStrategy build(Supplier clientSupplier) { + return new LegacyFlowControlStrategy(clientSupplier, this.initialCredits, this.additionalCredits); + } + + @Override + public ConsumerBuilder builder() { + return this.consumerBuilder; + } + + public LegacyFlowControlStrategyBuilder additionalCredits(int additionalCredits) { + this.additionalCredits = additionalCredits; + return this; + } + + public LegacyFlowControlStrategyBuilder initialCredits(int initialCredits) { + this.initialCredits = initialCredits; + return this; + } + } + +} diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 5be323db01..e35a1fedf2 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -21,16 +21,8 @@ import static com.rabbitmq.stream.impl.Utils.namedRunnable; import static com.rabbitmq.stream.impl.Utils.quote; -import com.rabbitmq.stream.BackOffDelayPolicy; -import com.rabbitmq.stream.Constants; -import com.rabbitmq.stream.Consumer; -import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.*; import com.rabbitmq.stream.MessageHandler.Context; -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.StreamDoesNotExistException; -import com.rabbitmq.stream.StreamException; -import com.rabbitmq.stream.StreamNotAvailableException; -import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.SubscriptionListener.SubscriptionContext; import com.rabbitmq.stream.impl.Client.Broker; import com.rabbitmq.stream.impl.Client.ChunkListener; @@ -121,9 +113,9 @@ Runnable subscribe( SubscriptionListener subscriptionListener, Runnable trackingClosingCallback, MessageHandler messageHandler, + ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, Map subscriptionProperties, - int initialCredits, - int additionalCredits) { + int initialCredits) { List candidates = findBrokersForStream(stream); Client.Broker newNode = pickBroker(candidates); if (newNode == null) { @@ -142,9 +134,9 @@ Runnable subscribe( subscriptionListener, trackingClosingCallback, messageHandler, + consumerFlowControlStrategyBuilder, subscriptionProperties, - initialCredits, - additionalCredits); + initialCredits); try { addToManager(newNode, subscriptionTracker, offsetSpecification, true); @@ -199,7 +191,7 @@ private void addToManager( if (pickedManager == null) { String name = keyForClientSubscription(node); LOGGER.debug("Creating subscription manager on {}", name); - pickedManager = new ClientSubscriptionsManager(node, clientParameters); + pickedManager = new ClientSubscriptionsManager(node, clientParameters, tracker.consumerFlowControlStrategyBuilder); LOGGER.debug("Created subscription manager on {}, id {}", name, pickedManager.id); } try { @@ -393,6 +385,7 @@ private static class SubscriptionTracker { private final OffsetSpecification initialOffsetSpecification; private final String offsetTrackingReference; private final MessageHandler messageHandler; + private final ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder; private final StreamConsumer consumer; private final SubscriptionListener subscriptionListener; private final Runnable trackingClosingCallback; @@ -404,7 +397,6 @@ private static class SubscriptionTracker { private volatile AtomicReference state = new AtomicReference<>(SubscriptionState.OPENING); private final int initialCredits; - private final int additionalCredits; private SubscriptionTracker( long id, @@ -415,9 +407,9 @@ private SubscriptionTracker( SubscriptionListener subscriptionListener, Runnable trackingClosingCallback, MessageHandler messageHandler, + ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, Map subscriptionProperties, - int initialCredits, - int additionalCredits) { + int initialCredits) { this.id = id; this.consumer = consumer; this.stream = stream; @@ -426,8 +418,8 @@ private SubscriptionTracker( this.subscriptionListener = subscriptionListener; this.trackingClosingCallback = trackingClosingCallback; this.messageHandler = messageHandler; + this.consumerFlowControlStrategyBuilder = consumerFlowControlStrategyBuilder; this.initialCredits = initialCredits; - this.additionalCredits = additionalCredits; if (this.offsetTrackingReference == null) { this.subscriptionProperties = subscriptionProperties; } else { @@ -551,33 +543,29 @@ private class ClientSubscriptionsManager implements Comparable> streamToStreamSubscriptions = new ConcurrentHashMap<>(); + private final ConsumerFlowControlStrategy consumerFlowControlStrategy; // trackers and tracker count must be kept in sync private volatile List subscriptionTrackers = new ArrayList<>(maxConsumersByConnection); private volatile int trackerCount = 0; private final AtomicBoolean closed = new AtomicBoolean(false); - private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientParameters) { + private ClientSubscriptionsManager( + Broker node, + Client.ClientParameters clientParameters, + ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder + ) { this.id = managerIdSequence.getAndIncrement(); this.node = node; this.name = keyForClientSubscription(node); LOGGER.debug("creating subscription manager on {}", name); IntStream.range(0, maxConsumersByConnection).forEach(i -> subscriptionTrackers.add(null)); this.trackerCount = 0; - AtomicBoolean clientInitializedInManager = new AtomicBoolean(false); - ChunkListener chunkListener = - (client, subscriptionId, offset, messageCount, dataSize) -> { - SubscriptionTracker subscriptionTracker = - subscriptionTrackers.get(subscriptionId & 0xFF); - if (subscriptionTracker != null && subscriptionTracker.consumer.isOpen()) { - client.credit(subscriptionId, subscriptionTracker.additionalCredits); - } else { - LOGGER.debug( - "Could not find stream subscription {} or subscription closing, not providing credits", - subscriptionId & 0xFF); - } - }; - + AtomicReference clientReference = new AtomicReference<>(); + ConsumerFlowControlStrategy localConsumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get); + this.consumerFlowControlStrategy = localConsumerFlowControlStrategy; + ChunkListener chunkListener = (ignoredClient, subscriptionId, offset, messageCount, dataSize) -> + localConsumerFlowControlStrategy.handleChunk(subscriptionId, offset, messageCount, dataSize); CreditNotification creditNotification = (subscriptionId, responseCode) -> { SubscriptionTracker subscriptionTracker = @@ -588,6 +576,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa subscriptionId & 0xFF, stream, Utils.formatConstant(responseCode)); + localConsumerFlowControlStrategy.handleCreditNotification(subscriptionId, responseCode); }; MessageListener messageListener = @@ -609,6 +598,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa this.id, this.name); } + localConsumerFlowControlStrategy.handleMessage(subscriptionId, offset, chunkTimestamp, committedOffset, message); }; ShutdownListener shutdownListener = shutdownContext -> { @@ -662,6 +652,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa "Consumers re-assignment after disconnection from %s", name)); } + localConsumerFlowControlStrategy.handleShutdown(shutdownContext); }; MetadataListener metadataListener = (stream, code) -> { @@ -713,6 +704,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa "Consumers re-assignment after metadata update on stream '%s'", stream)); } + localConsumerFlowControlStrategy.handleMetadata(stream, code); }; ConsumerUpdateListener consumerUpdateListener = (client, subscriptionId, active) -> { @@ -731,6 +723,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa LOGGER.debug( "Could not find stream subscription {} for consumer update", subscriptionId); } + localConsumerFlowControlStrategy.handleConsumerUpdate(subscriptionId, active); return result; }; String connectionName = connectionNamingStrategy.apply(ClientConnectionType.CONSUMER); @@ -745,9 +738,10 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa .metadataListener(metadataListener) .consumerUpdateListener(consumerUpdateListener)) .key(name); - this.client = clientFactory.client(clientFactoryContext); + Client localClient = clientFactory.client(clientFactoryContext); + this.client = localClient; + clientReference.set(localClient); LOGGER.debug("Created consumer connection '{}'", connectionName); - clientInitializedInManager.set(true); } private void assignConsumersToStream( @@ -962,6 +956,12 @@ synchronized void add( checkNotClosed(); byte subId = subscriptionId; + int initialCredits = this.consumerFlowControlStrategy.handleSubscribe( + subId, + subscriptionTracker.stream, + subscriptionContext.offsetSpecification(), + subscriptionTracker.subscriptionProperties + ); Client.Response subscribeResponse = Utils.callAndMaybeRetry( () -> @@ -969,7 +969,7 @@ synchronized void add( subId, subscriptionTracker.stream, subscriptionContext.offsetSpecification(), - subscriptionTracker.initialCredits, + initialCredits, subscriptionTracker.subscriptionProperties), RETRY_ON_TIMEOUT, "Subscribe request for consumer %s on stream '%s'", @@ -1033,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, @@ -1098,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( diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java index 2e63b11e00..c7602815f0 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java @@ -18,15 +18,8 @@ import static com.rabbitmq.stream.impl.Utils.offsetBefore; import static java.time.Duration.ofMillis; -import com.rabbitmq.stream.Constants; -import com.rabbitmq.stream.Consumer; -import com.rabbitmq.stream.ConsumerUpdateListener; -import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.*; import com.rabbitmq.stream.MessageHandler.Context; -import com.rabbitmq.stream.NoOffsetException; -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.StreamException; -import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; import com.rabbitmq.stream.impl.StreamEnvironment.TrackingConsumerRegistration; @@ -74,6 +67,7 @@ class StreamConsumer implements Consumer { String stream, OffsetSpecification offsetSpecification, MessageHandler messageHandler, + ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, String name, StreamEnvironment environment, TrackingConfiguration trackingConfiguration, @@ -255,9 +249,9 @@ class StreamConsumer implements Consumer { subscriptionListener, trackingClosingCallback, closedAwareMessageHandler, + consumerFlowControlStrategyBuilder, Collections.unmodifiableMap(subscriptionProperties), - initialCredits, - additionalCredits); + initialCredits); this.status = Status.RUNNING; }; diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java index 7dd98a05fe..a32e1c343f 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java @@ -13,13 +13,8 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.Consumer; -import com.rabbitmq.stream.ConsumerBuilder; -import com.rabbitmq.stream.ConsumerUpdateListener; -import com.rabbitmq.stream.MessageHandler; -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.StreamException; -import com.rabbitmq.stream.SubscriptionListener; +import com.rabbitmq.stream.*; + import java.lang.reflect.Field; import java.lang.reflect.Modifier; import java.time.Duration; @@ -44,6 +39,7 @@ class StreamConsumerBuilder implements ConsumerBuilder { private ConsumerUpdateListener consumerUpdateListener; private int initialCredits = 1; private int additionalCredits = 1; + private ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder = LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(this); public StreamConsumerBuilder(StreamEnvironment environment) { this.environment = environment; @@ -73,6 +69,13 @@ public ConsumerBuilder messageHandler(MessageHandler messageHandler) { return this; } + @Override + public > T flowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory) { + T localConsumerFlowControlStrategyBuilder = consumerFlowControlStrategyBuilderFactory.builder(this); + this.consumerFlowControlStrategyBuilder = localConsumerFlowControlStrategyBuilder; + return localConsumerFlowControlStrategyBuilder; + } + MessageHandler messageHandler() { return this.messageHandler; } @@ -132,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; } @@ -197,6 +211,7 @@ public Consumer build() { this.stream, this.offsetSpecification, this.messageHandler, + this.consumerFlowControlStrategyBuilder, this.name, this.environment, trackingConfiguration, diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java index 7fcf5b8bda..50fec81d3e 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java @@ -19,21 +19,8 @@ import static com.rabbitmq.stream.impl.Utils.namedRunnable; import static java.util.concurrent.TimeUnit.SECONDS; -import com.rabbitmq.stream.Address; -import com.rabbitmq.stream.AddressResolver; -import com.rabbitmq.stream.BackOffDelayPolicy; -import com.rabbitmq.stream.Codec; -import com.rabbitmq.stream.ConsumerBuilder; -import com.rabbitmq.stream.Environment; -import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.*; import com.rabbitmq.stream.MessageHandler.Context; -import com.rabbitmq.stream.NoOffsetException; -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.ProducerBuilder; -import com.rabbitmq.stream.StreamCreator; -import com.rabbitmq.stream.StreamException; -import com.rabbitmq.stream.StreamStats; -import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.compression.CompressionCodecFactory; import com.rabbitmq.stream.impl.Client.ClientParameters; import com.rabbitmq.stream.impl.Client.ShutdownListener; @@ -658,22 +645,20 @@ Runnable registerConsumer( SubscriptionListener subscriptionListener, Runnable trackingClosingCallback, MessageHandler messageHandler, + ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, Map subscriptionProperties, - int initialCredits, - int additionalCredits) { - Runnable closingCallback = - this.consumersCoordinator.subscribe( - consumer, - stream, - offsetSpecification, - trackingReference, - subscriptionListener, - trackingClosingCallback, - messageHandler, - subscriptionProperties, - initialCredits, - additionalCredits); - return closingCallback; + int initialCredits) { + return this.consumersCoordinator.subscribe( + consumer, + stream, + offsetSpecification, + trackingReference, + subscriptionListener, + trackingClosingCallback, + messageHandler, + consumerFlowControlStrategyBuilder, + subscriptionProperties, + initialCredits); } Runnable registerProducer(StreamProducer producer, String reference, String stream) { diff --git a/src/test/java/com/rabbitmq/stream/Host.java b/src/test/java/com/rabbitmq/stream/Host.java index c6f68fcbd1..00b2ec536e 100644 --- a/src/test/java/com/rabbitmq/stream/Host.java +++ b/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"; diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index a7634e13ea..f9485c7cba 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -33,12 +33,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import com.rabbitmq.stream.BackOffDelayPolicy; -import com.rabbitmq.stream.Constants; -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.StreamDoesNotExistException; -import com.rabbitmq.stream.StreamException; -import com.rabbitmq.stream.SubscriptionListener; +import com.rabbitmq.stream.*; import com.rabbitmq.stream.codec.WrapperMessageBuilder; import com.rabbitmq.stream.impl.Client.MessageListener; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; @@ -76,8 +71,7 @@ public class ConsumersCoordinatorTest { private static final SubscriptionListener NO_OP_SUBSCRIPTION_LISTENER = subscriptionContext -> {}; private static final Runnable NO_OP_TRACKING_CLOSING_CALLBACK = () -> {}; - private int initialCredits = 10; - private int additionalCredits = 1; + private final int initialCredits = 10; @Mock StreamEnvironment environment; @Mock StreamConsumer consumer; @@ -205,9 +199,9 @@ void tearDown() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(2)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -247,9 +241,9 @@ void shouldGetExactNodeImmediatelyWithAdvertisedHostNameClientFactoryAndExactNod NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -278,9 +272,9 @@ void shouldSubscribeWithEmptyPropertiesWithUnamedConsumer() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -300,9 +294,9 @@ void subscribeShouldThrowExceptionWhenNoMetadataForTheStream() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits)) + initialCredits)) .isInstanceOf(StreamDoesNotExistException.class); } @@ -320,9 +314,9 @@ void subscribeShouldThrowExceptionWhenStreamDoesNotExist() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits)) + initialCredits)) .isInstanceOf(StreamDoesNotExistException.class); } @@ -350,9 +344,9 @@ void subscribePropagateExceptionWhenClientSubscriptionFails() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits)) + initialCredits)) .isInstanceOf(StreamException.class) .hasMessage(exceptionMessage); assertThat(MonitoringTestUtils.extract(coordinator).isEmpty()).isTrue(); @@ -372,9 +366,9 @@ void subscribeShouldThrowExceptionWhenMetadataResponseIsNotOk() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits)) + initialCredits)) .isInstanceOf(IllegalStateException.class); } @@ -391,9 +385,9 @@ void subscribeShouldThrowExceptionIfNoNodeAvailableForStream() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits)) + initialCredits)) .isInstanceOf(IllegalStateException.class); } @@ -433,9 +427,9 @@ void subscribeShouldSubscribeToStreamAndDispatchMessage_UnsubscribeShouldUnsubsc NO_OP_SUBSCRIPTION_LISTENER, () -> trackingClosingCallbackCalls.incrementAndGet(), (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -484,9 +478,9 @@ void subscribeShouldSubscribeToStreamAndDispatchMessageWithManySubscriptions() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.compute(subId, (k, v) -> (v == null) ? 1 : ++v), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); closingRunnables.add(closingRunnable); } @@ -560,9 +554,9 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -580,9 +574,9 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(client, times(1 + 1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -661,9 +655,9 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -717,9 +711,9 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -732,9 +726,9 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(client, times(1 + 1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -821,9 +815,9 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -887,9 +881,9 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -941,9 +935,9 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -996,9 +990,9 @@ void shouldUseNewClientsForMoreThanMaxSubscriptionsAndCloseClientAfterUnsubscrip NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits)) + initialCredits)) .collect(Collectors.toList()); verify(clientFactory, times(2)).client(any()); @@ -1058,9 +1052,9 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); }); // the extra is allocated on another client from the same pool verify(clientFactory, times(2)).client(any()); @@ -1084,9 +1078,9 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(2 + 1)).client(any()); verify(client, times(subscriptionCount + ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT + 1)) @@ -1125,9 +1119,9 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); }); // the extra is allocated on another client from the same pool verify(clientFactory, times(2)).client(any()); @@ -1157,9 +1151,9 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); // no more client creation verify(clientFactory, times(2)).client(any()); @@ -1208,9 +1202,9 @@ void shouldRestartWhereItLeftOffAfterDisruption(Consumer {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1281,9 +1275,9 @@ void shouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesReceived( NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1356,9 +1350,9 @@ void shouldUseStoredOffsetOnRecovery(Consumer configur NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1443,9 +1437,9 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1458,9 +1452,9 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1 + 1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1524,9 +1518,9 @@ void shouldRetryAssignmentOnRecoveryStreamNotAvailableFailure() throws Exception NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1590,9 +1584,9 @@ void shouldRetryAssignmentOnRecoveryCandidateLookupFailure() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1633,9 +1627,9 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, + LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), - initialCredits, - additionalCredits); + initialCredits); closingRunnable.run(); }; diff --git a/src/test/java/com/rabbitmq/stream/impl/RetentionClientTest.java b/src/test/java/com/rabbitmq/stream/impl/RetentionClientTest.java index 252d0b4522..abf5c6fb63 100644 --- a/src/test/java/com/rabbitmq/stream/impl/RetentionClientTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/RetentionClientTest.java @@ -80,7 +80,7 @@ static RetentionTestConfig[] retention() { // small retention in policy String policyCommand = String.format( - "set_policy stream-retention-test \"%s\" " + "set_policy stream-retention-test '%s' " + "'{\"max-length-bytes\":%d,\"stream-max-segment-size-bytes\":%d }' " + "--priority 1 --apply-to queues", stream, maxLengthBytes, maxSegmentSizeBytes); diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamEnvironmentTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamEnvironmentTest.java index 2a9c275824..cd987e6dc9 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StreamEnvironmentTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StreamEnvironmentTest.java @@ -148,8 +148,7 @@ void environmentCreationShouldFailWithUrlUsingWrongPort() { .build() .close()) .isInstanceOf(StreamException.class) - .hasCauseInstanceOf(ConnectException.class) - .hasRootCauseMessage("Connection refused"); + .hasCauseInstanceOf(ConnectException.class); } @Test From 2fbd85fa8873a8b1a15ca8b86f79e29da5c56611 Mon Sep 17 00:00:00 2001 From: henry701 Date: Sat, 27 May 2023 22:48:49 -0300 Subject: [PATCH 02/14] Add test for flow control strategy main handler methods --- .../stream/impl/ConsumersCoordinatorTest.java | 129 ++++++++++++++++-- 1 file changed, 119 insertions(+), 10 deletions(-) diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index f9485c7cba..85160194e2 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -22,12 +22,7 @@ import static java.lang.String.format; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyByte; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyMap; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.ArgumentMatchers.*; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -54,6 +49,7 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -62,10 +58,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; -import org.mockito.ArgumentCaptor; -import org.mockito.Captor; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; +import org.mockito.*; public class ConsumersCoordinatorTest { @@ -1657,6 +1650,122 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() { } } + @ParameterizedTest + @MethodSource("disruptionArguments") + @SuppressWarnings("unchecked") + void shouldCallConsumerFlowControlHandlers(Consumer configurator) + throws Exception { + + scheduledExecutorService = createScheduledExecutorService(); + when(environment.scheduledExecutorService()).thenReturn(scheduledExecutorService); + Duration retryDelay = Duration.ofMillis(100); + when(environment.recoveryBackOffDelayPolicy()).thenReturn(BackOffDelayPolicy.fixed(retryDelay)); + when(environment.topologyUpdateBackOffDelayPolicy()) + .thenReturn(BackOffDelayPolicy.fixed(retryDelay)); + when(consumer.isOpen()).thenReturn(true); + when(locator.metadata("stream")) + .thenReturn(metadata(null, replicas())) + .thenReturn(metadata(null, Collections.emptyList())) + .thenReturn(metadata(null, replicas())); + + when(clientFactory.client(any())).thenReturn(client); + + String consumerName = "consumer-name"; + long lastStoredOffset = 5; + long lastReceivedOffset = 10; + when(client.queryOffset(consumerName, "stream")) + .thenReturn(new QueryOffsetResponse(Constants.RESPONSE_CODE_OK, 0L)) + .thenReturn(new QueryOffsetResponse(Constants.RESPONSE_CODE_OK, lastStoredOffset)); + + ArgumentCaptor offsetSpecificationArgumentCaptor = + ArgumentCaptor.forClass(OffsetSpecification.class); + ArgumentCaptor> subscriptionPropertiesArgumentCaptor = + ArgumentCaptor.forClass(Map.class); + when(client.subscribe( + subscriptionIdCaptor.capture(), + anyString(), + offsetSpecificationArgumentCaptor.capture(), + anyInt(), + subscriptionPropertiesArgumentCaptor.capture())) + .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); + + ConsumerFlowControlStrategy mockedConsumerFlowControlStrategy = Mockito.mock(ConsumerFlowControlStrategy.class); + + int numberOfInitialCreditsOnSubscribe = 7; + + when(mockedConsumerFlowControlStrategy.handleSubscribe(anyByte(), anyString(), any(), anyMap())).thenReturn(numberOfInitialCreditsOnSubscribe); + + ConsumerFlowControlStrategyBuilder mockedConsumerFlowControlStrategyBuilder = Mockito.mock(ConsumerFlowControlStrategyBuilder.class); + when(mockedConsumerFlowControlStrategyBuilder.build(any())).thenReturn(mockedConsumerFlowControlStrategy); + + Runnable closingRunnable = + coordinator.subscribe( + consumer, + "stream", + null, + consumerName, + NO_OP_SUBSCRIPTION_LISTENER, + NO_OP_TRACKING_CLOSING_CALLBACK, + (offset, message) -> {}, + mockedConsumerFlowControlStrategyBuilder, + Collections.emptyMap(), + initialCredits); + verify(clientFactory, times(1)).client(any()); + verify(client, times(1)) + .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), eq(numberOfInitialCreditsOnSubscribe), anyMap()); + verify(mockedConsumerFlowControlStrategy, times(1)) + .handleSubscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); + assertThat(offsetSpecificationArgumentCaptor.getAllValues()) + .element(0) + .isEqualTo(OffsetSpecification.next()); + assertThat(subscriptionPropertiesArgumentCaptor.getAllValues()) + .element(0) + .isEqualTo(Collections.singletonMap("name", "consumer-name")); + + Message message = new WrapperMessageBuilder().build(); + + messageListener.handle( + subscriptionIdCaptor.getValue(), + lastReceivedOffset, + 0, + 0, + message); + + verify(mockedConsumerFlowControlStrategy).handleMessage( + subscriptionIdCaptor.getValue(), + lastReceivedOffset, + 0, + 0, + message + ); + + configurator.accept(this); + + Thread.sleep(retryDelay.toMillis() * 5); + + verify(client, times(2)) + .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + + verify(mockedConsumerFlowControlStrategy, times(2)) + .handleSubscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); + + assertThat(offsetSpecificationArgumentCaptor.getAllValues()) + .element(1) + .isEqualTo(OffsetSpecification.offset(lastStoredOffset + 1)) + .isNotEqualTo(OffsetSpecification.offset(lastReceivedOffset)); + assertThat(subscriptionPropertiesArgumentCaptor.getAllValues()) + .element(1) + .isEqualTo(Collections.singletonMap("name", "consumer-name")); + when(client.unsubscribe(subscriptionIdCaptor.getValue())) + .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); + + closingRunnable.run(); + + verify(client, times(1)).unsubscribe(subscriptionIdCaptor.getValue()); + verify(mockedConsumerFlowControlStrategy, times(1)) + .handleUnsubscribe(subscriptionIdCaptor.getValue()); + } + Client.Broker leader() { return new Client.Broker("leader", -1); } From 7175c3d30166fc55e6bc274cf1f4f4d6a1039f44 Mon Sep 17 00:00:00 2001 From: henry701 Date: Sat, 10 Jun 2023 00:28:46 -0300 Subject: [PATCH 03/14] Abstract Client instance away from flow control strategy, create ConsumerStatisticRecorder and a flow control strategy that uses it --- .../stream/AbstractFlowControlStrategy.java | 30 -- .../com/rabbitmq/stream/ConsumerBuilder.java | 3 + .../stream/ConsumerFlowControlStrategy.java | 49 --- ...gacyFlowControlStrategyBuilderFactory.java | 49 --- .../AbstractConsumerFlowControlStrategy.java | 35 ++ .../flow/ConsumerFlowControlStrategy.java | 46 +++ .../ConsumerFlowControlStrategyBuilder.java | 11 +- ...umerFlowControlStrategyBuilderFactory.java | 4 +- .../com/rabbitmq/stream/flow/CreditAsker.java | 13 + .../stream/flow/MessageHandlingAware.java | 16 + .../java/com/rabbitmq/stream/impl/Client.java | 109 ++---- .../stream/{ => impl}/ClientDataHandler.java | 32 +- .../impl/ConsumerStatisticRecorder.java | 342 ++++++++++++++++++ .../stream/impl/ConsumersCoordinator.java | 47 +-- .../rabbitmq/stream/impl/StreamConsumer.java | 16 +- .../stream/impl/StreamConsumerBuilder.java | 7 +- .../stream/impl/StreamEnvironment.java | 32 +- ...cRecordingConsumerFlowControlStrategy.java | 132 +++++++ .../LegacyConsumerFlowControlStrategy.java} | 22 +- ...umerFlowControlStrategyBuilderFactory.java | 52 +++ ...bscriptionConsumerFlowControlStrategy.java | 82 +++++ ...umerFlowControlStrategyBuilderFactory.java | 46 +++ .../stream/impl/ConsumersCoordinatorTest.java | 118 +++--- 23 files changed, 924 insertions(+), 369 deletions(-) delete mode 100644 src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java delete mode 100644 src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java delete mode 100644 src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java create mode 100644 src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java create mode 100644 src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java rename src/main/java/com/rabbitmq/stream/{ => flow}/ConsumerFlowControlStrategyBuilder.java (56%) rename src/main/java/com/rabbitmq/stream/{ => flow}/ConsumerFlowControlStrategyBuilderFactory.java (89%) create mode 100644 src/main/java/com/rabbitmq/stream/flow/CreditAsker.java create mode 100644 src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java rename src/main/java/com/rabbitmq/stream/{ => impl}/ClientDataHandler.java (73%) create mode 100644 src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java create mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java rename src/main/java/com/rabbitmq/stream/{LegacyFlowControlStrategy.java => impl/flow/LegacyConsumerFlowControlStrategy.java} (59%) create mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java create mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java create mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java diff --git a/src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java deleted file mode 100644 index 08ed1980ab..0000000000 --- a/src/main/java/com/rabbitmq/stream/AbstractFlowControlStrategy.java +++ /dev/null @@ -1,30 +0,0 @@ -package com.rabbitmq.stream; - -import com.rabbitmq.stream.impl.Client; - -import java.util.Objects; -import java.util.function.Supplier; - -public abstract class AbstractFlowControlStrategy implements ConsumerFlowControlStrategy { - - private final Supplier clientSupplier; - private volatile Client client; - - protected AbstractFlowControlStrategy(Supplier clientSupplier) { - this.clientSupplier = Objects.requireNonNull(clientSupplier, "clientSupplier"); - } - - protected Client mandatoryClient() { - Client localClient = this.client; - if(localClient != null) { - return localClient; - } - localClient = clientSupplier.get(); - if(localClient == null) { - throw new IllegalStateException("Requested client, but client is not yet available! Supplier: " + this.clientSupplier); - } - this.client = localClient; - return localClient; - } - -} diff --git a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java index c5f0b9c511..59201c12d3 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java @@ -13,6 +13,9 @@ // info@rabbitmq.com. package com.rabbitmq.stream; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; + import java.time.Duration; /** API to configure and create a {@link Consumer}. */ diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java deleted file mode 100644 index 7681f1e118..0000000000 --- a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategy.java +++ /dev/null @@ -1,49 +0,0 @@ -package com.rabbitmq.stream; - -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 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(); - } - - @Override - default void close() { - // Override with cleanup logic, if applicable - } - -} diff --git a/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java deleted file mode 100644 index 1dd4226522..0000000000 --- a/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategyBuilderFactory.java +++ /dev/null @@ -1,49 +0,0 @@ -package com.rabbitmq.stream; - -import com.rabbitmq.stream.impl.Client; - -import java.util.function.Supplier; - -public class LegacyFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { - - public static final LegacyFlowControlStrategyBuilderFactory INSTANCE = new LegacyFlowControlStrategyBuilderFactory(); - - @Override - public LegacyFlowControlStrategyBuilder builder(ConsumerBuilder consumerBuilder) { - return new LegacyFlowControlStrategyBuilder(consumerBuilder); - } - - public static class LegacyFlowControlStrategyBuilder implements ConsumerFlowControlStrategyBuilder { - - private final ConsumerBuilder consumerBuilder; - - private int initialCredits = 1; - - private int additionalCredits = 1; - - public LegacyFlowControlStrategyBuilder(ConsumerBuilder consumerBuilder) { - this.consumerBuilder = consumerBuilder; - } - - @Override - public LegacyFlowControlStrategy build(Supplier clientSupplier) { - return new LegacyFlowControlStrategy(clientSupplier, this.initialCredits, this.additionalCredits); - } - - @Override - public ConsumerBuilder builder() { - return this.consumerBuilder; - } - - public LegacyFlowControlStrategyBuilder additionalCredits(int additionalCredits) { - this.additionalCredits = additionalCredits; - return this; - } - - public LegacyFlowControlStrategyBuilder initialCredits(int initialCredits) { - this.initialCredits = initialCredits; - return this; - } - } - -} diff --git a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..fd34cf784d --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java @@ -0,0 +1,35 @@ +package com.rabbitmq.stream.flow; + +import com.rabbitmq.stream.impl.Client; + +import java.util.Objects; +import java.util.function.Supplier; + +/** + * Abstract class for Consumer Flow Control Strategies which keeps a cached lazily-initialized + * {@link Client} ready for retrieval by its inheritors. + */ +public abstract class AbstractConsumerFlowControlStrategy implements ConsumerFlowControlStrategy { + + private final Supplier creditAskerSupplier; + private volatile CreditAsker creditAsker; + + + protected AbstractConsumerFlowControlStrategy(Supplier creditAskerSupplier) { + this.creditAskerSupplier = Objects.requireNonNull(creditAskerSupplier, "creditAskerSupplier"); + } + + protected CreditAsker mandatoryClient() { + CreditAsker localSupplied = this.creditAsker; + if(localSupplied != null) { + return localSupplied; + } + localSupplied = creditAskerSupplier.get(); + if(localSupplied == null) { + throw new IllegalStateException("Requested client, but client is not yet available! Supplier: " + this.creditAskerSupplier); + } + this.creditAsker = localSupplied; + return localSupplied; + } + +} diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..3ce242ef00 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java @@ -0,0 +1,46 @@ +package com.rabbitmq.stream.flow; + +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.impl.ClientDataHandler; + +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 CreditAsker#credit} to ask for credits. + */ +// TODO: Decouple from ClientDataHandler. Maybe create an adapter pattern for handling this, or something. +public interface ConsumerFlowControlStrategy extends ClientDataHandler { + + /** + * Callback for handling a new stream subscription. + * Called right before the subscription is sent to the actual client. + *

+ * Either this variant or {@link ClientDataHandler#handleSubscribe(byte, String, OffsetSpecification, Map)} should be called, NOT both. + *

+ * + * @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 handleSubscribeReturningInitialCredits( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ); + + @Override + default void handleSubscribe(byte subscriptionId, String stream, OffsetSpecification offsetSpecification, Map subscriptionProperties) { + handleSubscribeReturningInitialCredits( + subscriptionId, + stream, + offsetSpecification, + subscriptionProperties + ); + } + +} diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java similarity index 56% rename from src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java rename to src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java index b9f2db30d3..d732451627 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilder.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java @@ -1,6 +1,6 @@ -package com.rabbitmq.stream; +package com.rabbitmq.stream.flow; -import com.rabbitmq.stream.impl.Client; +import com.rabbitmq.stream.ConsumerBuilder; import java.util.function.Supplier; @@ -13,9 +13,8 @@ public interface ConsumerFlowControlStrategyBuilder} for retrieving the {@link Client}. - * Is not a {@link Client} instance because the {@link Client} may be lazily initialized. - * @return the FlowControlStrategy + * @param creditAskerSupplier {@link Supplier} for retrieving the instance (which may be lazily initialized). + * @return {@link T} the built {@link ConsumerFlowControlStrategy} */ - T build(Supplier clientSupplier); + T build(Supplier creditAskerSupplier); } diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java similarity index 89% rename from src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java rename to src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java index 3f52ab2d1c..766257dafd 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerFlowControlStrategyBuilderFactory.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java @@ -1,4 +1,6 @@ -package com.rabbitmq.stream; +package com.rabbitmq.stream.flow; + +import com.rabbitmq.stream.ConsumerBuilder; /** * A strategy for regulating consumer flow when consuming from a RabbitMQ Stream. diff --git a/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java b/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java new file mode 100644 index 0000000000..d367451173 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java @@ -0,0 +1,13 @@ +package com.rabbitmq.stream.flow; + +public interface CreditAsker { + + /** + * Asks for credits for a given subscription. + * @param subscriptionId the subscription ID + * @param credit how many credits to ask for + * @throws IllegalArgumentException if credits are below 0 or above {@link Short#MAX_VALUE} + */ + void credit(byte subscriptionId, int credit); + +} diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java new file mode 100644 index 0000000000..f21d370f26 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java @@ -0,0 +1,16 @@ +package com.rabbitmq.stream.flow; + +import com.rabbitmq.stream.MessageHandler; + +public interface MessageHandlingAware { + + /** + * Marks a message as handled + * + * @param messageContext The {@link MessageHandler.Context} of the handled message + * @return Whether the message was marked as handled (returning {@code true}) + * or was not found (either because it was already marked as handled, or wasn't tracked) + */ + boolean markHandled(MessageHandler.Context messageContext); + +} diff --git a/src/main/java/com/rabbitmq/stream/impl/Client.java b/src/main/java/com/rabbitmq/stream/impl/Client.java index d9061799eb..ffc91ed963 100644 --- a/src/main/java/com/rabbitmq/stream/impl/Client.java +++ b/src/main/java/com/rabbitmq/stream/impl/Client.java @@ -13,85 +13,25 @@ // info@rabbitmq.com. 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; -import static com.rabbitmq.stream.Constants.COMMAND_DELETE_PUBLISHER; -import static com.rabbitmq.stream.Constants.COMMAND_DELETE_STREAM; -import static com.rabbitmq.stream.Constants.COMMAND_EXCHANGE_COMMAND_VERSIONS; -import static com.rabbitmq.stream.Constants.COMMAND_HEARTBEAT; -import static com.rabbitmq.stream.Constants.COMMAND_METADATA; -import static com.rabbitmq.stream.Constants.COMMAND_OPEN; -import static com.rabbitmq.stream.Constants.COMMAND_PARTITIONS; -import static com.rabbitmq.stream.Constants.COMMAND_PEER_PROPERTIES; -import static com.rabbitmq.stream.Constants.COMMAND_PUBLISH; -import static com.rabbitmq.stream.Constants.COMMAND_QUERY_OFFSET; -import static com.rabbitmq.stream.Constants.COMMAND_QUERY_PUBLISHER_SEQUENCE; -import static com.rabbitmq.stream.Constants.COMMAND_ROUTE; -import static com.rabbitmq.stream.Constants.COMMAND_SASL_AUTHENTICATE; -import static com.rabbitmq.stream.Constants.COMMAND_SASL_HANDSHAKE; -import static com.rabbitmq.stream.Constants.COMMAND_STORE_OFFSET; -import static com.rabbitmq.stream.Constants.COMMAND_STREAM_STATS; -import static com.rabbitmq.stream.Constants.COMMAND_SUBSCRIBE; -import static com.rabbitmq.stream.Constants.COMMAND_UNSUBSCRIBE; -import static com.rabbitmq.stream.Constants.RESPONSE_CODE_AUTHENTICATION_FAILURE; -import static com.rabbitmq.stream.Constants.RESPONSE_CODE_AUTHENTICATION_FAILURE_LOOPBACK; -import static com.rabbitmq.stream.Constants.RESPONSE_CODE_OK; -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 com.rabbitmq.stream.impl.Utils.noOpConsumer; -import static java.lang.String.format; -import static java.lang.String.join; -import static java.util.concurrent.TimeUnit.SECONDS; - -import com.rabbitmq.stream.AuthenticationFailureException; -import com.rabbitmq.stream.ByteCapacity; -import com.rabbitmq.stream.ChunkChecksum; -import com.rabbitmq.stream.Codec; +import com.rabbitmq.stream.*; import com.rabbitmq.stream.Codec.EncodedMessage; -import com.rabbitmq.stream.Constants; -import com.rabbitmq.stream.Environment; -import com.rabbitmq.stream.Message; -import com.rabbitmq.stream.MessageBuilder; -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.Producer; import com.rabbitmq.stream.StreamCreator.LeaderLocator; -import com.rabbitmq.stream.StreamException; import com.rabbitmq.stream.compression.Compression; import com.rabbitmq.stream.compression.CompressionCodec; import com.rabbitmq.stream.compression.CompressionCodecFactory; +import com.rabbitmq.stream.flow.CreditAsker; import com.rabbitmq.stream.impl.Client.ShutdownContext.ShutdownReason; import com.rabbitmq.stream.impl.ServerFrameHandler.FrameHandler; import com.rabbitmq.stream.impl.ServerFrameHandler.FrameHandlerInfo; -import com.rabbitmq.stream.impl.Utils.NamedThreadFactory; +import com.rabbitmq.stream.impl.Utils.*; import com.rabbitmq.stream.metrics.MetricsCollector; import com.rabbitmq.stream.metrics.NoOpMetricsCollector; -import com.rabbitmq.stream.sasl.CredentialsProvider; -import com.rabbitmq.stream.sasl.DefaultSaslConfiguration; -import com.rabbitmq.stream.sasl.DefaultUsernamePasswordCredentialsProvider; -import com.rabbitmq.stream.sasl.SaslConfiguration; -import com.rabbitmq.stream.sasl.SaslMechanism; -import com.rabbitmq.stream.sasl.UsernamePasswordCredentialsProvider; +import com.rabbitmq.stream.sasl.*; import io.netty.bootstrap.Bootstrap; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufAllocator; import io.netty.buffer.ByteBufOutputStream; -import io.netty.channel.Channel; -import io.netty.channel.ChannelFuture; -import io.netty.channel.ChannelHandlerContext; -import io.netty.channel.ChannelInboundHandlerAdapter; -import io.netty.channel.ChannelInitializer; -import io.netty.channel.ChannelOption; -import io.netty.channel.ChannelOutboundHandlerAdapter; -import io.netty.channel.ChannelPromise; -import io.netty.channel.ConnectTimeoutException; -import io.netty.channel.EventLoopGroup; +import io.netty.channel.*; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioSocketChannel; @@ -103,6 +43,12 @@ import io.netty.handler.timeout.IdleState; import io.netty.handler.timeout.IdleStateEvent; import io.netty.handler.timeout.IdleStateHandler; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.net.ssl.SSLEngine; +import javax.net.ssl.SSLHandshakeException; +import javax.net.ssl.SSLParameters; import java.io.IOException; import java.io.OutputStream; import java.lang.reflect.Field; @@ -111,34 +57,20 @@ import java.net.SocketAddress; import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; +import java.util.*; +import java.util.concurrent.*; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.ToLongFunction; -import javax.net.ssl.SSLEngine; -import javax.net.ssl.SSLHandshakeException; -import javax.net.ssl.SSLParameters; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; + +import static com.rabbitmq.stream.Constants.*; +import static com.rabbitmq.stream.impl.Utils.*; +import static java.lang.String.format; +import static java.lang.String.join; +import static java.util.concurrent.TimeUnit.SECONDS; /** * This is low-level client API to communicate with the broker. @@ -151,7 +83,7 @@ *

People wanting very fine control over their interaction with the broker can use {@link Client} * but at their own risk. */ -public class Client implements AutoCloseable { +public class Client implements CreditAsker, AutoCloseable { public static final int DEFAULT_PORT = 5552; public static final int DEFAULT_TLS_PORT = 5551; @@ -1072,6 +1004,7 @@ public MessageBuilder messageBuilder() { return this.codec.messageBuilder(); } + @Override public void credit(byte subscriptionId, int credit) { if (credit < 0 || credit > Short.MAX_VALUE) { throw new IllegalArgumentException("Credit value must be between 0 and " + Short.MAX_VALUE); diff --git a/src/main/java/com/rabbitmq/stream/ClientDataHandler.java b/src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java similarity index 73% rename from src/main/java/com/rabbitmq/stream/ClientDataHandler.java rename to src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java index 7a116a0f68..c91cbdbc39 100644 --- a/src/main/java/com/rabbitmq/stream/ClientDataHandler.java +++ b/src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java @@ -1,6 +1,9 @@ -package com.rabbitmq.stream; +package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.impl.Client; +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.OffsetSpecification; + +import java.util.Map; /** * Exposes callbacks to handle events from a particular {@link Client}, @@ -89,4 +92,29 @@ default void handleMetadata(String stream, short code) { // No-op by default } + /** + * Callback for handling a new stream subscription. + * + * @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 + */ + default void handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ) { + // No-op by default + } + + /** + * 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 + } + } diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java new file mode 100644 index 0000000000..7897caa308 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java @@ -0,0 +1,342 @@ +package com.rabbitmq.stream.impl; + +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.OffsetSpecification; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.Collections; +import java.util.Map; +import java.util.NavigableMap; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.atomic.AtomicInteger; + +public class ConsumerStatisticRecorder implements ClientDataHandler { + + private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerStatisticRecorder.class); + + private final ConcurrentMap> streamNameToSubscriptionIdMap = new ConcurrentHashMap<>(); + private final ConcurrentMap subscriptionStatisticsMap = new ConcurrentHashMap<>(); + + @Override + public void handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ) { + this.streamNameToSubscriptionIdMap.compute( + stream, + (k, v) -> { + if(v == null) { + v = Collections.newSetFromMap(new ConcurrentHashMap<>()); + } + boolean isNewElement = v.add(subscriptionId); + if(!isNewElement) { + LOGGER.warn( + "handleSubscribed called for stream that already had same associated subscription! " + + "subscriptionId={}, stream={}", + subscriptionId, + stream + ); + } + return v; + } + ); + this.subscriptionStatisticsMap.compute( + subscriptionId, + (k, v) -> { + if(v != null) { + LOGGER.warn( + "handleSubscribed called for subscription that already exists! subscriptionId={}", + subscriptionId + ); + } + return new SubscriptionStatistics(subscriptionId, stream, subscriptionProperties); + } + ); + } + + @Override + public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + this.subscriptionStatisticsMap.compute( + subscriptionId, + (k, v) -> { + if(v == null) { + LOGGER.warn( + "handleChunk called for subscription that does not exist! subscriptionId={}", + subscriptionId + ); + return null; + } + v.pendingChunks.decrementAndGet(); + v.unprocessedChunksByOffset.put(offset, new ChunkStatistics(offset, messageCount, dataSize)); + return v; + } + ); + } + + @Override + public void handleMessage( + byte subscriptionId, + long offset, + long chunkTimestamp, + long committedChunkId, + Message message + ) { + this.subscriptionStatisticsMap.compute( + subscriptionId, + (k, v) -> { + if(v == null) { + LOGGER.warn( + "handleMessage called for subscription that does not exist! subscriptionId={}", + subscriptionId + ); + return null; + } + NavigableMap subHeadMap = v.unprocessedChunksByOffset.headMap(offset, true); + Map.Entry lastEntry = subHeadMap.pollLastEntry(); + if(lastEntry == null) { + LOGGER.warn( + "handleMessage called but chunk was not found! subscriptionId={} offset={}", + subscriptionId, + offset + ); + return v; + } + ChunkStatistics statistics = lastEntry.getValue(); + statistics.unprocessedMessagesByOffset.put(offset, message); + return v; + } + ); + } + + @Override + public void handleUnsubscribe(byte subscriptionId) { + Object removed = this.subscriptionStatisticsMap.remove(subscriptionId); + if(removed == null) { + LOGGER.warn( + "handleUnsubscribe called for subscriptionId that does not exist! subscriptionId={}", + subscriptionId + ); + } + } + + /** + * Marks a message as handled, changing internal statistics. + * + * @param messageContext The {@link MessageHandler.Context} of the handled message + * @return Whether the message was marked as handled (returning {@code true}) + * or was not found (either because it was already marked as handled, or wasn't tracked) + */ + public boolean markHandled(MessageHandler.Context messageContext) { + AggregatedMessageStatistics entry = retrieveStatistics(messageContext); + if (entry == null) { + return false; + } + return markHandled(entry); + } + + /** + * Marks a message as handled, changing internal statistics. + * + * @param aggregatedMessageStatistics The {@link AggregatedMessageStatistics} of the handled message + * @return Whether the message was marked as handled (returning {@code true}) + * or was not found (either because it was already marked as handled, or wasn't tracked) + */ + public boolean markHandled(AggregatedMessageStatistics aggregatedMessageStatistics) { + // Can't remove, not enough information + if (aggregatedMessageStatistics.chunkStatistics == null + || aggregatedMessageStatistics.messageEntry == null + || aggregatedMessageStatistics.chunkHeadMap == null) { + return false; + } + Message removedMessage = aggregatedMessageStatistics.chunkStatistics + .unprocessedMessagesByOffset + .remove(aggregatedMessageStatistics.offset); + if (removedMessage == null) { + return false; + } + // Remove chunk from list of unprocessed chunks if all its messages have been processed + if (aggregatedMessageStatistics.chunkStatistics.unprocessedMessagesByOffset.isEmpty()) { + aggregatedMessageStatistics.chunkHeadMap.remove(aggregatedMessageStatistics.messageEntry.getKey(), aggregatedMessageStatistics.chunkStatistics); + } + return true; + } + + public AggregatedMessageStatistics retrieveStatistics(String stream, long offset) { + Set possibleSubscriptionIds = this.streamNameToSubscriptionIdMap.get(stream); + AggregatedMessageStatistics entry = null; + for (Byte subscriptionId : possibleSubscriptionIds) { + entry = retrieveStatistics(subscriptionId, offset); + if (entry == null) { + continue; + } + // We have all the info we need, we found the specific chunk. Stop right here + if (entry.chunkHeadMap != null && entry.chunkStatistics != null) { + return entry; + } + } + // Return the next-best result, because we might find the subscription but not the message + return entry; + } + + public AggregatedMessageStatistics retrieveStatistics(byte subscriptionId, long offset) { + SubscriptionStatistics subscriptionStatistics = this.subscriptionStatisticsMap.get(subscriptionId); + if (subscriptionStatistics == null) { + return null; + } + NavigableMap chunkStatisticsHeadMap = subscriptionStatistics.unprocessedChunksByOffset.headMap(offset, true); + Map.Entry messageEntry = chunkStatisticsHeadMap.lastEntry(); + ChunkStatistics chunkStatistics = messageEntry == null ? null : messageEntry.getValue(); + return new AggregatedMessageStatistics(offset, subscriptionStatistics, chunkStatisticsHeadMap, chunkStatistics, messageEntry); + } + + public AggregatedMessageStatistics retrieveStatistics(MessageHandler.Context messageContext) { + return retrieveStatistics(messageContext.stream(), messageContext.offset()); + } + + public static class AggregatedMessageStatistics { + + private final long offset; + private final SubscriptionStatistics subscriptionStatistics; + private final NavigableMap chunkHeadMap; + private final ChunkStatistics chunkStatistics; + private final Map.Entry messageEntry; + + public AggregatedMessageStatistics( + long offset, + @Nonnull SubscriptionStatistics subscriptionStatistics, + @Nullable NavigableMap chunkHeadMap, + @Nullable ChunkStatistics chunkStatistics, + @Nullable Map.Entry messageEntry) { + this.subscriptionStatistics = subscriptionStatistics; + this.chunkStatistics = chunkStatistics; + this.chunkHeadMap = chunkHeadMap; + this.messageEntry = messageEntry; + this.offset = offset; + } + + @Nonnull + public SubscriptionStatistics getSubscriptionStatistics() { + return subscriptionStatistics; + } + + @Nullable + public ChunkStatistics getChunkStatistics() { + return chunkStatistics; + } + + @Nullable + public NavigableMap getChunkHeadMap() { + return chunkHeadMap; + } + + @Nullable + public Map.Entry getMessageEntry() { + return messageEntry; + } + + public long getOffset() { + return offset; + } + + } + + public static class SubscriptionStatistics { + + private final byte subscriptionId; + private final String stream; + private final AtomicInteger pendingChunks = new AtomicInteger(0); + private final Map subscriptionProperties; + private final NavigableMap unprocessedChunksByOffset; + + public SubscriptionStatistics(byte subscriptionId, String stream, Map subscriptionProperties) { + this(subscriptionId, stream, subscriptionProperties, new ConcurrentSkipListMap<>()); + } + + public SubscriptionStatistics( + byte subscriptionId, + String stream, + Map subscriptionProperties, + NavigableMap unprocessedChunksByOffset + ) { + this.subscriptionId = subscriptionId; + this.stream = stream; + this.subscriptionProperties = subscriptionProperties; + this.unprocessedChunksByOffset = unprocessedChunksByOffset; + } + + public byte getSubscriptionId() { + return subscriptionId; + } + + public String getStream() { + return stream; + } + + public AtomicInteger getPendingChunks() { + return pendingChunks; + } + + public Map getSubscriptionProperties() { + return Collections.unmodifiableMap(subscriptionProperties); + } + + public NavigableMap getUnprocessedChunksByOffset() { + return Collections.unmodifiableNavigableMap(unprocessedChunksByOffset); + } + + } + + public static class ChunkStatistics { + + private final long offset; + private final long messageCount; + private final long dataSize; + private final Map unprocessedMessagesByOffset; + + public ChunkStatistics(long offset, long messageCount, long dataSize) { + this(offset, messageCount, dataSize, new ConcurrentHashMap<>()); + } + + public ChunkStatistics(long offset, long messageCount, long dataSize, Map unprocessedMessagesByOffset) { + this.offset = offset; + this.messageCount = messageCount; + this.dataSize = dataSize; + this.unprocessedMessagesByOffset = unprocessedMessagesByOffset; + } + + public long getOffset() { + return offset; + } + + public long getMessageCount() { + return messageCount; + } + + public long getDataSize() { + return dataSize; + } + + public Map getUnprocessedMessagesByOffset() { + return Collections.unmodifiableMap(unprocessedMessagesByOffset); + } + + } + + public Map> getStreamNameToSubscriptionIdMap() { + return Collections.unmodifiableMap(streamNameToSubscriptionIdMap); + } + + public Map getSubscriptionStatisticsMap() { + return Collections.unmodifiableMap(subscriptionStatisticsMap); + } + +} diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index e35a1fedf2..7958d52a23 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -13,40 +13,19 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.impl.Utils.convertCodeToException; -import static com.rabbitmq.stream.impl.Utils.formatConstant; -import static com.rabbitmq.stream.impl.Utils.isSac; -import static com.rabbitmq.stream.impl.Utils.jsonField; -import static com.rabbitmq.stream.impl.Utils.namedFunction; -import static com.rabbitmq.stream.impl.Utils.namedRunnable; -import static com.rabbitmq.stream.impl.Utils.quote; - import com.rabbitmq.stream.*; import com.rabbitmq.stream.MessageHandler.Context; import com.rabbitmq.stream.SubscriptionListener.SubscriptionContext; -import com.rabbitmq.stream.impl.Client.Broker; -import com.rabbitmq.stream.impl.Client.ChunkListener; -import com.rabbitmq.stream.impl.Client.ClientParameters; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.ConsumerUpdateListener; -import com.rabbitmq.stream.impl.Client.CreditNotification; -import com.rabbitmq.stream.impl.Client.MessageListener; -import com.rabbitmq.stream.impl.Client.MetadataListener; -import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; -import com.rabbitmq.stream.impl.Client.ShutdownListener; -import com.rabbitmq.stream.impl.Utils.ClientConnectionType; -import com.rabbitmq.stream.impl.Utils.ClientFactory; -import com.rabbitmq.stream.impl.Utils.ClientFactoryContext; -import java.util.ArrayList; -import java.util.Collection; -import java.util.Collections; -import java.util.Iterator; -import java.util.List; -import java.util.Map; +import com.rabbitmq.stream.impl.Client.*; +import com.rabbitmq.stream.impl.Utils.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.*; import java.util.Map.Entry; -import java.util.NavigableSet; -import java.util.Objects; -import java.util.Random; -import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentSkipListSet; import java.util.concurrent.CopyOnWriteArrayList; @@ -57,8 +36,8 @@ import java.util.function.Predicate; import java.util.stream.Collectors; import java.util.stream.IntStream; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; + +import static com.rabbitmq.stream.impl.Utils.*; class ConsumersCoordinator { @@ -564,8 +543,6 @@ private ClientSubscriptionsManager( AtomicReference clientReference = new AtomicReference<>(); ConsumerFlowControlStrategy localConsumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get); this.consumerFlowControlStrategy = localConsumerFlowControlStrategy; - ChunkListener chunkListener = (ignoredClient, subscriptionId, offset, messageCount, dataSize) -> - localConsumerFlowControlStrategy.handleChunk(subscriptionId, offset, messageCount, dataSize); CreditNotification creditNotification = (subscriptionId, responseCode) -> { SubscriptionTracker subscriptionTracker = @@ -731,7 +708,7 @@ private ClientSubscriptionsManager( ClientFactoryContext.fromParameters( clientParameters .clientProperty("connection_name", connectionName) - .chunkListener(chunkListener) + .chunkListener(localConsumerFlowControlStrategy) .creditNotification(creditNotification) .messageListener(messageListener) .shutdownListener(shutdownListener) @@ -956,7 +933,7 @@ synchronized void add( checkNotClosed(); byte subId = subscriptionId; - int initialCredits = this.consumerFlowControlStrategy.handleSubscribe( + int initialCredits = this.consumerFlowControlStrategy.handleSubscribeReturningInitialCredits( subId, subscriptionTracker.stream, subscriptionContext.offsetSpecification(), diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java index c7602815f0..0e71fe9256 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java @@ -13,17 +13,16 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.BackOffDelayPolicy.fixedWithInitialDelay; -import static com.rabbitmq.stream.impl.AsyncRetry.asyncRetry; -import static com.rabbitmq.stream.impl.Utils.offsetBefore; -import static java.time.Duration.ofMillis; - import com.rabbitmq.stream.*; import com.rabbitmq.stream.MessageHandler.Context; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; import com.rabbitmq.stream.impl.StreamEnvironment.TrackingConsumerRegistration; import com.rabbitmq.stream.impl.Utils.CompositeConsumerUpdateListener; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.util.Collections; import java.util.Map; import java.util.Objects; @@ -37,8 +36,11 @@ import java.util.function.LongConsumer; import java.util.function.LongSupplier; import java.util.function.Supplier; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; + +import static com.rabbitmq.stream.BackOffDelayPolicy.fixedWithInitialDelay; +import static com.rabbitmq.stream.impl.AsyncRetry.asyncRetry; +import static com.rabbitmq.stream.impl.Utils.offsetBefore; +import static java.time.Duration.ofMillis; class StreamConsumer implements Consumer { diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java index a32e1c343f..240269e02f 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java @@ -14,6 +14,9 @@ package com.rabbitmq.stream.impl; import com.rabbitmq.stream.*; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.impl.flow.LegacyConsumerFlowControlStrategyBuilderFactory; import java.lang.reflect.Field; import java.lang.reflect.Modifier; @@ -39,7 +42,7 @@ class StreamConsumerBuilder implements ConsumerBuilder { private ConsumerUpdateListener consumerUpdateListener; private int initialCredits = 1; private int additionalCredits = 1; - private ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder = LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(this); + private ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder = LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(this); public StreamConsumerBuilder(StreamEnvironment environment) { this.environment = environment; @@ -148,7 +151,7 @@ public StreamConsumerBuilder credits(int initial, int onChunkDelivery) { if (initial <= 0 || onChunkDelivery <= 0) { throw new IllegalArgumentException("Credits must be positive"); } - this.consumerFlowControlStrategyBuilder = LegacyFlowControlStrategyBuilderFactory.INSTANCE + this.consumerFlowControlStrategyBuilder = LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE .builder(this) .initialCredits(initial) .additionalCredits(additionalCredits); diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java index 50fec81d3e..eed5f52c3c 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java @@ -13,27 +13,26 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.impl.Utils.convertCodeToException; -import static com.rabbitmq.stream.impl.Utils.exceptionMessage; -import static com.rabbitmq.stream.impl.Utils.formatConstant; -import static com.rabbitmq.stream.impl.Utils.namedRunnable; -import static java.util.concurrent.TimeUnit.SECONDS; - import com.rabbitmq.stream.*; import com.rabbitmq.stream.MessageHandler.Context; import com.rabbitmq.stream.compression.CompressionCodecFactory; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.ClientParameters; import com.rabbitmq.stream.impl.Client.ShutdownListener; import com.rabbitmq.stream.impl.Client.StreamStatsResponse; import com.rabbitmq.stream.impl.OffsetTrackingCoordinator.Registration; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; import com.rabbitmq.stream.impl.StreamEnvironmentBuilder.DefaultTlsConfiguration; -import com.rabbitmq.stream.impl.Utils.ClientConnectionType; +import com.rabbitmq.stream.impl.Utils.*; import io.netty.buffer.ByteBufAllocator; import io.netty.channel.EventLoopGroup; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.handler.ssl.SslContext; import io.netty.handler.ssl.SslContextBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.net.ssl.SSLException; import java.io.IOException; import java.net.URI; import java.net.URLDecoder; @@ -42,24 +41,15 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeoutException; +import java.util.concurrent.*; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BiFunction; import java.util.function.Consumer; -import java.util.function.Function; -import java.util.function.LongConsumer; -import java.util.function.LongSupplier; -import java.util.function.Supplier; +import java.util.function.*; import java.util.stream.Collectors; -import javax.net.ssl.SSLException; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; + +import static com.rabbitmq.stream.impl.Utils.*; +import static java.util.concurrent.TimeUnit.SECONDS; class StreamEnvironment implements Environment { diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..888a22d389 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java @@ -0,0 +1,132 @@ +package com.rabbitmq.stream.impl.flow; + +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.flow.AbstractConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.CreditAsker; +import com.rabbitmq.stream.flow.MessageHandlingAware; +import com.rabbitmq.stream.impl.ConsumerStatisticRecorder; + +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.IntUnaryOperator; +import java.util.function.Supplier; + +/** + * Abstract class that calls an instance of {@link ConsumerStatisticRecorder} and exposes it to child implementations + * that may use its statistics to control flow as they see fit. + */ +public abstract class AbstractStatisticRecordingConsumerFlowControlStrategy + extends AbstractConsumerFlowControlStrategy + implements MessageHandlingAware { + + protected final ConsumerStatisticRecorder consumerStatisticRecorder = new ConsumerStatisticRecorder(); + + protected AbstractStatisticRecordingConsumerFlowControlStrategy(Supplier creditAskerSupplier) { + super(creditAskerSupplier); + } + + /** + * Note for implementors: This method MUST be called from the implementation of + * {@link ConsumerFlowControlStrategy#handleSubscribeReturningInitialCredits}, + * otherwise statistics will not be registered! + *

+ * {@inheritDoc} + */ + @Override + public void handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ) { + super.handleSubscribe(subscriptionId, stream, offsetSpecification, subscriptionProperties); + this.consumerStatisticRecorder.handleSubscribe( + subscriptionId, + stream, + offsetSpecification, + subscriptionProperties + ); + } + + @Override + public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + super.handleChunk(subscriptionId, offset, messageCount, dataSize); + this.consumerStatisticRecorder.handleChunk( + subscriptionId, + offset, + messageCount, + dataSize + ); + } + + @Override + public void handleMessage( + byte subscriptionId, + long offset, + long chunkTimestamp, + long committedChunkId, + Message message + ) { + super.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); + this.consumerStatisticRecorder.handleMessage( + subscriptionId, + offset, + chunkTimestamp, + committedChunkId, + message + ); + } + + @Override + public void handleCreditNotification(byte subscriptionId, short responseCode) { + super.handleCreditNotification(subscriptionId, responseCode); + this.consumerStatisticRecorder.handleCreditNotification(subscriptionId, responseCode); + } + + @Override + public void handleUnsubscribe(byte subscriptionId) { + super.handleUnsubscribe(subscriptionId); + this.consumerStatisticRecorder.handleUnsubscribe(subscriptionId); + } + + protected int registerCredits(byte subscriptionId, IntUnaryOperator askedToAsk, boolean askForCredits) { + AtomicInteger outerCreditsToAsk = new AtomicInteger(); + ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStatistics = this.consumerStatisticRecorder + .getSubscriptionStatisticsMap() + .get(subscriptionId); + subscriptionStatistics.getPendingChunks().updateAndGet(credits -> { + int creditsToAsk = askedToAsk.applyAsInt(credits); + outerCreditsToAsk.set(creditsToAsk); + return credits + creditsToAsk; + }); + int finalCreditsToAsk = outerCreditsToAsk.get(); + if(askForCredits && finalCreditsToAsk > 0) { + mandatoryClient().credit(subscriptionId, finalCreditsToAsk); + } + return finalCreditsToAsk; + } + + @Override + public boolean markHandled(MessageHandler.Context messageContext) { + ConsumerStatisticRecorder.AggregatedMessageStatistics messageStatistics = this.consumerStatisticRecorder + .retrieveStatistics(messageContext); + if(messageStatistics == null) { + return false; + } + boolean markedAsHandled = this.consumerStatisticRecorder.markHandled(messageStatistics); + if(!markedAsHandled) { + return false; + } + afterMarkHandledStateChanged(messageContext, messageStatistics); + return true; + } + + protected void afterMarkHandledStateChanged( + MessageHandler.Context messageContext, + ConsumerStatisticRecorder.AggregatedMessageStatistics messageStatistics) { + // Default no-op callback + } +} diff --git a/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java similarity index 59% rename from src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java rename to src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java index 27a5966061..4b6598f42b 100644 --- a/src/main/java/com/rabbitmq/stream/LegacyFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java @@ -1,6 +1,8 @@ -package com.rabbitmq.stream; +package com.rabbitmq.stream.impl.flow; -import com.rabbitmq.stream.impl.Client; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.flow.AbstractConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.CreditAsker; import java.util.Map; import java.util.function.Supplier; @@ -9,27 +11,19 @@ * The flow control strategy that was always applied before the flow control strategy mechanism existed in the codebase. * Requests a set amount of credits after each chunk arrives. */ -public class LegacyFlowControlStrategy extends AbstractFlowControlStrategy { +public class LegacyConsumerFlowControlStrategy extends AbstractConsumerFlowControlStrategy { private final int initialCredits; private final int additionalCredits; - public LegacyFlowControlStrategy(Supplier clientSupplier) { - this(clientSupplier, 1); - } - - public LegacyFlowControlStrategy(Supplier clientSupplier, int initialCredits) { - this(clientSupplier, initialCredits, 1); - } - - public LegacyFlowControlStrategy(Supplier clientSupplier, int initialCredits, int additionalCredits) { - super(clientSupplier); + public LegacyConsumerFlowControlStrategy(Supplier creditAskerSupplier, int initialCredits, int additionalCredits) { + super(creditAskerSupplier); this.initialCredits = initialCredits; this.additionalCredits = additionalCredits; } @Override - public int handleSubscribe( + public int handleSubscribeReturningInitialCredits( byte subscriptionId, String stream, OffsetSpecification offsetSpecification, diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java new file mode 100644 index 0000000000..5ca5aaaa43 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java @@ -0,0 +1,52 @@ +package com.rabbitmq.stream.impl.flow; + +import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.flow.CreditAsker; + +import java.util.function.Supplier; + +public class LegacyConsumerFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { + + public static final LegacyConsumerFlowControlStrategyBuilderFactory INSTANCE = new LegacyConsumerFlowControlStrategyBuilderFactory(); + + @Override + public Builder builder(ConsumerBuilder consumerBuilder) { + return new Builder(consumerBuilder); + } + + public static class Builder implements ConsumerFlowControlStrategyBuilder { + + private final ConsumerBuilder consumerBuilder; + + private int initialCredits = 1; + + private int additionalCredits = 1; + + public Builder(ConsumerBuilder consumerBuilder) { + this.consumerBuilder = consumerBuilder; + } + + @Override + public LegacyConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { + return new LegacyConsumerFlowControlStrategy(creditAskerSupplier, this.initialCredits, this.additionalCredits); + } + + @Override + public ConsumerBuilder builder() { + return this.consumerBuilder; + } + + public Builder additionalCredits(int additionalCredits) { + this.additionalCredits = additionalCredits; + return this; + } + + public Builder initialCredits(int initialCredits) { + this.initialCredits = initialCredits; + return this; + } + } + +} diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..16404c2d1a --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java @@ -0,0 +1,82 @@ +package com.rabbitmq.stream.impl.flow; + +import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.flow.CreditAsker; +import com.rabbitmq.stream.flow.MessageHandlingAware; +import com.rabbitmq.stream.impl.ConsumerStatisticRecorder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.function.IntUnaryOperator; +import java.util.function.Supplier; + +/** + * A flow control strategy that enforces a maximum amount of Inflight chunks per registered subscription. + * Based on {@link MessageHandlingAware message acknowledgement}, asking for the maximum number of chunks possible, given the limit. + */ +public class MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy extends AbstractStatisticRecordingConsumerFlowControlStrategy { + + private static final Logger LOGGER = LoggerFactory.getLogger(MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.class); + + private final int maximumSimultaneousChunksPerSubscription; + + public MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy( + Supplier creditAskerSupplier, + int maximumSimultaneousChunksPerSubscription + ) { + super(creditAskerSupplier); + if(maximumSimultaneousChunksPerSubscription <= 0) { + throw new IllegalArgumentException( + "maximumSimultaneousChunksPerSubscription must be greater than 0. Was: " + maximumSimultaneousChunksPerSubscription + ); + } + this.maximumSimultaneousChunksPerSubscription = maximumSimultaneousChunksPerSubscription; + } + + @Override + public int handleSubscribeReturningInitialCredits( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties) { + this.handleSubscribe( + subscriptionId, + stream, + offsetSpecification, + subscriptionProperties + ); + return registerCredits(subscriptionId, getCreditAsker(subscriptionId), false); + } + + @Override + protected void afterMarkHandledStateChanged( + MessageHandler.Context messageContext, + ConsumerStatisticRecorder.AggregatedMessageStatistics messageStatistics) { + byte subscriptionId = messageStatistics.getSubscriptionStatistics().getSubscriptionId(); + registerCredits(subscriptionId, getCreditAsker(subscriptionId), true); + } + + private IntUnaryOperator getCreditAsker(byte subscriptionId) { + return pendingChunks -> { + int inProcessingChunks = extractInProcessingChunks(subscriptionId); + return Math.max(0, this.maximumSimultaneousChunksPerSubscription - (pendingChunks + inProcessingChunks)); + }; + } + + private int extractInProcessingChunks(byte subscriptionId) { + int inProcessingChunks; + ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStats = this.consumerStatisticRecorder + .getSubscriptionStatisticsMap() + .get(subscriptionId); + if(subscriptionStats == null) { + LOGGER.warn("Subscription data not found while calculating credits to ask! subscriptionId: {}", subscriptionId); + inProcessingChunks = 0; + } else { + inProcessingChunks = subscriptionStats.getUnprocessedChunksByOffset().size(); + } + return inProcessingChunks; + } + +} diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java new file mode 100644 index 0000000000..e426645548 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java @@ -0,0 +1,46 @@ +package com.rabbitmq.stream.impl.flow; + +import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.flow.CreditAsker; + +import java.util.function.Supplier; + +public class MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { + + public static final MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory INSTANCE = new MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory(); + + @Override + public Builder builder(ConsumerBuilder consumerBuilder) { + return new Builder(consumerBuilder); + } + + public static class Builder implements ConsumerFlowControlStrategyBuilder { + + private final ConsumerBuilder consumerBuilder; + + private int maximumInflightChunksPerSubscription = 1; + + public Builder(ConsumerBuilder consumerBuilder) { + this.consumerBuilder = consumerBuilder; + } + + @Override + public MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { + return new MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy(creditAskerSupplier, this.maximumInflightChunksPerSubscription); + } + + @Override + public ConsumerBuilder builder() { + return this.consumerBuilder; + } + + public Builder maximumInflightChunksPerSubscription(int maximumInflightChunksPerSubscription) { + this.maximumInflightChunksPerSubscription = maximumInflightChunksPerSubscription; + return this; + } + + } + +} diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index 85160194e2..d23eecd66d 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -13,46 +13,16 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.BackOffDelayPolicy.fixedWithInitialDelay; -import static com.rabbitmq.stream.impl.TestUtils.b; -import static com.rabbitmq.stream.impl.TestUtils.latchAssert; -import static com.rabbitmq.stream.impl.TestUtils.metadata; -import static com.rabbitmq.stream.impl.TestUtils.namedConsumer; -import static com.rabbitmq.stream.impl.TestUtils.waitAtMost; -import static java.lang.String.format; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - import com.rabbitmq.stream.*; import com.rabbitmq.stream.codec.WrapperMessageBuilder; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.MessageListener; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.Client.Response; import com.rabbitmq.stream.impl.MonitoringTestUtils.ConsumerCoordinatorInfo; import com.rabbitmq.stream.impl.Utils.ClientFactory; -import java.time.Duration; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.function.Consumer; -import java.util.function.Supplier; -import java.util.stream.Collectors; -import java.util.stream.IntStream; -import java.util.stream.Stream; +import com.rabbitmq.stream.impl.flow.LegacyConsumerFlowControlStrategyBuilderFactory; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -60,6 +30,23 @@ import org.junit.jupiter.params.provider.MethodSource; import org.mockito.*; +import java.time.Duration; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static com.rabbitmq.stream.BackOffDelayPolicy.fixedWithInitialDelay; +import static com.rabbitmq.stream.impl.TestUtils.*; +import static java.lang.String.format; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + public class ConsumersCoordinatorTest { private static final SubscriptionListener NO_OP_SUBSCRIPTION_LISTENER = subscriptionContext -> {}; @@ -192,7 +179,7 @@ void tearDown() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(2)).client(any()); @@ -234,7 +221,7 @@ void shouldGetExactNodeImmediatelyWithAdvertisedHostNameClientFactoryAndExactNod NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -265,7 +252,7 @@ void shouldSubscribeWithEmptyPropertiesWithUnamedConsumer() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -287,7 +274,7 @@ void subscribeShouldThrowExceptionWhenNoMetadataForTheStream() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits)) .isInstanceOf(StreamDoesNotExistException.class); @@ -307,7 +294,7 @@ void subscribeShouldThrowExceptionWhenStreamDoesNotExist() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits)) .isInstanceOf(StreamDoesNotExistException.class); @@ -337,7 +324,7 @@ void subscribePropagateExceptionWhenClientSubscriptionFails() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits)) .isInstanceOf(StreamException.class) @@ -359,7 +346,7 @@ void subscribeShouldThrowExceptionWhenMetadataResponseIsNotOk() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits)) .isInstanceOf(IllegalStateException.class); @@ -378,7 +365,7 @@ void subscribeShouldThrowExceptionIfNoNodeAvailableForStream() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits)) .isInstanceOf(IllegalStateException.class); @@ -420,7 +407,7 @@ void subscribeShouldSubscribeToStreamAndDispatchMessage_UnsubscribeShouldUnsubsc NO_OP_SUBSCRIPTION_LISTENER, () -> trackingClosingCallbackCalls.incrementAndGet(), (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -471,7 +458,7 @@ void subscribeShouldSubscribeToStreamAndDispatchMessageWithManySubscriptions() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.compute(subId, (k, v) -> (v == null) ? 1 : ++v), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); closingRunnables.add(closingRunnable); @@ -547,7 +534,7 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -567,7 +554,7 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); @@ -648,7 +635,7 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -704,7 +691,7 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -719,7 +706,7 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); @@ -808,7 +795,7 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -874,7 +861,7 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -928,7 +915,7 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -983,7 +970,7 @@ void shouldUseNewClientsForMoreThanMaxSubscriptionsAndCloseClientAfterUnsubscrip NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits)) .collect(Collectors.toList()); @@ -1045,7 +1032,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); }); @@ -1071,7 +1058,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); @@ -1112,7 +1099,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); }); @@ -1144,7 +1131,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); @@ -1195,7 +1182,7 @@ void shouldRestartWhereItLeftOffAfterDisruption(Consumer {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1268,7 +1255,7 @@ void shouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesReceived( NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1343,7 +1330,7 @@ void shouldUseStoredOffsetOnRecovery(Consumer configur NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1430,7 +1417,7 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1445,7 +1432,7 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1511,7 +1498,7 @@ void shouldRetryAssignmentOnRecoveryStreamNotAvailableFailure() throws Exception NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1577,7 +1564,7 @@ void shouldRetryAssignmentOnRecoveryCandidateLookupFailure() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); verify(clientFactory, times(1)).client(any()); @@ -1620,7 +1607,7 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), Collections.emptyMap(), initialCredits); @@ -1693,7 +1680,8 @@ void shouldCallConsumerFlowControlHandlers(Consumer co int numberOfInitialCreditsOnSubscribe = 7; - when(mockedConsumerFlowControlStrategy.handleSubscribe(anyByte(), anyString(), any(), anyMap())).thenReturn(numberOfInitialCreditsOnSubscribe); + when(mockedConsumerFlowControlStrategy.handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(), anyMap())) + .thenReturn(numberOfInitialCreditsOnSubscribe); ConsumerFlowControlStrategyBuilder mockedConsumerFlowControlStrategyBuilder = Mockito.mock(ConsumerFlowControlStrategyBuilder.class); when(mockedConsumerFlowControlStrategyBuilder.build(any())).thenReturn(mockedConsumerFlowControlStrategy); From 9a08c72154109373ad6e9d7cd9c8b10faf84cbb2 Mon Sep 17 00:00:00 2001 From: henry701 Date: Sun, 11 Jun 2023 16:06:48 -0300 Subject: [PATCH 04/14] Abstract dependency on Client away from public Flow Control classes by using an adapter for the callback interface --- .../stream/CallbackStreamDataHandler.java | 64 ++++++++++ .../AbstractConsumerFlowControlStrategy.java | 6 +- .../flow/ConsumerFlowControlStrategy.java | 7 +- ...lientCallbackStreamDataHandlerAdapter.java | 76 +++++++++++ .../stream/impl/ClientDataHandler.java | 120 ------------------ .../impl/ConsumerStatisticRecorder.java | 3 +- .../stream/impl/ConsumersCoordinator.java | 8 +- ...cRecordingConsumerFlowControlStrategy.java | 2 +- .../LegacyConsumerFlowControlStrategy.java | 2 +- 9 files changed, 154 insertions(+), 134 deletions(-) create mode 100644 src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java create mode 100644 src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java delete mode 100644 src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java diff --git a/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java b/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java new file mode 100644 index 0000000000..b26cec38ad --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java @@ -0,0 +1,64 @@ +package com.rabbitmq.stream; + +import java.util.Map; + +/** + * Exposes callbacks to handle events from a particular Stream connection, + * with specific names for methods and no connection-oriented parameter. + */ +public interface CallbackStreamDataHandler { + + default void handlePublishConfirm(byte publisherId, long publishingId) { + // No-op by default + } + + default void handlePublishError(byte publisherId, long publishingId, short errorCode) { + // No-op by default + } + + default void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + // No-op by default + } + + default void handleMessage(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { + // No-op by default + } + + default void handleCreditNotification(byte subscriptionId, short responseCode) { + // No-op by default + } + + default void handleConsumerUpdate(byte subscriptionId, boolean active) { + // No-op by default + } + + default void handleMetadata(String stream, short code) { + // No-op by default + } + + /** + * Callback for handling a new stream subscription. + * + * @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 + */ + default void handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ) { + // No-op by default + } + + /** + * 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 + } + +} diff --git a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java index fd34cf784d..e45015ca0a 100644 --- a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java @@ -1,13 +1,11 @@ package com.rabbitmq.stream.flow; -import com.rabbitmq.stream.impl.Client; - import java.util.Objects; import java.util.function.Supplier; /** * Abstract class for Consumer Flow Control Strategies which keeps a cached lazily-initialized - * {@link Client} ready for retrieval by its inheritors. + * {@link CreditAsker} ready for retrieval by its inheritors. */ public abstract class AbstractConsumerFlowControlStrategy implements ConsumerFlowControlStrategy { @@ -19,7 +17,7 @@ protected AbstractConsumerFlowControlStrategy(Supplier creditAskerS this.creditAskerSupplier = Objects.requireNonNull(creditAskerSupplier, "creditAskerSupplier"); } - protected CreditAsker mandatoryClient() { + protected CreditAsker mandatoryCreditAsker() { CreditAsker localSupplied = this.creditAsker; if(localSupplied != null) { return localSupplied; diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java index 3ce242ef00..096ed95bd1 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java @@ -1,7 +1,7 @@ package com.rabbitmq.stream.flow; +import com.rabbitmq.stream.CallbackStreamDataHandler; import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.impl.ClientDataHandler; import java.util.Map; @@ -10,14 +10,13 @@ * Implementations may freely implement reactions to the various client callbacks. * When defined by each implementation, it may internally call {@link CreditAsker#credit} to ask for credits. */ -// TODO: Decouple from ClientDataHandler. Maybe create an adapter pattern for handling this, or something. -public interface ConsumerFlowControlStrategy extends ClientDataHandler { +public interface ConsumerFlowControlStrategy extends CallbackStreamDataHandler { /** * Callback for handling a new stream subscription. * Called right before the subscription is sent to the actual client. *

- * Either this variant or {@link ClientDataHandler#handleSubscribe(byte, String, OffsetSpecification, Map)} should be called, NOT both. + * Either this variant or {@link CallbackStreamDataHandler#handleSubscribe(byte, String, OffsetSpecification, Map)} should be called, NOT both. *

* * @param subscriptionId The subscriptionId as specified by the Stream Protocol diff --git a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java new file mode 100644 index 0000000000..1bc320830d --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java @@ -0,0 +1,76 @@ +package com.rabbitmq.stream.impl; + +import com.rabbitmq.stream.CallbackStreamDataHandler; +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.OffsetSpecification; + +public class ClientCallbackStreamDataHandlerAdapter implements + CallbackStreamDataHandler, + Client.PublishConfirmListener, + Client.PublishErrorListener, + Client.ChunkListener, + Client.MessageListener, + Client.CreditNotification, + Client.ConsumerUpdateListener, + Client.ShutdownListener, + Client.MetadataListener { + + private final CallbackStreamDataHandler callbackStreamDataHandler; + + public ClientCallbackStreamDataHandlerAdapter(CallbackStreamDataHandler callbackStreamDataHandler) { + this.callbackStreamDataHandler = callbackStreamDataHandler; + } + + + @Override + public void handle(byte publisherId, long publishingId) { + this.handlePublishConfirm(publisherId, publishingId); + } + + @Override + public void handle(byte publisherId, long publishingId, short errorCode) { + this.handlePublishError(publisherId, publishingId, errorCode); + } + + @Override + public void handle(Client client, byte subscriptionId, long offset, long messageCount, long dataSize) { + this.handleChunk(subscriptionId, offset, messageCount, dataSize); + } + + @Override + public void handle(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { + this.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); + } + + @Override + public void handle(byte subscriptionId, short responseCode) { + this.handleCreditNotification(subscriptionId, responseCode); + } + + @Override + public OffsetSpecification handle(Client client, byte subscriptionId, boolean active) { + this.handleConsumerUpdate(subscriptionId, active); + return null; + } + + @Override + public void handle(Client.ShutdownContext shutdownContext) { + this.handleShutdown(shutdownContext); + } + + void handleShutdown(Client.ShutdownContext shutdownContext) { + if(callbackStreamDataHandler instanceof AutoCloseable) { + try { + ((AutoCloseable) callbackStreamDataHandler).close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + } + + @Override + public void handle(String stream, short code) { + this.handleMetadata(stream, code); + } + +} diff --git a/src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java b/src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java deleted file mode 100644 index c91cbdbc39..0000000000 --- a/src/main/java/com/rabbitmq/stream/impl/ClientDataHandler.java +++ /dev/null @@ -1,120 +0,0 @@ -package com.rabbitmq.stream.impl; - -import com.rabbitmq.stream.Message; -import com.rabbitmq.stream.OffsetSpecification; - -import java.util.Map; - -/** - * 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, - Client.PublishErrorListener, - Client.ChunkListener, - Client.MessageListener, - Client.CreditNotification, - Client.ConsumerUpdateListener, - Client.ShutdownListener, - Client.MetadataListener { - - @Override - default void handle(byte publisherId, long publishingId) { - this.handlePublishConfirm(publisherId, publishingId); - } - - default void handlePublishConfirm(byte publisherId, long publishingId) { - // No-op by default - } - - @Override - default void handle(byte publisherId, long publishingId, short errorCode) { - this.handlePublishError(publisherId, publishingId, errorCode); - } - - default void handlePublishError(byte publisherId, long publishingId, short errorCode) { - // No-op by default - } - - @Override - default void handle(Client client, byte subscriptionId, long offset, long messageCount, long dataSize) { - this.handleChunk(subscriptionId, offset, messageCount, dataSize); - } - - default void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - // No-op by default - } - - @Override - default void handle(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { - this.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); - } - - default void handleMessage(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { - // No-op by default - } - - @Override - default void handle(byte subscriptionId, short responseCode) { - this.handleCreditNotification(subscriptionId, responseCode); - } - - default void handleCreditNotification(byte subscriptionId, short responseCode) { - // No-op by default - } - - @Override - default OffsetSpecification handle(Client client, byte subscriptionId, boolean active) { - this.handleConsumerUpdate(subscriptionId, active); - return null; - } - - default void handleConsumerUpdate(byte subscriptionId, boolean active) { - // No-op by default - } - - @Override - default void handle(Client.ShutdownContext shutdownContext) { - this.handleShutdown(shutdownContext); - } - - default void handleShutdown(Client.ShutdownContext shutdownContext) { - // No-op by default - } - - @Override - default void handle(String stream, short code) { - this.handleMetadata(stream, code); - } - - default void handleMetadata(String stream, short code) { - // No-op by default - } - - /** - * Callback for handling a new stream subscription. - * - * @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 - */ - default void handleSubscribe( - byte subscriptionId, - String stream, - OffsetSpecification offsetSpecification, - Map subscriptionProperties - ) { - // No-op by default - } - - /** - * 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 - } - -} diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java index 7897caa308..75addf3674 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java @@ -1,5 +1,6 @@ package com.rabbitmq.stream.impl; +import com.rabbitmq.stream.CallbackStreamDataHandler; import com.rabbitmq.stream.Message; import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.OffsetSpecification; @@ -17,7 +18,7 @@ import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicInteger; -public class ConsumerStatisticRecorder implements ClientDataHandler { +public class ConsumerStatisticRecorder implements CallbackStreamDataHandler { private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerStatisticRecorder.class); diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 7958d52a23..4fc4dd962d 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -20,7 +20,6 @@ import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.ConsumerUpdateListener; import com.rabbitmq.stream.impl.Client.*; -import com.rabbitmq.stream.impl.Utils.*; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -542,6 +541,9 @@ private ClientSubscriptionsManager( this.trackerCount = 0; AtomicReference clientReference = new AtomicReference<>(); ConsumerFlowControlStrategy localConsumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get); + ClientCallbackStreamDataHandlerAdapter clientListenerAdaptedConsumerFlowControlStrategy = new ClientCallbackStreamDataHandlerAdapter( + localConsumerFlowControlStrategy + ); this.consumerFlowControlStrategy = localConsumerFlowControlStrategy; CreditNotification creditNotification = (subscriptionId, responseCode) -> { @@ -629,7 +631,7 @@ private ClientSubscriptionsManager( "Consumers re-assignment after disconnection from %s", name)); } - localConsumerFlowControlStrategy.handleShutdown(shutdownContext); + clientListenerAdaptedConsumerFlowControlStrategy.handleShutdown(shutdownContext); }; MetadataListener metadataListener = (stream, code) -> { @@ -708,7 +710,7 @@ private ClientSubscriptionsManager( ClientFactoryContext.fromParameters( clientParameters .clientProperty("connection_name", connectionName) - .chunkListener(localConsumerFlowControlStrategy) + .chunkListener(clientListenerAdaptedConsumerFlowControlStrategy) .creditNotification(creditNotification) .messageListener(messageListener) .shutdownListener(shutdownListener) diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java index 888a22d389..992f970c17 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java @@ -104,7 +104,7 @@ protected int registerCredits(byte subscriptionId, IntUnaryOperator askedToAsk, }); int finalCreditsToAsk = outerCreditsToAsk.get(); if(askForCredits && finalCreditsToAsk > 0) { - mandatoryClient().credit(subscriptionId, finalCreditsToAsk); + mandatoryCreditAsker().credit(subscriptionId, finalCreditsToAsk); } return finalCreditsToAsk; } diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java index 4b6598f42b..f92c25fb24 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java @@ -34,7 +34,7 @@ public int handleSubscribeReturningInitialCredits( @Override public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - mandatoryClient().credit(subscriptionId, this.additionalCredits); + mandatoryCreditAsker().credit(subscriptionId, this.additionalCredits); } } From 704e7b0be740e314564ed45c50bb743c2be51b0b Mon Sep 17 00:00:00 2001 From: henry701 Date: Tue, 13 Jun 2023 19:26:17 -0300 Subject: [PATCH 05/14] Remove obsolete 'initialCredits' and 'additionalCredits' parameters --- .../stream/impl/ConsumersCoordinator.java | 40 +++-- .../rabbitmq/stream/impl/StreamConsumer.java | 17 +- .../stream/impl/StreamConsumerBuilder.java | 21 ++- .../stream/impl/StreamEnvironment.java | 36 +++- .../stream/impl/ConsumersCoordinatorTest.java | 155 ++++++++++-------- 5 files changed, 167 insertions(+), 102 deletions(-) diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 4fc4dd962d..4ee632cdc3 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -13,18 +13,41 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.BackOffDelayPolicy; +import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.Consumer; +import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.MessageHandler.Context; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.StreamDoesNotExistException; +import com.rabbitmq.stream.StreamException; +import com.rabbitmq.stream.StreamNotAvailableException; +import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.SubscriptionListener.SubscriptionContext; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.impl.Client.Broker; +import com.rabbitmq.stream.impl.Client.ClientParameters; import com.rabbitmq.stream.impl.Client.ConsumerUpdateListener; -import com.rabbitmq.stream.impl.Client.*; +import com.rabbitmq.stream.impl.Client.CreditNotification; +import com.rabbitmq.stream.impl.Client.MessageListener; +import com.rabbitmq.stream.impl.Client.MetadataListener; +import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; +import com.rabbitmq.stream.impl.Client.ShutdownListener; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.*; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; import java.util.Map.Entry; +import java.util.NavigableSet; +import java.util.Objects; +import java.util.Random; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentSkipListSet; import java.util.concurrent.CopyOnWriteArrayList; @@ -92,8 +115,7 @@ Runnable subscribe( Runnable trackingClosingCallback, MessageHandler messageHandler, ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, - Map subscriptionProperties, - int initialCredits) { + Map subscriptionProperties) { List candidates = findBrokersForStream(stream); Client.Broker newNode = pickBroker(candidates); if (newNode == null) { @@ -113,8 +135,7 @@ Runnable subscribe( trackingClosingCallback, messageHandler, consumerFlowControlStrategyBuilder, - subscriptionProperties, - initialCredits); + subscriptionProperties); try { addToManager(newNode, subscriptionTracker, offsetSpecification, true); @@ -374,7 +395,6 @@ private static class SubscriptionTracker { private volatile ClientSubscriptionsManager manager; private volatile AtomicReference state = new AtomicReference<>(SubscriptionState.OPENING); - private final int initialCredits; private SubscriptionTracker( long id, @@ -386,8 +406,7 @@ private SubscriptionTracker( Runnable trackingClosingCallback, MessageHandler messageHandler, ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, - Map subscriptionProperties, - int initialCredits) { + Map subscriptionProperties) { this.id = id; this.consumer = consumer; this.stream = stream; @@ -397,7 +416,6 @@ private SubscriptionTracker( this.trackingClosingCallback = trackingClosingCallback; this.messageHandler = messageHandler; this.consumerFlowControlStrategyBuilder = consumerFlowControlStrategyBuilder; - this.initialCredits = initialCredits; if (this.offsetTrackingReference == null) { this.subscriptionProperties = subscriptionProperties; } else { diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java index 0e71fe9256..4472f42742 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java @@ -13,8 +13,15 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.Consumer; +import com.rabbitmq.stream.ConsumerUpdateListener; +import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.MessageHandler.Context; +import com.rabbitmq.stream.NoOffsetException; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.StreamException; +import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; @@ -76,9 +83,7 @@ class StreamConsumer implements Consumer { boolean lazyInit, SubscriptionListener subscriptionListener, Map subscriptionProperties, - ConsumerUpdateListener consumerUpdateListener, - int initialCredits, - int additionalCredits) { + ConsumerUpdateListener consumerUpdateListener) { this.id = ID_SEQUENCE.getAndIncrement(); Runnable trackingClosingCallback; @@ -252,8 +257,8 @@ class StreamConsumer implements Consumer { trackingClosingCallback, closedAwareMessageHandler, consumerFlowControlStrategyBuilder, - Collections.unmodifiableMap(subscriptionProperties), - initialCredits); + Collections.unmodifiableMap(subscriptionProperties) + ); this.status = Status.RUNNING; }; diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java index 240269e02f..ac8f81413d 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java @@ -13,7 +13,13 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.Consumer; +import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.ConsumerUpdateListener; +import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.StreamException; +import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; import com.rabbitmq.stream.impl.flow.LegacyConsumerFlowControlStrategyBuilderFactory; @@ -38,10 +44,8 @@ class StreamConsumerBuilder implements ConsumerBuilder { private boolean noTrackingStrategy = false; private boolean lazyInit = false; private SubscriptionListener subscriptionListener = subscriptionContext -> {}; - private Map subscriptionProperties = new ConcurrentHashMap<>(); + private final Map subscriptionProperties = new ConcurrentHashMap<>(); private ConsumerUpdateListener consumerUpdateListener; - private int initialCredits = 1; - private int additionalCredits = 1; private ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder = LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(this); public StreamConsumerBuilder(StreamEnvironment environment) { @@ -143,7 +147,7 @@ public ConsumerBuilder noTrackingStrategy() { * @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)} + * @deprecated Prefer using {@link ConsumerBuilder#flowControlStrategy} * to define flow control strategies instead. */ @Deprecated @@ -154,7 +158,7 @@ public StreamConsumerBuilder credits(int initial, int onChunkDelivery) { this.consumerFlowControlStrategyBuilder = LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE .builder(this) .initialCredits(initial) - .additionalCredits(additionalCredits); + .additionalCredits(onChunkDelivery); return this; } @@ -221,9 +225,8 @@ public Consumer build() { this.lazyInit, this.subscriptionListener, this.subscriptionProperties, - this.consumerUpdateListener, - this.initialCredits, - this.additionalCredits); + this.consumerUpdateListener + ); environment.addConsumer((StreamConsumer) consumer); } else { if (Utils.isSac(this.subscriptionProperties)) { diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java index eed5f52c3c..b8831d558a 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java @@ -13,8 +13,21 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.Address; +import com.rabbitmq.stream.AddressResolver; +import com.rabbitmq.stream.BackOffDelayPolicy; +import com.rabbitmq.stream.Codec; +import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.Environment; +import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.MessageHandler.Context; +import com.rabbitmq.stream.NoOffsetException; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.ProducerBuilder; +import com.rabbitmq.stream.StreamCreator; +import com.rabbitmq.stream.StreamException; +import com.rabbitmq.stream.StreamStats; +import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.compression.CompressionCodecFactory; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.impl.Client.ClientParameters; @@ -23,7 +36,6 @@ import com.rabbitmq.stream.impl.OffsetTrackingCoordinator.Registration; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; import com.rabbitmq.stream.impl.StreamEnvironmentBuilder.DefaultTlsConfiguration; -import com.rabbitmq.stream.impl.Utils.*; import io.netty.buffer.ByteBufAllocator; import io.netty.channel.EventLoopGroup; import io.netty.channel.nio.NioEventLoopGroup; @@ -41,11 +53,20 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.*; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; import java.util.function.Consumer; -import java.util.function.*; +import java.util.function.Function; +import java.util.function.LongConsumer; +import java.util.function.LongSupplier; +import java.util.function.Supplier; import java.util.stream.Collectors; import static com.rabbitmq.stream.impl.Utils.*; @@ -636,8 +657,7 @@ Runnable registerConsumer( Runnable trackingClosingCallback, MessageHandler messageHandler, ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, - Map subscriptionProperties, - int initialCredits) { + Map subscriptionProperties) { return this.consumersCoordinator.subscribe( consumer, stream, @@ -647,8 +667,8 @@ Runnable registerConsumer( trackingClosingCallback, messageHandler, consumerFlowControlStrategyBuilder, - subscriptionProperties, - initialCredits); + subscriptionProperties + ); } Runnable registerProducer(StreamProducer producer, String reference, String stream) { diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index d23eecd66d..614c73f6e0 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -13,7 +13,13 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.BackOffDelayPolicy; +import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.StreamDoesNotExistException; +import com.rabbitmq.stream.StreamException; +import com.rabbitmq.stream.SubscriptionListener; import com.rabbitmq.stream.codec.WrapperMessageBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; @@ -28,11 +34,24 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; -import org.mockito.*; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; import java.time.Duration; -import java.util.*; -import java.util.concurrent.*; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.stream.Collectors; @@ -180,8 +199,8 @@ void tearDown() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(2)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -222,8 +241,8 @@ void shouldGetExactNodeImmediatelyWithAdvertisedHostNameClientFactoryAndExactNod NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -253,8 +272,8 @@ void shouldSubscribeWithEmptyPropertiesWithUnamedConsumer() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -275,8 +294,8 @@ void subscribeShouldThrowExceptionWhenNoMetadataForTheStream() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits)) + Collections.emptyMap() + )) .isInstanceOf(StreamDoesNotExistException.class); } @@ -295,8 +314,8 @@ void subscribeShouldThrowExceptionWhenStreamDoesNotExist() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits)) + Collections.emptyMap() + )) .isInstanceOf(StreamDoesNotExistException.class); } @@ -325,8 +344,8 @@ void subscribePropagateExceptionWhenClientSubscriptionFails() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits)) + Collections.emptyMap() + )) .isInstanceOf(StreamException.class) .hasMessage(exceptionMessage); assertThat(MonitoringTestUtils.extract(coordinator).isEmpty()).isTrue(); @@ -347,8 +366,8 @@ void subscribeShouldThrowExceptionWhenMetadataResponseIsNotOk() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits)) + Collections.emptyMap() + )) .isInstanceOf(IllegalStateException.class); } @@ -366,8 +385,8 @@ void subscribeShouldThrowExceptionIfNoNodeAvailableForStream() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits)) + Collections.emptyMap() + )) .isInstanceOf(IllegalStateException.class); } @@ -408,8 +427,8 @@ void subscribeShouldSubscribeToStreamAndDispatchMessage_UnsubscribeShouldUnsubsc () -> trackingClosingCallbackCalls.incrementAndGet(), (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -459,8 +478,8 @@ void subscribeShouldSubscribeToStreamAndDispatchMessageWithManySubscriptions() { (offset, message) -> messageHandlerCalls.compute(subId, (k, v) -> (v == null) ? 1 : ++v), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); closingRunnables.add(closingRunnable); } @@ -535,8 +554,8 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -555,8 +574,8 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(client, times(1 + 1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -636,8 +655,8 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -692,8 +711,8 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -707,8 +726,8 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(client, times(1 + 1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -796,8 +815,8 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -862,8 +881,8 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -916,8 +935,8 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -971,8 +990,8 @@ void shouldUseNewClientsForMoreThanMaxSubscriptionsAndCloseClientAfterUnsubscrip NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits)) + Collections.emptyMap() + )) .collect(Collectors.toList()); verify(clientFactory, times(2)).client(any()); @@ -1033,8 +1052,8 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); }); // the extra is allocated on another client from the same pool verify(clientFactory, times(2)).client(any()); @@ -1059,8 +1078,8 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(2 + 1)).client(any()); verify(client, times(subscriptionCount + ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT + 1)) @@ -1100,8 +1119,8 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); }); // the extra is allocated on another client from the same pool verify(clientFactory, times(2)).client(any()); @@ -1132,8 +1151,8 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); // no more client creation verify(clientFactory, times(2)).client(any()); @@ -1183,8 +1202,8 @@ void shouldRestartWhereItLeftOffAfterDisruption(Consumer {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1256,8 +1275,8 @@ void shouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesReceived( NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1331,8 +1350,8 @@ void shouldUseStoredOffsetOnRecovery(Consumer configur NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1418,8 +1437,8 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1433,8 +1452,8 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1 + 1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1499,8 +1518,8 @@ void shouldRetryAssignmentOnRecoveryStreamNotAvailableFailure() throws Exception NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1565,8 +1584,8 @@ void shouldRetryAssignmentOnRecoveryCandidateLookupFailure() throws Exception { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); @@ -1608,8 +1627,8 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); closingRunnable.run(); }; @@ -1696,8 +1715,8 @@ void shouldCallConsumerFlowControlHandlers(Consumer co NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, mockedConsumerFlowControlStrategyBuilder, - Collections.emptyMap(), - initialCredits); + Collections.emptyMap() + ); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), eq(numberOfInitialCreditsOnSubscribe), anyMap()); From f94be9597e0ae6874ea7d9b913a8a4e49c89a0c3 Mon Sep 17 00:00:00 2001 From: henry701 Date: Tue, 13 Jun 2023 20:06:38 -0300 Subject: [PATCH 06/14] Fix tests by fixing control flow callback adapter delegation --- ...lientCallbackStreamDataHandlerAdapter.java | 50 ++++++++++++++++++- .../stream/impl/ConsumersCoordinatorTest.java | 4 +- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java index 1bc320830d..9a25d9e574 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java +++ b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java @@ -4,6 +4,8 @@ import com.rabbitmq.stream.Message; import com.rabbitmq.stream.OffsetSpecification; +import java.util.Map; + public class ClientCallbackStreamDataHandlerAdapter implements CallbackStreamDataHandler, Client.PublishConfirmListener, @@ -21,44 +23,73 @@ public ClientCallbackStreamDataHandlerAdapter(CallbackStreamDataHandler callback this.callbackStreamDataHandler = callbackStreamDataHandler; } - @Override public void handle(byte publisherId, long publishingId) { this.handlePublishConfirm(publisherId, publishingId); } + @Override + public void handlePublishConfirm(byte publisherId, long publishingId) { + this.callbackStreamDataHandler.handlePublishConfirm(publisherId, publishingId); + } + @Override public void handle(byte publisherId, long publishingId, short errorCode) { this.handlePublishError(publisherId, publishingId, errorCode); } + @Override + public void handlePublishError(byte publisherId, long publishingId, short errorCode) { + this.callbackStreamDataHandler.handlePublishError(publisherId, publishingId, errorCode); + } + @Override public void handle(Client client, byte subscriptionId, long offset, long messageCount, long dataSize) { this.handleChunk(subscriptionId, offset, messageCount, dataSize); } + @Override + public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + this.callbackStreamDataHandler.handleChunk(subscriptionId, offset, messageCount, dataSize); + } + @Override public void handle(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { this.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); } + @Override + public void handleMessage(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { + this.callbackStreamDataHandler.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); + } + @Override public void handle(byte subscriptionId, short responseCode) { this.handleCreditNotification(subscriptionId, responseCode); } + @Override + public void handleCreditNotification(byte subscriptionId, short responseCode) { + this.callbackStreamDataHandler.handleCreditNotification(subscriptionId, responseCode); + } + @Override public OffsetSpecification handle(Client client, byte subscriptionId, boolean active) { this.handleConsumerUpdate(subscriptionId, active); return null; } + @Override + public void handleConsumerUpdate(byte subscriptionId, boolean active) { + this.callbackStreamDataHandler.handleConsumerUpdate(subscriptionId, active); + } + @Override public void handle(Client.ShutdownContext shutdownContext) { this.handleShutdown(shutdownContext); } - void handleShutdown(Client.ShutdownContext shutdownContext) { + public void handleShutdown(Client.ShutdownContext shutdownContext) { if(callbackStreamDataHandler instanceof AutoCloseable) { try { ((AutoCloseable) callbackStreamDataHandler).close(); @@ -73,4 +104,19 @@ public void handle(String stream, short code) { this.handleMetadata(stream, code); } + @Override + public void handleMetadata(String stream, short code) { + this.callbackStreamDataHandler.handleMetadata(stream, code); + } + + @Override + public void handleSubscribe(byte subscriptionId, String stream, OffsetSpecification offsetSpecification, Map subscriptionProperties) { + this.callbackStreamDataHandler.handleSubscribe(subscriptionId, stream, offsetSpecification, subscriptionProperties); + } + + @Override + public void handleUnsubscribe(byte subscriptionId) { + this.callbackStreamDataHandler.handleUnsubscribe(subscriptionId); + } + } diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index 614c73f6e0..1b4decff14 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -1721,7 +1721,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), eq(numberOfInitialCreditsOnSubscribe), anyMap()); verify(mockedConsumerFlowControlStrategy, times(1)) - .handleSubscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); + .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(0) .isEqualTo(OffsetSpecification.next()); @@ -1754,7 +1754,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); verify(mockedConsumerFlowControlStrategy, times(2)) - .handleSubscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); + .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(1) From b4a0b44934b631330255d77be121a51e59852156 Mon Sep 17 00:00:00 2001 From: henry701 Date: Thu, 15 Jun 2023 21:54:32 -0300 Subject: [PATCH 07/14] Add integrated test and fix issues discovered in the initial flow control implementation --- .../com/rabbitmq/stream/ConsumerBuilder.java | 2 +- .../flow/MessageHandlingListenerAware.java | 7 + .../impl/ConsumerStatisticRecorder.java | 18 ++- .../stream/impl/ConsumersCoordinator.java | 8 +- ...cRecordingConsumerFlowControlStrategy.java | 1 - ...umerFlowControlStrategyBuilderFactory.java | 29 ++++- .../stream/impl/StreamConsumerTest.java | 123 ++++++++++++++---- 7 files changed, 152 insertions(+), 36 deletions(-) create mode 100644 src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java diff --git a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java index 59201c12d3..eb8c14c776 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java @@ -66,7 +66,7 @@ public interface ConsumerBuilder { * Factory for the flow control strategy to be used when consuming messages. * @param consumerFlowControlStrategyBuilderFactory the factory * @return a fluent configurable builder for the flow control strategy - * @param + * @param The type of the builder for the provided factory */ > T flowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory); diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java new file mode 100644 index 0000000000..913718036d --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java @@ -0,0 +1,7 @@ +package com.rabbitmq.stream.flow; + +import com.rabbitmq.stream.ConsumerBuilder; + +public interface MessageHandlingListenerAware extends ConsumerBuilder.ConsumerBuilderAccessor { + MessageHandlingAware messageHandlingListener(); +} \ No newline at end of file diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java index 75addf3674..9cbdb19a74 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java @@ -17,6 +17,7 @@ import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; public class ConsumerStatisticRecorder implements CallbackStreamDataHandler { @@ -101,9 +102,9 @@ public void handleMessage( ); return null; } - NavigableMap subHeadMap = v.unprocessedChunksByOffset.headMap(offset, true); - Map.Entry lastEntry = subHeadMap.pollLastEntry(); - if(lastEntry == null) { + NavigableMap subHeadMapByOffset = v.unprocessedChunksByOffset.headMap(offset, true); + Map.Entry lastOffsetToChunkEntry = subHeadMapByOffset.lastEntry(); + if(lastOffsetToChunkEntry == null) { LOGGER.warn( "handleMessage called but chunk was not found! subscriptionId={} offset={}", subscriptionId, @@ -111,8 +112,8 @@ public void handleMessage( ); return v; } - ChunkStatistics statistics = lastEntry.getValue(); - statistics.unprocessedMessagesByOffset.put(offset, message); + ChunkStatistics chunkStatistics = lastOffsetToChunkEntry.getValue(); + chunkStatistics.unprocessedMessagesByOffset.put(offset, message); return v; } ); @@ -165,7 +166,8 @@ public boolean markHandled(AggregatedMessageStatistics aggregatedMessageStatisti return false; } // Remove chunk from list of unprocessed chunks if all its messages have been processed - if (aggregatedMessageStatistics.chunkStatistics.unprocessedMessagesByOffset.isEmpty()) { + aggregatedMessageStatistics.chunkStatistics.processedMessages.incrementAndGet(); + if (aggregatedMessageStatistics.chunkStatistics.isDone()) { aggregatedMessageStatistics.chunkHeadMap.remove(aggregatedMessageStatistics.messageEntry.getKey(), aggregatedMessageStatistics.chunkStatistics); } return true; @@ -299,6 +301,7 @@ public NavigableMap getUnprocessedChunksByOffset() { public static class ChunkStatistics { private final long offset; + private AtomicLong processedMessages = new AtomicLong(); private final long messageCount; private final long dataSize; private final Map unprocessedMessagesByOffset; @@ -330,6 +333,9 @@ public Map getUnprocessedMessagesByOffset() { return Collections.unmodifiableMap(unprocessedMessagesByOffset); } + public boolean isDone() { + return processedMessages.get() == messageCount && unprocessedMessagesByOffset.isEmpty(); + } } public Map> getStreamNameToSubscriptionIdMap() { diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 4ee632cdc3..5aa4b31030 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -565,6 +565,7 @@ private ClientSubscriptionsManager( this.consumerFlowControlStrategy = localConsumerFlowControlStrategy; CreditNotification creditNotification = (subscriptionId, responseCode) -> { + localConsumerFlowControlStrategy.handleCreditNotification(subscriptionId, responseCode); SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); String stream = subscriptionTracker == null ? "?" : subscriptionTracker.stream; @@ -573,11 +574,11 @@ private ClientSubscriptionsManager( subscriptionId & 0xFF, stream, Utils.formatConstant(responseCode)); - localConsumerFlowControlStrategy.handleCreditNotification(subscriptionId, responseCode); }; MessageListener messageListener = (subscriptionId, offset, chunkTimestamp, committedOffset, message) -> { + localConsumerFlowControlStrategy.handleMessage(subscriptionId, offset, chunkTimestamp, committedOffset, message); SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); if (subscriptionTracker != null) { @@ -595,7 +596,6 @@ private ClientSubscriptionsManager( this.id, this.name); } - localConsumerFlowControlStrategy.handleMessage(subscriptionId, offset, chunkTimestamp, committedOffset, message); }; ShutdownListener shutdownListener = shutdownContext -> { @@ -656,6 +656,7 @@ private ClientSubscriptionsManager( LOGGER.debug( "Received metadata notification for '{}', stream is likely to have become unavailable", stream); + localConsumerFlowControlStrategy.handleMetadata(stream, code); Set affectedSubscriptions; synchronized (this) { Set subscriptions = streamToStreamSubscriptions.remove(stream); @@ -701,10 +702,10 @@ private ClientSubscriptionsManager( "Consumers re-assignment after metadata update on stream '%s'", stream)); } - localConsumerFlowControlStrategy.handleMetadata(stream, code); }; ConsumerUpdateListener consumerUpdateListener = (client, subscriptionId, active) -> { + localConsumerFlowControlStrategy.handleConsumerUpdate(subscriptionId, active); OffsetSpecification result = null; SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); @@ -720,7 +721,6 @@ private ClientSubscriptionsManager( LOGGER.debug( "Could not find stream subscription {} for consumer update", subscriptionId); } - localConsumerFlowControlStrategy.handleConsumerUpdate(subscriptionId, active); return result; }; String connectionName = connectionNamingStrategy.apply(ClientConnectionType.CONSUMER); diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java index 992f970c17..112b46351b 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java @@ -42,7 +42,6 @@ public void handleSubscribe( OffsetSpecification offsetSpecification, Map subscriptionProperties ) { - super.handleSubscribe(subscriptionId, stream, offsetSpecification, subscriptionProperties); this.consumerStatisticRecorder.handleSubscribe( subscriptionId, stream, diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java index e426645548..2b9a0cd356 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java @@ -1,10 +1,16 @@ package com.rabbitmq.stream.impl.flow; import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; import com.rabbitmq.stream.flow.CreditAsker; +import com.rabbitmq.stream.flow.MessageHandlingAware; +import com.rabbitmq.stream.flow.MessageHandlingListenerAware; +import java.util.Collections; +import java.util.Set; +import java.util.WeakHashMap; import java.util.function.Supplier; public class MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { @@ -16,19 +22,26 @@ public Builder builder(ConsumerBuilder consumerBuilder) { return new Builder(consumerBuilder); } - public static class Builder implements ConsumerFlowControlStrategyBuilder { + public static class Builder implements ConsumerFlowControlStrategyBuilder, MessageHandlingListenerAware { private final ConsumerBuilder consumerBuilder; private int maximumInflightChunksPerSubscription = 1; + private final Set instances = Collections.newSetFromMap(new WeakHashMap<>()); + public Builder(ConsumerBuilder consumerBuilder) { this.consumerBuilder = consumerBuilder; } @Override public MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { - return new MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy(creditAskerSupplier, this.maximumInflightChunksPerSubscription); + MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy built = new MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy( + creditAskerSupplier, + this.maximumInflightChunksPerSubscription + ); + instances.add(built); + return built; } @Override @@ -41,6 +54,18 @@ public Builder maximumInflightChunksPerSubscription(int maximumInflightChunksPer return this; } + @Override + public MessageHandlingAware messageHandlingListener() { + return this::messageHandlingMulticaster; + } + + private boolean messageHandlingMulticaster(MessageHandler.Context context) { + boolean changed = false; + for(MessageHandlingAware instance : instances) { + changed = changed || instance.markHandled(context); + } + return changed; + } } } diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java index 070864a753..2bd59b6047 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java @@ -13,25 +13,37 @@ // info@rabbitmq.com. 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.localhost; -import static com.rabbitmq.stream.impl.TestUtils.streamName; -import static com.rabbitmq.stream.impl.TestUtils.waitAtMost; -import static com.rabbitmq.stream.impl.TestUtils.waitMs; -import static java.lang.String.format; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.Address; +import com.rabbitmq.stream.BackOffDelayPolicy; +import com.rabbitmq.stream.ConfirmationHandler; +import com.rabbitmq.stream.Consumer; +import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.Environment; +import com.rabbitmq.stream.EnvironmentBuilder; +import com.rabbitmq.stream.Host; +import com.rabbitmq.stream.MessageHandler; +import com.rabbitmq.stream.NoOffsetException; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.Producer; +import com.rabbitmq.stream.StreamDoesNotExistException; +import com.rabbitmq.stream.flow.MessageHandlingAware; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.MonitoringTestUtils.ConsumerInfo; -import com.rabbitmq.stream.impl.TestUtils.BrokerVersion; -import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast; -import com.rabbitmq.stream.impl.TestUtils.DisabledIfRabbitMqCtlNotSet; +import com.rabbitmq.stream.impl.TestUtils.*; +import com.rabbitmq.stream.impl.flow.MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory; import io.netty.channel.EventLoopGroup; +import org.assertj.core.api.ThrowableAssert.ThrowingCallable; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + import java.net.UnknownHostException; import java.time.Duration; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -48,14 +60,11 @@ import java.util.function.UnaryOperator; import java.util.stream.IntStream; import java.util.stream.Stream; -import org.assertj.core.api.ThrowableAssert.ThrowingCallable; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInfo; -import org.junit.jupiter.api.extension.ExtendWith; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.MethodSource; + +import static com.rabbitmq.stream.impl.TestUtils.*; +import static java.lang.String.format; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; @ExtendWith(TestUtils.StreamTestInfrastructureExtension.class) public class StreamConsumerTest { @@ -200,6 +209,76 @@ void consume() throws Exception { consumer.close(); } + @Test + void consumeWithAsyncConsumerFlowControl() throws Exception { + int messageCount = 100_000; + CountDownLatch publishLatch = new CountDownLatch(messageCount); + Client client = + cf.get( + new Client.ClientParameters() + .publishConfirmListener((publisherId, publishingId) -> publishLatch.countDown())); + + client.declarePublisher(b(1), null, stream); + IntStream.range(0, messageCount) + .forEach( + i -> + client.publish( + b(1), + Collections.singletonList( + client.messageBuilder().addData("".getBytes()).build()))); + + assertThat(publishLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + CountDownLatch firstConsumeLatch = new CountDownLatch(1); + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + + AtomicLong chunkTimestamp = new AtomicLong(); + + ConsumerBuilder consumerBuilder = environment.consumerBuilder().stream(stream) + .offset(OffsetSpecification.first()); + + MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.Builder flowControlStrategyBuilder = consumerBuilder + .flowControlStrategy(MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.INSTANCE) + .maximumInflightChunksPerSubscription(1); + MessageHandlingAware messageHandlingListener = flowControlStrategyBuilder.messageHandlingListener(); + + List messageContexts = new ArrayList<>(); + + AtomicBoolean shouldInstaConsume = new AtomicBoolean(false); + AtomicBoolean unhandledOnInstaConsume = new AtomicBoolean(false); + + consumerBuilder = flowControlStrategyBuilder + .builder() + .messageHandler( + (context, message) -> { + if(shouldInstaConsume.get()) { + if(!messageHandlingListener.markHandled(context)) { + unhandledOnInstaConsume.set(true); + } + } else { + messageContexts.add(context); + } + firstConsumeLatch.countDown(); + chunkTimestamp.set(context.timestamp()); + consumeLatch.countDown(); + }); + Consumer consumer = consumerBuilder.build(); + + assertThat(firstConsumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isFalse(); + assertThat(chunkTimestamp.get()).isNotZero(); + + shouldInstaConsume.set(true); + boolean allMarkedHandled = messageContexts.parallelStream().allMatch(messageHandlingListener::markHandled); + assertThat(allMarkedHandled).isTrue(); + + assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + assertThat(unhandledOnInstaConsume.get()).isFalse(); + + consumer.close(); + } + @Test void closeOnCondition() throws Exception { int messageCount = 50_000; From 13102e7c663e42033aad61512e6f3ed8a2718569 Mon Sep 17 00:00:00 2001 From: henry701 Date: Sat, 17 Jun 2023 00:55:15 -0300 Subject: [PATCH 08/14] Fix consumer flow on reconnection situations and add integrated test for it. Also add shortcuts on Consumer to avoid users having to use impl classes for flow control. Also, remove the verbose builder factories by using functional interfaces. --- .../stream/CallbackStreamDataHandler.java | 7 +- .../com/rabbitmq/stream/ConsumerBuilder.java | 27 ++- .../AsyncConsumerFlowControlStrategy.java | 9 + .../flow/ConsumerFlowControlStrategy.java | 21 +- ...umerFlowControlStrategyBuilderFactory.java | 1 + .../com/rabbitmq/stream/flow/CreditAsker.java | 1 + ...ware.java => MessageHandlingListener.java} | 2 +- .../flow/MessageHandlingListenerAware.java | 7 - ...ndlingListenerConsumerBuilderAccessor.java | 7 + ...lientCallbackStreamDataHandlerAdapter.java | 16 +- .../impl/ConsumerStatisticRecorder.java | 127 ++++++++++-- .../stream/impl/ConsumersCoordinator.java | 13 +- .../stream/impl/StreamConsumerBuilder.java | 33 ++- ...cRecordingConsumerFlowControlStrategy.java | 22 +- .../LegacyConsumerFlowControlStrategy.java | 40 ---- ...umerFlowControlStrategyBuilderFactory.java | 52 ----- ...tionAsyncConsumerFlowControlStrategy.java} | 72 ++++++- ...umerFlowControlStrategyBuilderFactory.java | 71 ------- ...ynchronousConsumerFlowControlStrategy.java | 83 ++++++++ .../stream/impl/ConsumersCoordinatorTest.java | 70 +++---- .../stream/impl/StreamConsumerTest.java | 194 +++++++++++++++++- 21 files changed, 605 insertions(+), 270 deletions(-) create mode 100644 src/main/java/com/rabbitmq/stream/flow/AsyncConsumerFlowControlStrategy.java rename src/main/java/com/rabbitmq/stream/flow/{MessageHandlingAware.java => MessageHandlingListener.java} (91%) delete mode 100644 src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java create mode 100644 src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java delete mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java delete mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java rename src/main/java/com/rabbitmq/stream/impl/flow/{MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java => MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java} (50%) delete mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java create mode 100644 src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java diff --git a/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java b/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java index b26cec38ad..e043a1ef1f 100644 --- a/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java +++ b/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java @@ -37,18 +37,21 @@ default void handleMetadata(String stream, short code) { } /** - * Callback for handling a new stream subscription. + * Callback for handling a stream subscription. * * @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 + * @param isInitialSubscription Whether this subscription is an initial subscription + * or a recovery for an existing subscription */ default void handleSubscribe( byte subscriptionId, String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties + Map subscriptionProperties, + boolean isInitialSubscription ) { // No-op by default } diff --git a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java index eb8c14c776..b18bb2e297 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java @@ -13,8 +13,11 @@ // info@rabbitmq.com. package com.rabbitmq.stream; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.flow.MessageHandlingListener; +import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import java.time.Duration; @@ -62,13 +65,35 @@ public interface ConsumerBuilder { */ ConsumerBuilder messageHandler(MessageHandler messageHandler); + /** + * Configure credit parameters for synchronous flow control. + * + * @param initial Credits to ask for with each new subscription + * @param onChunkDelivery Credits to ask for after a chunk is delivered + * @return this {@link ConsumerBuilder} + */ + ConsumerBuilder synchronousControlFlow(int initial, int onChunkDelivery); + + /** + * Configure credit parameters for asynchronous flow control. + * + * @param concurrencyLevel Maximum chunks to have in-processing at a given moment + * @return A {@link MessageHandlingListenerConsumerBuilderAccessor} for obtaining the {@link MessageHandlingListener} + * and navigating fluently back to this {@link ConsumerBuilder} + */ + MessageHandlingListenerConsumerBuilderAccessor asynchronousControlFlow(int concurrencyLevel); + /** * Factory for the flow control strategy to be used when consuming messages. + * When there is no need to use a custom strategy, which is the majority of cases, + * prefer using {@link ConsumerBuilder#synchronousControlFlow} or {@link ConsumerBuilder#asynchronousControlFlow} instead. + * * @param consumerFlowControlStrategyBuilderFactory the factory * @return a fluent configurable builder for the flow control strategy * @param The type of the builder for the provided factory */ - > T flowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory); + , S extends ConsumerFlowControlStrategy> + T customFlowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory); /** * The logical name of the {@link Consumer}. diff --git a/src/main/java/com/rabbitmq/stream/flow/AsyncConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/AsyncConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..14a27ae9b5 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/AsyncConsumerFlowControlStrategy.java @@ -0,0 +1,9 @@ +package com.rabbitmq.stream.flow; + +/** + * Variant of {@link ConsumerFlowControlStrategy} that implements {@link MessageHandlingListener} to asynchronously + * mark messages as handled. + */ +public interface AsyncConsumerFlowControlStrategy extends ConsumerFlowControlStrategy, MessageHandlingListener { + +} diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java index 096ed95bd1..c9e1dba68e 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java @@ -13,32 +13,41 @@ public interface ConsumerFlowControlStrategy extends CallbackStreamDataHandler { /** - * Callback for handling a new stream subscription. - * Called right before the subscription is sent to the actual client. + * Callback for handling a stream subscription. + * Called right before the subscription is sent to the broker. *

- * Either this variant or {@link CallbackStreamDataHandler#handleSubscribe(byte, String, OffsetSpecification, Map)} should be called, NOT both. + * Either this variant or {@link CallbackStreamDataHandler#handleSubscribe} should be called, NOT both. *

* * @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 + * @param isInitialSubscription Whether this subscription is an initial subscription + * or a recovery for an existing subscription * @return The initial credits that should be granted to this new subscription */ int handleSubscribeReturningInitialCredits( byte subscriptionId, String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties + Map subscriptionProperties, + boolean isInitialSubscription ); @Override - default void handleSubscribe(byte subscriptionId, String stream, OffsetSpecification offsetSpecification, Map subscriptionProperties) { + default void handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties, + boolean isInitialSubscription) { handleSubscribeReturningInitialCredits( subscriptionId, stream, offsetSpecification, - subscriptionProperties + subscriptionProperties, + isInitialSubscription ); } diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java index 766257dafd..28450f54bd 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilderFactory.java @@ -7,6 +7,7 @@ * @param the type of {@link ConsumerFlowControlStrategy} to be built * @param the type of fluent builder exposed by this factory. Must subclass {@link ConsumerFlowControlStrategyBuilder} */ +@FunctionalInterface public interface ConsumerFlowControlStrategyBuilderFactory> { /** * Accessor for configuration builder with settings specific to each implementing strategy diff --git a/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java b/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java index d367451173..0a906597ad 100644 --- a/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java +++ b/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java @@ -1,5 +1,6 @@ package com.rabbitmq.stream.flow; +@FunctionalInterface public interface CreditAsker { /** diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java similarity index 91% rename from src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java rename to src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java index f21d370f26..39edb6a119 100644 --- a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingAware.java +++ b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java @@ -2,7 +2,7 @@ import com.rabbitmq.stream.MessageHandler; -public interface MessageHandlingAware { +public interface MessageHandlingListener { /** * Marks a message as handled diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java deleted file mode 100644 index 913718036d..0000000000 --- a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerAware.java +++ /dev/null @@ -1,7 +0,0 @@ -package com.rabbitmq.stream.flow; - -import com.rabbitmq.stream.ConsumerBuilder; - -public interface MessageHandlingListenerAware extends ConsumerBuilder.ConsumerBuilderAccessor { - MessageHandlingAware messageHandlingListener(); -} \ No newline at end of file diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java new file mode 100644 index 0000000000..6c14552e65 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java @@ -0,0 +1,7 @@ +package com.rabbitmq.stream.flow; + +import com.rabbitmq.stream.ConsumerBuilder; + +public interface MessageHandlingListenerConsumerBuilderAccessor extends ConsumerBuilder.ConsumerBuilderAccessor { + MessageHandlingListener messageHandlingListener(); +} \ No newline at end of file diff --git a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java index 9a25d9e574..6f7268e9fb 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java +++ b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java @@ -110,8 +110,20 @@ public void handleMetadata(String stream, short code) { } @Override - public void handleSubscribe(byte subscriptionId, String stream, OffsetSpecification offsetSpecification, Map subscriptionProperties) { - this.callbackStreamDataHandler.handleSubscribe(subscriptionId, stream, offsetSpecification, subscriptionProperties); + public void handleSubscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties, + boolean isInitialSubscription + ) { + this.callbackStreamDataHandler.handleSubscribe( + subscriptionId, + stream, + offsetSpecification, + subscriptionProperties, + isInitialSubscription + ); } @Override diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java index 9cbdb19a74..d12e53379b 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java @@ -4,12 +4,14 @@ import com.rabbitmq.stream.Message; import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.flow.MessageHandlingListener; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; import javax.annotation.Nullable; import java.util.Collections; +import java.util.Iterator; import java.util.Map; import java.util.NavigableMap; import java.util.Set; @@ -19,7 +21,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; -public class ConsumerStatisticRecorder implements CallbackStreamDataHandler { +public class ConsumerStatisticRecorder implements CallbackStreamDataHandler, MessageHandlingListener { private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerStatisticRecorder.class); @@ -31,7 +33,8 @@ public void handleSubscribe( byte subscriptionId, String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties + Map subscriptionProperties, + boolean isInitialSubscription ) { this.streamNameToSubscriptionIdMap.compute( stream, @@ -40,12 +43,13 @@ public void handleSubscribe( v = Collections.newSetFromMap(new ConcurrentHashMap<>()); } boolean isNewElement = v.add(subscriptionId); - if(!isNewElement) { + if(isInitialSubscription && !isNewElement) { LOGGER.warn( - "handleSubscribed called for stream that already had same associated subscription! " + - "subscriptionId={}, stream={}", + "handleSubscribe called for stream that already had same associated subscription! " + + "subscriptionId={} stream={} offsetSpecification={}", subscriptionId, - stream + stream, + offsetSpecification ); } return v; @@ -54,17 +58,60 @@ public void handleSubscribe( this.subscriptionStatisticsMap.compute( subscriptionId, (k, v) -> { - if(v != null) { + if(v != null && isInitialSubscription) { LOGGER.warn( - "handleSubscribed called for subscription that already exists! subscriptionId={}", - subscriptionId + "handleSubscribe called for subscription that already exists! " + + "subscriptionId={} stream={} offsetSpecification={}", + subscriptionId, + stream, + offsetSpecification + ); + } + // Only overwrite if is a de-facto initial subscription + if(v == null) { + return new SubscriptionStatistics( + subscriptionId, + stream, + offsetSpecification, + subscriptionProperties ); } - return new SubscriptionStatistics(subscriptionId, stream, subscriptionProperties); + v.offsetSpecification = offsetSpecification; + v.pendingChunks.set(0); + v.subscriptionProperties = subscriptionProperties; + cleanupOldTrackingData(v); + return v; } ); } + private static void cleanupOldTrackingData(SubscriptionStatistics subscriptionStatistics) { + if (!subscriptionStatistics.offsetSpecification.isOffset()) { + LOGGER.debug("Can't cleanup old tracking data: offsetSpecification is not an offset! {}", subscriptionStatistics.offsetSpecification); + return; + } + // Mark messages before the initial offset as handled + long newSubscriptionInitialOffset = subscriptionStatistics.offsetSpecification.getOffset(); + NavigableMap chunksHeadMap = subscriptionStatistics.unprocessedChunksByOffset.headMap(newSubscriptionInitialOffset, false); + Iterator> chunksHeadMapEntryIterator = chunksHeadMap.entrySet().iterator(); + while(chunksHeadMapEntryIterator.hasNext()) { + Map.Entry chunksHeadMapEntry = chunksHeadMapEntryIterator.next(); + ChunkStatistics chunkStatistics = chunksHeadMapEntry.getValue(); + Iterator> chunkMessagesIterator = chunkStatistics.unprocessedMessagesByOffset.entrySet().iterator(); + while(chunkMessagesIterator.hasNext()) { + Map.Entry chunkMessageEntry = chunkMessagesIterator.next(); + long messageOffset = chunkMessageEntry.getKey(); + if(messageOffset < newSubscriptionInitialOffset) { + chunkMessagesIterator.remove(); + chunkStatistics.processedMessages.incrementAndGet(); + } + } + if(chunkStatistics.isDone()) { + chunksHeadMapEntryIterator.remove(); + } + } + } + @Override public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { this.subscriptionStatisticsMap.compute( @@ -72,8 +119,9 @@ public void handleChunk(byte subscriptionId, long offset, long messageCount, lon (k, v) -> { if(v == null) { LOGGER.warn( - "handleChunk called for subscription that does not exist! subscriptionId={}", - subscriptionId + "handleChunk called for subscription that does not exist! subscriptionId={} offset={}", + subscriptionId, + offset ); return null; } @@ -97,8 +145,9 @@ public void handleMessage( (k, v) -> { if(v == null) { LOGGER.warn( - "handleMessage called for subscription that does not exist! subscriptionId={}", - subscriptionId + "handleMessage called for subscription that does not exist! subscriptionId={} offset={}", + subscriptionId, + offset ); return null; } @@ -121,13 +170,32 @@ public void handleMessage( @Override public void handleUnsubscribe(byte subscriptionId) { - Object removed = this.subscriptionStatisticsMap.remove(subscriptionId); - if(removed == null) { + ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStatistics = this.subscriptionStatisticsMap.remove(subscriptionId); + if(subscriptionStatistics == null) { LOGGER.warn( "handleUnsubscribe called for subscriptionId that does not exist! subscriptionId={}", subscriptionId ); + return; } + this.streamNameToSubscriptionIdMap.compute(subscriptionStatistics.stream, (k, v) -> { + if(v == null) { + LOGGER.warn( + "handleUnsubscribe called and stream name '{}' did not contain subscriptions!", + subscriptionStatistics.stream + ); + return null; + } + boolean removed = v.remove(subscriptionId); + if(!removed) { + LOGGER.warn( + "handleUnsubscribe called and stream name '{}' did not contain subscriptionId {}!", + subscriptionStatistics.stream, + subscriptionId + ); + } + return v.isEmpty() ? null : v; + }); } /** @@ -137,6 +205,7 @@ public void handleUnsubscribe(byte subscriptionId) { * @return Whether the message was marked as handled (returning {@code true}) * or was not found (either because it was already marked as handled, or wasn't tracked) */ + @Override public boolean markHandled(MessageHandler.Context messageContext) { AggregatedMessageStatistics entry = retrieveStatistics(messageContext); if (entry == null) { @@ -159,6 +228,14 @@ public boolean markHandled(AggregatedMessageStatistics aggregatedMessageStatisti || aggregatedMessageStatistics.chunkHeadMap == null) { return false; } + if(aggregatedMessageStatistics.subscriptionStatistics.offsetSpecification.isOffset()) { + long initialOffset = aggregatedMessageStatistics.subscriptionStatistics.offsetSpecification.getOffset(); + // Old tracked message, should already be handled, probably a late acknowledgment of a defunct connection. + if(aggregatedMessageStatistics.offset < initialOffset) { + LOGGER.debug("Old message registered as consumed. Message Offset: {}, Start Offset: {}", aggregatedMessageStatistics.offset, initialOffset); + return true; + } + } Message removedMessage = aggregatedMessageStatistics.chunkStatistics .unprocessedMessagesByOffset .remove(aggregatedMessageStatistics.offset); @@ -257,21 +334,29 @@ public static class SubscriptionStatistics { private final byte subscriptionId; private final String stream; private final AtomicInteger pendingChunks = new AtomicInteger(0); - private final Map subscriptionProperties; + private OffsetSpecification offsetSpecification; + private Map subscriptionProperties; private final NavigableMap unprocessedChunksByOffset; - public SubscriptionStatistics(byte subscriptionId, String stream, Map subscriptionProperties) { - this(subscriptionId, stream, subscriptionProperties, new ConcurrentSkipListMap<>()); + public SubscriptionStatistics( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties + ) { + this(subscriptionId, stream, offsetSpecification, subscriptionProperties, new ConcurrentSkipListMap<>()); } public SubscriptionStatistics( byte subscriptionId, String stream, + OffsetSpecification offsetSpecification, Map subscriptionProperties, NavigableMap unprocessedChunksByOffset ) { this.subscriptionId = subscriptionId; this.stream = stream; + this.offsetSpecification = offsetSpecification; this.subscriptionProperties = subscriptionProperties; this.unprocessedChunksByOffset = unprocessedChunksByOffset; } @@ -288,6 +373,10 @@ public AtomicInteger getPendingChunks() { return pendingChunks; } + public OffsetSpecification getOffsetSpecification() { + return offsetSpecification; + } + public Map getSubscriptionProperties() { return Collections.unmodifiableMap(subscriptionProperties); } diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 5aa4b31030..015916561f 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -952,18 +952,19 @@ synchronized void add( subscriptionContext.offsetSpecification()); checkNotClosed(); - byte subId = subscriptionId; int initialCredits = this.consumerFlowControlStrategy.handleSubscribeReturningInitialCredits( - subId, - subscriptionTracker.stream, - subscriptionContext.offsetSpecification(), - subscriptionTracker.subscriptionProperties + subscriptionId, + subscriptionTracker.stream, + subscriptionContext.offsetSpecification(), + subscriptionTracker.subscriptionProperties, + isInitialSubscription ); + final byte finalSubscriptionId = subscriptionId; Client.Response subscribeResponse = Utils.callAndMaybeRetry( () -> client.subscribe( - subId, + finalSubscriptionId, subscriptionTracker.stream, subscriptionContext.offsetSpecification(), initialCredits, diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java index ac8f81413d..fb2133b71a 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java @@ -20,9 +20,12 @@ import com.rabbitmq.stream.OffsetSpecification; import com.rabbitmq.stream.StreamException; import com.rabbitmq.stream.SubscriptionListener; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; -import com.rabbitmq.stream.impl.flow.LegacyConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; +import com.rabbitmq.stream.impl.flow.MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy; +import com.rabbitmq.stream.impl.flow.SynchronousConsumerFlowControlStrategy; import java.lang.reflect.Field; import java.lang.reflect.Modifier; @@ -46,7 +49,7 @@ class StreamConsumerBuilder implements ConsumerBuilder { private SubscriptionListener subscriptionListener = subscriptionContext -> {}; private final Map subscriptionProperties = new ConcurrentHashMap<>(); private ConsumerUpdateListener consumerUpdateListener; - private ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder = LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(this); + private ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder = SynchronousConsumerFlowControlStrategy.builder(this); public StreamConsumerBuilder(StreamEnvironment environment) { this.environment = environment; @@ -77,7 +80,9 @@ public ConsumerBuilder messageHandler(MessageHandler messageHandler) { } @Override - public > T flowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory) { + public + , S extends ConsumerFlowControlStrategy> + T customFlowControlStrategy(ConsumerFlowControlStrategyBuilderFactory consumerFlowControlStrategyBuilderFactory) { T localConsumerFlowControlStrategyBuilder = consumerFlowControlStrategyBuilderFactory.builder(this); this.consumerFlowControlStrategyBuilder = localConsumerFlowControlStrategyBuilder; return localConsumerFlowControlStrategyBuilder; @@ -143,25 +148,37 @@ public ConsumerBuilder noTrackingStrategy() { } /** + * Configure credit parameters for flow control. + * Implies usage of a traditional {@link SynchronousConsumerFlowControlStrategy}. * * @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} - * to define flow control strategies instead. */ - @Deprecated - public StreamConsumerBuilder credits(int initial, int onChunkDelivery) { + @Override + public StreamConsumerBuilder synchronousControlFlow(int initial, int onChunkDelivery) { if (initial <= 0 || onChunkDelivery <= 0) { throw new IllegalArgumentException("Credits must be positive"); } - this.consumerFlowControlStrategyBuilder = LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE + this.consumerFlowControlStrategyBuilder = SynchronousConsumerFlowControlStrategy .builder(this) .initialCredits(initial) .additionalCredits(onChunkDelivery); return this; } + @Override + public MessageHandlingListenerConsumerBuilderAccessor asynchronousControlFlow(int concurrencyLevel) { + if (concurrencyLevel <= 0) { + throw new IllegalArgumentException("ConcurrencyLevel must be positive"); + } + MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.Builder localBuilder = MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy + .builder(this) + .maximumInflightChunksPerSubscription(concurrencyLevel); + this.consumerFlowControlStrategyBuilder = localBuilder; + return localBuilder; + } + StreamConsumerBuilder lazyInit(boolean lazyInit) { this.lazyInit = lazyInit; return this; diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java index 112b46351b..17089b2198 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java @@ -4,10 +4,12 @@ import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.OffsetSpecification; import com.rabbitmq.stream.flow.AbstractConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.AsyncConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.CreditAsker; -import com.rabbitmq.stream.flow.MessageHandlingAware; import com.rabbitmq.stream.impl.ConsumerStatisticRecorder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; @@ -20,7 +22,9 @@ */ public abstract class AbstractStatisticRecordingConsumerFlowControlStrategy extends AbstractConsumerFlowControlStrategy - implements MessageHandlingAware { + implements AsyncConsumerFlowControlStrategy { + + private static final Logger LOGGER = LoggerFactory.getLogger(AbstractStatisticRecordingConsumerFlowControlStrategy.class); protected final ConsumerStatisticRecorder consumerStatisticRecorder = new ConsumerStatisticRecorder(); @@ -40,13 +44,15 @@ public void handleSubscribe( byte subscriptionId, String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties + Map subscriptionProperties, + boolean isInitialSubscription ) { this.consumerStatisticRecorder.handleSubscribe( subscriptionId, stream, offsetSpecification, - subscriptionProperties + subscriptionProperties, + isInitialSubscription ); } @@ -96,6 +102,10 @@ protected int registerCredits(byte subscriptionId, IntUnaryOperator askedToAsk, ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStatistics = this.consumerStatisticRecorder .getSubscriptionStatisticsMap() .get(subscriptionId); + if(subscriptionStatistics == null) { + LOGGER.warn("Lost subscription {}, returning no credits. askForCredits={}", subscriptionId, askForCredits); + return 0; + } subscriptionStatistics.getPendingChunks().updateAndGet(credits -> { int creditsToAsk = askedToAsk.applyAsInt(credits); outerCreditsToAsk.set(creditsToAsk); @@ -103,8 +113,10 @@ protected int registerCredits(byte subscriptionId, IntUnaryOperator askedToAsk, }); int finalCreditsToAsk = outerCreditsToAsk.get(); if(askForCredits && finalCreditsToAsk > 0) { + LOGGER.debug("Asking for {} credits for subscriptionId {}", finalCreditsToAsk, subscriptionId); mandatoryCreditAsker().credit(subscriptionId, finalCreditsToAsk); } + LOGGER.debug("Returning {} credits for subscriptionId {} with askForCredits={}", finalCreditsToAsk, subscriptionId, askForCredits); return finalCreditsToAsk; } @@ -113,10 +125,12 @@ public boolean markHandled(MessageHandler.Context messageContext) { ConsumerStatisticRecorder.AggregatedMessageStatistics messageStatistics = this.consumerStatisticRecorder .retrieveStatistics(messageContext); if(messageStatistics == null) { + LOGGER.warn("Message statistics not found for offset {} on stream '{}'", messageContext.offset(), messageContext.stream()); return false; } boolean markedAsHandled = this.consumerStatisticRecorder.markHandled(messageStatistics); if(!markedAsHandled) { + LOGGER.warn("Message not marked as handled for offset {} on stream '{}'", messageContext.offset(), messageContext.stream()); return false; } afterMarkHandledStateChanged(messageContext, messageStatistics); diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java deleted file mode 100644 index f92c25fb24..0000000000 --- a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategy.java +++ /dev/null @@ -1,40 +0,0 @@ -package com.rabbitmq.stream.impl.flow; - -import com.rabbitmq.stream.OffsetSpecification; -import com.rabbitmq.stream.flow.AbstractConsumerFlowControlStrategy; -import com.rabbitmq.stream.flow.CreditAsker; - -import java.util.Map; -import java.util.function.Supplier; - -/** - * The flow control strategy that was always applied before the flow control strategy mechanism existed in the codebase. - * Requests a set amount of credits after each chunk arrives. - */ -public class LegacyConsumerFlowControlStrategy extends AbstractConsumerFlowControlStrategy { - - private final int initialCredits; - private final int additionalCredits; - - public LegacyConsumerFlowControlStrategy(Supplier creditAskerSupplier, int initialCredits, int additionalCredits) { - super(creditAskerSupplier); - this.initialCredits = initialCredits; - this.additionalCredits = additionalCredits; - } - - @Override - public int handleSubscribeReturningInitialCredits( - byte subscriptionId, - String stream, - OffsetSpecification offsetSpecification, - Map subscriptionProperties - ) { - return this.initialCredits; - } - - @Override - public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - mandatoryCreditAsker().credit(subscriptionId, this.additionalCredits); - } - -} diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java deleted file mode 100644 index 5ca5aaaa43..0000000000 --- a/src/main/java/com/rabbitmq/stream/impl/flow/LegacyConsumerFlowControlStrategyBuilderFactory.java +++ /dev/null @@ -1,52 +0,0 @@ -package com.rabbitmq.stream.impl.flow; - -import com.rabbitmq.stream.ConsumerBuilder; -import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; -import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; -import com.rabbitmq.stream.flow.CreditAsker; - -import java.util.function.Supplier; - -public class LegacyConsumerFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { - - public static final LegacyConsumerFlowControlStrategyBuilderFactory INSTANCE = new LegacyConsumerFlowControlStrategyBuilderFactory(); - - @Override - public Builder builder(ConsumerBuilder consumerBuilder) { - return new Builder(consumerBuilder); - } - - public static class Builder implements ConsumerFlowControlStrategyBuilder { - - private final ConsumerBuilder consumerBuilder; - - private int initialCredits = 1; - - private int additionalCredits = 1; - - public Builder(ConsumerBuilder consumerBuilder) { - this.consumerBuilder = consumerBuilder; - } - - @Override - public LegacyConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { - return new LegacyConsumerFlowControlStrategy(creditAskerSupplier, this.initialCredits, this.additionalCredits); - } - - @Override - public ConsumerBuilder builder() { - return this.consumerBuilder; - } - - public Builder additionalCredits(int additionalCredits) { - this.additionalCredits = additionalCredits; - return this; - } - - public Builder initialCredits(int initialCredits) { - this.initialCredits = initialCredits; - return this; - } - } - -} diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java similarity index 50% rename from src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java rename to src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java index 16404c2d1a..e5b3c41408 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java @@ -1,28 +1,34 @@ package com.rabbitmq.stream.impl.flow; +import com.rabbitmq.stream.ConsumerBuilder; import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.CreditAsker; -import com.rabbitmq.stream.flow.MessageHandlingAware; +import com.rabbitmq.stream.flow.MessageHandlingListener; +import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import com.rabbitmq.stream.impl.ConsumerStatisticRecorder; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collections; import java.util.Map; +import java.util.Set; +import java.util.WeakHashMap; import java.util.function.IntUnaryOperator; import java.util.function.Supplier; /** * A flow control strategy that enforces a maximum amount of Inflight chunks per registered subscription. - * Based on {@link MessageHandlingAware message acknowledgement}, asking for the maximum number of chunks possible, given the limit. + * Based on {@link MessageHandlingListener message acknowledgement}, asking for the maximum number of chunks possible, given the limit. */ -public class MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy extends AbstractStatisticRecordingConsumerFlowControlStrategy { +public class MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy extends AbstractStatisticRecordingConsumerFlowControlStrategy { - private static final Logger LOGGER = LoggerFactory.getLogger(MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy.class); + private static final Logger LOGGER = LoggerFactory.getLogger(MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.class); private final int maximumSimultaneousChunksPerSubscription; - public MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy( + public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( Supplier creditAskerSupplier, int maximumSimultaneousChunksPerSubscription ) { @@ -40,12 +46,14 @@ public int handleSubscribeReturningInitialCredits( byte subscriptionId, String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties) { + Map subscriptionProperties, + boolean isInitialSubscription) { this.handleSubscribe( subscriptionId, stream, offsetSpecification, - subscriptionProperties + subscriptionProperties, + isInitialSubscription ); return registerCredits(subscriptionId, getCreditAsker(subscriptionId), false); } @@ -79,4 +87,54 @@ private int extractInProcessingChunks(byte subscriptionId) { return inProcessingChunks; } + public static Builder builder(ConsumerBuilder consumerBuilder) { + return new Builder(consumerBuilder); + } + + public static class Builder implements ConsumerFlowControlStrategyBuilder, MessageHandlingListenerConsumerBuilderAccessor { + + private final ConsumerBuilder consumerBuilder; + + private int maximumInflightChunksPerSubscription = 1; + + private final Set instances = Collections.newSetFromMap(new WeakHashMap<>()); + + public Builder(ConsumerBuilder consumerBuilder) { + this.consumerBuilder = consumerBuilder; + } + + @Override + public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { + MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy built = new MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( + creditAskerSupplier, + this.maximumInflightChunksPerSubscription + ); + instances.add(built); + return built; + } + + @Override + public ConsumerBuilder builder() { + return this.consumerBuilder; + } + + public Builder maximumInflightChunksPerSubscription(int maximumInflightChunksPerSubscription) { + this.maximumInflightChunksPerSubscription = maximumInflightChunksPerSubscription; + return this; + } + + @Override + public MessageHandlingListener messageHandlingListener() { + return this::messageHandlingMulticaster; + } + + private boolean messageHandlingMulticaster(MessageHandler.Context context) { + boolean changed = false; + for(MessageHandlingListener instance : instances) { + changed = changed || instance.markHandled(context); + } + return changed; + } + } + } diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java deleted file mode 100644 index 2b9a0cd356..0000000000 --- a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.java +++ /dev/null @@ -1,71 +0,0 @@ -package com.rabbitmq.stream.impl.flow; - -import com.rabbitmq.stream.ConsumerBuilder; -import com.rabbitmq.stream.MessageHandler; -import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; -import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; -import com.rabbitmq.stream.flow.CreditAsker; -import com.rabbitmq.stream.flow.MessageHandlingAware; -import com.rabbitmq.stream.flow.MessageHandlingListenerAware; - -import java.util.Collections; -import java.util.Set; -import java.util.WeakHashMap; -import java.util.function.Supplier; - -public class MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory implements ConsumerFlowControlStrategyBuilderFactory { - - public static final MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory INSTANCE = new MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory(); - - @Override - public Builder builder(ConsumerBuilder consumerBuilder) { - return new Builder(consumerBuilder); - } - - public static class Builder implements ConsumerFlowControlStrategyBuilder, MessageHandlingListenerAware { - - private final ConsumerBuilder consumerBuilder; - - private int maximumInflightChunksPerSubscription = 1; - - private final Set instances = Collections.newSetFromMap(new WeakHashMap<>()); - - public Builder(ConsumerBuilder consumerBuilder) { - this.consumerBuilder = consumerBuilder; - } - - @Override - public MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { - MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy built = new MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategy( - creditAskerSupplier, - this.maximumInflightChunksPerSubscription - ); - instances.add(built); - return built; - } - - @Override - public ConsumerBuilder builder() { - return this.consumerBuilder; - } - - public Builder maximumInflightChunksPerSubscription(int maximumInflightChunksPerSubscription) { - this.maximumInflightChunksPerSubscription = maximumInflightChunksPerSubscription; - return this; - } - - @Override - public MessageHandlingAware messageHandlingListener() { - return this::messageHandlingMulticaster; - } - - private boolean messageHandlingMulticaster(MessageHandler.Context context) { - boolean changed = false; - for(MessageHandlingAware instance : instances) { - changed = changed || instance.markHandled(context); - } - return changed; - } - } - -} diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java new file mode 100644 index 0000000000..f652f6de99 --- /dev/null +++ b/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java @@ -0,0 +1,83 @@ +package com.rabbitmq.stream.impl.flow; + +import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.flow.AbstractConsumerFlowControlStrategy; +import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.flow.CreditAsker; + +import java.util.Map; +import java.util.function.Supplier; + +/** + * The default flow control strategy. + * Requests a set amount of credits after each chunk arrives. + * Ideal for usage when the message is consumed synchronously inside the message handler, + * which is the case for most Consumers. + */ +public class SynchronousConsumerFlowControlStrategy extends AbstractConsumerFlowControlStrategy { + + private final int initialCredits; + private final int additionalCredits; + + public SynchronousConsumerFlowControlStrategy(Supplier creditAskerSupplier, int initialCredits, int additionalCredits) { + super(creditAskerSupplier); + this.initialCredits = initialCredits; + this.additionalCredits = additionalCredits; + } + + @Override + public int handleSubscribeReturningInitialCredits( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + Map subscriptionProperties, + boolean isInitialSubscription + ) { + return this.initialCredits; + } + + @Override + public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + mandatoryCreditAsker().credit(subscriptionId, this.additionalCredits); + } + + public static SynchronousConsumerFlowControlStrategy.Builder builder(ConsumerBuilder consumerBuilder) { + return new SynchronousConsumerFlowControlStrategy.Builder(consumerBuilder); + } + + public static class Builder implements ConsumerFlowControlStrategyBuilder { + + private final ConsumerBuilder consumerBuilder; + + private int initialCredits = 1; + + private int additionalCredits = 1; + + public Builder(ConsumerBuilder consumerBuilder) { + this.consumerBuilder = consumerBuilder; + } + + @Override + public SynchronousConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { + return new SynchronousConsumerFlowControlStrategy(creditAskerSupplier, this.initialCredits, this.additionalCredits); + } + + @Override + public ConsumerBuilder builder() { + return this.consumerBuilder; + } + + public Builder additionalCredits(int additionalCredits) { + this.additionalCredits = additionalCredits; + return this; + } + + public Builder initialCredits(int initialCredits) { + this.initialCredits = initialCredits; + return this; + } + + } + +} diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index 1b4decff14..97396f2189 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -28,7 +28,7 @@ import com.rabbitmq.stream.impl.Client.Response; import com.rabbitmq.stream.impl.MonitoringTestUtils.ConsumerCoordinatorInfo; import com.rabbitmq.stream.impl.Utils.ClientFactory; -import com.rabbitmq.stream.impl.flow.LegacyConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.impl.flow.SynchronousConsumerFlowControlStrategy; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -198,7 +198,7 @@ void tearDown() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(2)).client(any()); @@ -240,7 +240,7 @@ void shouldGetExactNodeImmediatelyWithAdvertisedHostNameClientFactoryAndExactNod NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -271,7 +271,7 @@ void shouldSubscribeWithEmptyPropertiesWithUnamedConsumer() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -293,7 +293,7 @@ void subscribeShouldThrowExceptionWhenNoMetadataForTheStream() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() )) .isInstanceOf(StreamDoesNotExistException.class); @@ -313,7 +313,7 @@ void subscribeShouldThrowExceptionWhenStreamDoesNotExist() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() )) .isInstanceOf(StreamDoesNotExistException.class); @@ -343,7 +343,7 @@ void subscribePropagateExceptionWhenClientSubscriptionFails() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() )) .isInstanceOf(StreamException.class) @@ -365,7 +365,7 @@ void subscribeShouldThrowExceptionWhenMetadataResponseIsNotOk() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() )) .isInstanceOf(IllegalStateException.class); @@ -384,7 +384,7 @@ void subscribeShouldThrowExceptionIfNoNodeAvailableForStream() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() )) .isInstanceOf(IllegalStateException.class); @@ -426,7 +426,7 @@ void subscribeShouldSubscribeToStreamAndDispatchMessage_UnsubscribeShouldUnsubsc NO_OP_SUBSCRIPTION_LISTENER, () -> trackingClosingCallbackCalls.incrementAndGet(), (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -477,7 +477,7 @@ void subscribeShouldSubscribeToStreamAndDispatchMessageWithManySubscriptions() { NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.compute(subId, (k, v) -> (v == null) ? 1 : ++v), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); closingRunnables.add(closingRunnable); @@ -553,7 +553,7 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -573,7 +573,7 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); @@ -654,7 +654,7 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -710,7 +710,7 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -725,7 +725,7 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); @@ -814,7 +814,7 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -880,7 +880,7 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -934,7 +934,7 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> messageHandlerCalls.incrementAndGet(), - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -989,7 +989,7 @@ void shouldUseNewClientsForMoreThanMaxSubscriptionsAndCloseClientAfterUnsubscrip NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() )) .collect(Collectors.toList()); @@ -1051,7 +1051,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); }); @@ -1077,7 +1077,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); @@ -1118,7 +1118,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); }); @@ -1150,7 +1150,7 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); @@ -1201,7 +1201,7 @@ void shouldRestartWhereItLeftOffAfterDisruption(Consumer {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1274,7 +1274,7 @@ void shouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesReceived( NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1349,7 +1349,7 @@ void shouldUseStoredOffsetOnRecovery(Consumer configur NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1436,7 +1436,7 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1451,7 +1451,7 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1517,7 +1517,7 @@ void shouldRetryAssignmentOnRecoveryStreamNotAvailableFailure() throws Exception NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1583,7 +1583,7 @@ void shouldRetryAssignmentOnRecoveryCandidateLookupFailure() throws Exception { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); verify(clientFactory, times(1)).client(any()); @@ -1626,7 +1626,7 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() { NO_OP_SUBSCRIPTION_LISTENER, NO_OP_TRACKING_CLOSING_CALLBACK, (offset, message) -> {}, - LegacyConsumerFlowControlStrategyBuilderFactory.INSTANCE.builder(null), + SynchronousConsumerFlowControlStrategy.builder(null), Collections.emptyMap() ); @@ -1699,7 +1699,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co int numberOfInitialCreditsOnSubscribe = 7; - when(mockedConsumerFlowControlStrategy.handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(), anyMap())) + when(mockedConsumerFlowControlStrategy.handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(), anyMap(), anyBoolean())) .thenReturn(numberOfInitialCreditsOnSubscribe); ConsumerFlowControlStrategyBuilder mockedConsumerFlowControlStrategyBuilder = Mockito.mock(ConsumerFlowControlStrategyBuilder.class); @@ -1721,7 +1721,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), eq(numberOfInitialCreditsOnSubscribe), anyMap()); verify(mockedConsumerFlowControlStrategy, times(1)) - .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); + .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap(), anyBoolean()); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(0) .isEqualTo(OffsetSpecification.next()); @@ -1754,7 +1754,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); verify(mockedConsumerFlowControlStrategy, times(2)) - .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap()); + .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap(), anyBoolean()); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(1) diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java index 2bd59b6047..a78e320900 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java @@ -26,11 +26,11 @@ import com.rabbitmq.stream.OffsetSpecification; import com.rabbitmq.stream.Producer; import com.rabbitmq.stream.StreamDoesNotExistException; -import com.rabbitmq.stream.flow.MessageHandlingAware; +import com.rabbitmq.stream.flow.MessageHandlingListener; +import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.MonitoringTestUtils.ConsumerInfo; -import com.rabbitmq.stream.impl.TestUtils.*; -import com.rabbitmq.stream.impl.flow.MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory; +import com.rabbitmq.stream.impl.flow.MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy; import io.netty.channel.EventLoopGroup; import org.assertj.core.api.ThrowableAssert.ThrowingCallable; import org.junit.jupiter.api.AfterEach; @@ -77,7 +77,7 @@ public class StreamConsumerTest { TestUtils.ClientFactory cf; Environment environment; - static Stream> consumerShouldKeepConsumingAfterDisruption() { + static Stream> consumerDisruptionTasks() { return Stream.of( TestUtils.namedTask( o -> { @@ -237,10 +237,79 @@ void consumeWithAsyncConsumerFlowControl() throws Exception { ConsumerBuilder consumerBuilder = environment.consumerBuilder().stream(stream) .offset(OffsetSpecification.first()); - MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.Builder flowControlStrategyBuilder = consumerBuilder - .flowControlStrategy(MaximumInflightChunksPerSubscriptionConsumerFlowControlStrategyBuilderFactory.INSTANCE) - .maximumInflightChunksPerSubscription(1); - MessageHandlingAware messageHandlingListener = flowControlStrategyBuilder.messageHandlingListener(); + MessageHandlingListenerConsumerBuilderAccessor messageHandlingListenerConsumerBuilderAccessor = consumerBuilder + .asynchronousControlFlow(5); + MessageHandlingListener messageHandlingListener = messageHandlingListenerConsumerBuilderAccessor.messageHandlingListener(); + + List messageContexts = new ArrayList<>(); + + AtomicBoolean shouldInstaConsume = new AtomicBoolean(false); + AtomicBoolean unhandledOnInstaConsume = new AtomicBoolean(false); + + consumerBuilder = messageHandlingListenerConsumerBuilderAccessor + .builder() + .messageHandler( + (context, message) -> { + if(shouldInstaConsume.get()) { + if(!messageHandlingListener.markHandled(context)) { + unhandledOnInstaConsume.set(true); + } + } else { + messageContexts.add(context); + } + firstConsumeLatch.countDown(); + chunkTimestamp.set(context.timestamp()); + consumeLatch.countDown(); + }); + Consumer consumer = consumerBuilder.build(); + + assertThat(firstConsumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isFalse(); + assertThat(chunkTimestamp.get()).isNotZero(); + + shouldInstaConsume.set(true); + boolean allMarkedHandled = messageContexts.parallelStream().allMatch(messageHandlingListener::markHandled); + assertThat(allMarkedHandled).isTrue(); + + assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + assertThat(unhandledOnInstaConsume.get()).isFalse(); + + consumer.close(); + } + + @Test + void consumeWithCustomAsyncConsumerFlowControl() throws Exception { + int messageCount = 100_000; + CountDownLatch publishLatch = new CountDownLatch(messageCount); + Client client = + cf.get( + new Client.ClientParameters() + .publishConfirmListener((publisherId, publishingId) -> publishLatch.countDown())); + + client.declarePublisher(b(1), null, stream); + IntStream.range(0, messageCount) + .forEach( + i -> + client.publish( + b(1), + Collections.singletonList( + client.messageBuilder().addData("".getBytes()).build()))); + + assertThat(publishLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + CountDownLatch firstConsumeLatch = new CountDownLatch(1); + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + + AtomicLong chunkTimestamp = new AtomicLong(); + + ConsumerBuilder consumerBuilder = environment.consumerBuilder().stream(stream) + .offset(OffsetSpecification.first()); + + MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.Builder flowControlStrategyBuilder = consumerBuilder + .customFlowControlStrategy(MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy::builder) + .maximumInflightChunksPerSubscription(5); + MessageHandlingListener messageHandlingListener = flowControlStrategyBuilder.messageHandlingListener(); List messageContexts = new ArrayList<>(); @@ -550,7 +619,7 @@ void consumerShouldReUseInitialOffsetSpecificationAfterDisruptionIfNoMessagesRec } @ParameterizedTest - @MethodSource + @MethodSource("consumerDisruptionTasks") @TestUtils.DisabledIfRabbitMqCtlNotSet void consumerShouldKeepConsumingAfterDisruption( java.util.function.Consumer disruption, TestInfo info) throws Exception { @@ -625,6 +694,113 @@ void consumerShouldKeepConsumingAfterDisruption( } } + @ParameterizedTest + @MethodSource("consumerDisruptionTasks") + @TestUtils.DisabledIfRabbitMqCtlNotSet + void consumerWithAsyncFlowControlShouldKeepConsumingAfterDisruption( + java.util.function.Consumer disruption, TestInfo info) throws Exception { + String s = streamName(info); + environment.streamCreator().stream(s).create(); + StreamConsumer consumer = null; + try { + int messageCount = 10_000; + CountDownLatch publishLatch = new CountDownLatch(messageCount); + Producer producer = environment.producerBuilder().stream(s).build(); + IntStream.range(0, messageCount) + .forEach( + i -> + producer.send( + producer.messageBuilder().addData("".getBytes()).build(), + confirmationStatus -> publishLatch.countDown())); + + assertThat(publishLatch.await(10, TimeUnit.SECONDS)).isTrue(); + producer.close(); + + AtomicInteger receivedMessageCount = new AtomicInteger(0); + CountDownLatch firstConsumeLatch = new CountDownLatch(1); + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + CountDownLatch consumeLatchSecondWave = new CountDownLatch(messageCount * 2); + + ConsumerBuilder consumerBuilder = environment.consumerBuilder().stream(s); + + MessageHandlingListenerConsumerBuilderAccessor messageHandlingListenerConsumerBuilderAccessor = consumerBuilder + .asynchronousControlFlow(5); + MessageHandlingListener messageHandlingListener = messageHandlingListenerConsumerBuilderAccessor.messageHandlingListener(); + consumerBuilder = messageHandlingListenerConsumerBuilderAccessor.builder(); + + List messageContexts = new ArrayList<>(); + + AtomicBoolean shouldInstaConsume = new AtomicBoolean(false); + AtomicBoolean unhandledOnInstaConsume = new AtomicBoolean(false); + + consumer = + (StreamConsumer) + consumerBuilder + .offset(OffsetSpecification.first()) + .messageHandler( + (context, message) -> { + if(shouldInstaConsume.get()) { + if(!messageHandlingListener.markHandled(context)) { + unhandledOnInstaConsume.set(true); + } + } else { + messageContexts.add(context); + } + receivedMessageCount.incrementAndGet(); + firstConsumeLatch.countDown(); + consumeLatch.countDown(); + consumeLatchSecondWave.countDown(); + }) + .build(); + + assertThat(firstConsumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isFalse(); + + assertThat(consumer.isOpen()).isTrue(); + + shouldInstaConsume.set(true); + boolean allMarkedHandled = messageContexts.parallelStream().allMatch(messageHandlingListener::markHandled); + assertThat(allMarkedHandled).isTrue(); + + assertThat(consumeLatch.await(20, TimeUnit.SECONDS)).isTrue(); + + assertThat(unhandledOnInstaConsume.get()).isFalse(); + + disruption.accept(s); + + Client client = cf.get(); + TestUtils.waitAtMost( + recoveryInitialDelay.plusSeconds(2), + () -> { + Client.StreamMetadata metadata = client.metadata(s).get(s); + return metadata.getLeader() != null || !metadata.getReplicas().isEmpty(); + }); + + CountDownLatch publishLatchSecondWave = new CountDownLatch(messageCount); + Producer producerSecondWave = environment.producerBuilder().stream(s).build(); + IntStream.range(0, messageCount) + .forEach( + i -> + producerSecondWave.send( + producerSecondWave.messageBuilder().addData("".getBytes()).build(), + confirmationStatus -> publishLatchSecondWave.countDown())); + + assertThat(publishLatchSecondWave.await(20, TimeUnit.SECONDS)).isTrue(); + producerSecondWave.close(); + + latchAssert(consumeLatchSecondWave).completes(recoveryInitialDelay.plusSeconds(30)); + assertThat(receivedMessageCount.get()) + .isBetween(messageCount * 2, messageCount * 2 + 1); // there can be a duplicate + assertThat(consumer.isOpen()).isTrue(); + + } finally { + if (consumer != null) { + consumer.close(); + } + environment.deleteStream(s); + } + } + @Test void autoTrackingShouldStorePeriodicallyAndAfterInactivity() throws Exception { AtomicInteger messageCount = new AtomicInteger(0); From a54100b9e52ef36861d93fff08effeb9e58f586d Mon Sep 17 00:00:00 2001 From: henry701 Date: Wed, 21 Jun 2023 20:06:10 -0300 Subject: [PATCH 09/14] Remove redundant 0 assignments to trackerCount field --- .../java/com/rabbitmq/stream/impl/ConsumersCoordinator.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 015916561f..6e4dcfd53d 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -543,7 +543,7 @@ private class ClientSubscriptionsManager implements Comparable subscriptionTrackers = new ArrayList<>(maxConsumersByConnection); - private volatile int trackerCount = 0; + private volatile int trackerCount; private final AtomicBoolean closed = new AtomicBoolean(false); private ClientSubscriptionsManager( @@ -556,7 +556,6 @@ private ClientSubscriptionsManager( this.name = keyForClientSubscription(node); LOGGER.debug("creating subscription manager on {}", name); IntStream.range(0, maxConsumersByConnection).forEach(i -> subscriptionTrackers.add(null)); - this.trackerCount = 0; AtomicReference clientReference = new AtomicReference<>(); ConsumerFlowControlStrategy localConsumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get); ClientCallbackStreamDataHandlerAdapter clientListenerAdaptedConsumerFlowControlStrategy = new ClientCallbackStreamDataHandlerAdapter( From 7250c8ae4c21d088add5464a9241f38da0a1a055 Mon Sep 17 00:00:00 2001 From: henry701 Date: Sun, 25 Jun 2023 17:39:33 -0300 Subject: [PATCH 10/14] Change lifecycle of each consumer flow control strategy instance to be per-subscription --- .../com/rabbitmq/stream/MessageHandler.java | 4 +- .../AbstractConsumerFlowControlStrategy.java | 29 +++-- .../flow/ConsumerFlowControlStrategy.java | 2 + .../ConsumerFlowControlStrategyBuilder.java | 2 + .../stream/flow/MessageHandlingListener.java | 1 + .../stream/impl/ConsumersCoordinator.java | 108 ++++++++++-------- .../stream/impl/SuperStreamConsumer.java | 20 ++-- .../stream/impl/ConsumersCoordinatorTest.java | 8 -- .../impl/OffsetTrackingCoordinatorTest.java | 28 +++-- .../rabbitmq/stream/impl/SacClientTest.java | 34 +++--- 10 files changed, 132 insertions(+), 104 deletions(-) diff --git a/src/main/java/com/rabbitmq/stream/MessageHandler.java b/src/main/java/com/rabbitmq/stream/MessageHandler.java index d9d39fdb63..1e21b412b1 100644 --- a/src/main/java/com/rabbitmq/stream/MessageHandler.java +++ b/src/main/java/com/rabbitmq/stream/MessageHandler.java @@ -13,6 +13,8 @@ // info@rabbitmq.com. package com.rabbitmq.stream; +import com.rabbitmq.stream.flow.MessageHandlingListener; + /** * Callback API for inbound messages. * @@ -30,7 +32,7 @@ public interface MessageHandler { void handle(Context context, Message message); /** Information about the message. */ - interface Context { + interface Context extends MessageHandlingListener { /** * The offset of the message in the stream. diff --git a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java index e45015ca0a..75ccdddd6e 100644 --- a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java @@ -10,24 +10,35 @@ public abstract class AbstractConsumerFlowControlStrategy implements ConsumerFlowControlStrategy { private final Supplier creditAskerSupplier; - private volatile CreditAsker creditAsker; - + private volatile CreditAsker lastRetrievedCreditAsker; protected AbstractConsumerFlowControlStrategy(Supplier creditAskerSupplier) { this.creditAskerSupplier = Objects.requireNonNull(creditAskerSupplier, "creditAskerSupplier"); } protected CreditAsker mandatoryCreditAsker() { - CreditAsker localSupplied = this.creditAsker; - if(localSupplied != null) { - return localSupplied; - } - localSupplied = creditAskerSupplier.get(); + CreditAsker localSupplied = creditAskerSupplier.get(); if(localSupplied == null) { - throw new IllegalStateException("Requested client, but client is not yet available! Supplier: " + this.creditAskerSupplier); + throw new IllegalStateException("Requested CreditAsker but it's not yet available! Supplier: " + this.creditAskerSupplier); } - this.creditAsker = localSupplied; + this.lastRetrievedCreditAsker = localSupplied; return localSupplied; } + public CreditAsker getLastRetrievedCreditAsker() { + return lastRetrievedCreditAsker; + } + + public Supplier getCreditAskerSupplier() { + return creditAskerSupplier; + } + + @Override + public String toString() { + return "AbstractConsumerFlowControlStrategy{" + + "creditAskerSupplier=" + creditAskerSupplier + + ", lastRetrievedCreditAsker=" + lastRetrievedCreditAsker + + '}'; + } + } diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java index c9e1dba68e..953d7b780d 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java @@ -9,6 +9,8 @@ * 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 CreditAsker#credit} to ask for credits. + * One instance of this is expected to be built for each separate subscription. + * A {@link com.rabbitmq.stream.Consumer} may have multiple subscriptions, and thus multiple instances of this. */ public interface ConsumerFlowControlStrategy extends CallbackStreamDataHandler { diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java index d732451627..3b5c4444b3 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java @@ -6,6 +6,8 @@ /** * Fluent builder for a {@link ConsumerFlowControlStrategyBuilderFactory}. + * One instance of this is set per {@link com.rabbitmq.stream.Consumer}. + * A {@link com.rabbitmq.stream.Consumer} may have multiple subscriptions, and thus multiple instances built by this. * * @param the type of {@link ConsumerFlowControlStrategy} to be built */ diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java index 39edb6a119..7884e52727 100644 --- a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java +++ b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListener.java @@ -2,6 +2,7 @@ import com.rabbitmq.stream.MessageHandler; +@FunctionalInterface public interface MessageHandlingListener { /** diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 6e4dcfd53d..4871440f00 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -26,6 +26,7 @@ import com.rabbitmq.stream.SubscriptionListener.SubscriptionContext; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; +import com.rabbitmq.stream.flow.MessageHandlingListener; import com.rabbitmq.stream.impl.Client.Broker; import com.rabbitmq.stream.impl.Client.ClientParameters; import com.rabbitmq.stream.impl.Client.ConsumerUpdateListener; @@ -136,7 +137,6 @@ Runnable subscribe( messageHandler, consumerFlowControlStrategyBuilder, subscriptionProperties); - try { addToManager(newNode, subscriptionTracker, offsetSpecification, true); } catch (ConnectionStreamException e) { @@ -160,10 +160,10 @@ Runnable subscribe( } private void addToManager( - Broker node, - SubscriptionTracker tracker, - OffsetSpecification offsetSpecification, - boolean isInitialSubscription) { + Broker node, + SubscriptionTracker tracker, + OffsetSpecification offsetSpecification, + boolean isInitialSubscription) { ClientParameters clientParameters = environment .clientParametersCopy() @@ -190,9 +190,10 @@ private void addToManager( if (pickedManager == null) { String name = keyForClientSubscription(node); LOGGER.debug("Creating subscription manager on {}", name); - pickedManager = new ClientSubscriptionsManager(node, clientParameters, tracker.consumerFlowControlStrategyBuilder); + pickedManager = new ClientSubscriptionsManager(node, clientParameters); LOGGER.debug("Created subscription manager on {}, id {}", name, pickedManager.id); } + tracker.clientReference.set(pickedManager.client); try { pickedManager.add(tracker, offsetSpecification, isInitialSubscription); LOGGER.debug( @@ -204,22 +205,22 @@ private void addToManager( tracker.subscriptionIdInClient); this.managers.add(pickedManager); } catch (IllegalStateException e) { + tracker.clientReference.set(null); pickedManager = null; } catch (ConnectionStreamException | ClientClosedException | StreamNotAvailableException e) { + tracker.clientReference.set(null); // manager connection is dead or stream not available // scheduling manager closing if necessary in another thread to avoid blocking this one if (pickedManager.isEmpty()) { - ClientSubscriptionsManager manager = pickedManager; - ConsumersCoordinator.this.environment.execute( - () -> { - manager.closeIfEmpty(); - }, + ConsumersCoordinator.this.environment.execute( + pickedManager::closeIfEmpty, "Consumer manager closing after timeout, consumer %d on stream '%s'", tracker.consumer.id(), tracker.stream); } throw e; } catch (RuntimeException e) { + tracker.clientReference.set(null); if (pickedManager != null) { pickedManager.closeIfEmpty(); } @@ -384,7 +385,8 @@ private static class SubscriptionTracker { private final OffsetSpecification initialOffsetSpecification; private final String offsetTrackingReference; private final MessageHandler messageHandler; - private final ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder; + private final ConsumerFlowControlStrategy consumerFlowControlStrategy; + private final AtomicReference clientReference; private final StreamConsumer consumer; private final SubscriptionListener subscriptionListener; private final Runnable trackingClosingCallback; @@ -397,16 +399,16 @@ private static class SubscriptionTracker { new AtomicReference<>(SubscriptionState.OPENING); private SubscriptionTracker( - long id, - StreamConsumer consumer, - String stream, - OffsetSpecification initialOffsetSpecification, - String offsetTrackingReference, - SubscriptionListener subscriptionListener, - Runnable trackingClosingCallback, - MessageHandler messageHandler, - ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, - Map subscriptionProperties) { + long id, + StreamConsumer consumer, + String stream, + OffsetSpecification initialOffsetSpecification, + String offsetTrackingReference, + SubscriptionListener subscriptionListener, + Runnable trackingClosingCallback, + MessageHandler messageHandler, + ConsumerFlowControlStrategyBuilder consumerFlowControlStrategyBuilder, + Map subscriptionProperties) { this.id = id; this.consumer = consumer; this.stream = stream; @@ -415,7 +417,8 @@ private SubscriptionTracker( this.subscriptionListener = subscriptionListener; this.trackingClosingCallback = trackingClosingCallback; this.messageHandler = messageHandler; - this.consumerFlowControlStrategyBuilder = consumerFlowControlStrategyBuilder; + this.clientReference = new AtomicReference<>(); + this.consumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(this.clientReference::get); if (this.offsetTrackingReference == null) { this.subscriptionProperties = subscriptionProperties; } else { @@ -485,13 +488,15 @@ private static final class MessageHandlerContext implements Context { private final long timestamp; private final long committedOffset; private final StreamConsumer consumer; + private final MessageHandlingListener handledCallback; private MessageHandlerContext( - long offset, long timestamp, long committedOffset, StreamConsumer consumer) { + long offset, long timestamp, long committedOffset, StreamConsumer consumer, MessageHandlingListener handledCallback) { this.offset = offset; this.timestamp = timestamp; this.committedOffset = committedOffset; this.consumer = consumer; + this.handledCallback = handledCallback; } @Override @@ -522,6 +527,12 @@ public String stream() { public Consumer consumer() { return this.consumer; } + + @Override + public boolean markHandled(Context messageContext) { + return this.handledCallback.markHandled(messageContext); + } + } /** @@ -539,7 +550,6 @@ private class ClientSubscriptionsManager implements Comparable> streamToStreamSubscriptions = new ConcurrentHashMap<>(); - private final ConsumerFlowControlStrategy consumerFlowControlStrategy; // trackers and tracker count must be kept in sync private volatile List subscriptionTrackers = new ArrayList<>(maxConsumersByConnection); @@ -548,23 +558,15 @@ private class ClientSubscriptionsManager implements Comparable consumerFlowControlStrategyBuilder + Client.ClientParameters clientParameters ) { this.id = managerIdSequence.getAndIncrement(); this.node = node; this.name = keyForClientSubscription(node); LOGGER.debug("creating subscription manager on {}", name); IntStream.range(0, maxConsumersByConnection).forEach(i -> subscriptionTrackers.add(null)); - AtomicReference clientReference = new AtomicReference<>(); - ConsumerFlowControlStrategy localConsumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(clientReference::get); - ClientCallbackStreamDataHandlerAdapter clientListenerAdaptedConsumerFlowControlStrategy = new ClientCallbackStreamDataHandlerAdapter( - localConsumerFlowControlStrategy - ); - this.consumerFlowControlStrategy = localConsumerFlowControlStrategy; CreditNotification creditNotification = (subscriptionId, responseCode) -> { - localConsumerFlowControlStrategy.handleCreditNotification(subscriptionId, responseCode); SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); String stream = subscriptionTracker == null ? "?" : subscriptionTracker.stream; @@ -577,15 +579,28 @@ private ClientSubscriptionsManager( MessageListener messageListener = (subscriptionId, offset, chunkTimestamp, committedOffset, message) -> { - localConsumerFlowControlStrategy.handleMessage(subscriptionId, offset, chunkTimestamp, committedOffset, message); SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); if (subscriptionTracker != null) { subscriptionTracker.offset = offset; subscriptionTracker.hasReceivedSomething = true; + subscriptionTracker.consumerFlowControlStrategy.handleMessage( + subscriptionId, + offset, + chunkTimestamp, + committedOffset, + message + ); subscriptionTracker.messageHandler.handle( new MessageHandlerContext( - offset, chunkTimestamp, committedOffset, subscriptionTracker.consumer), + offset, + chunkTimestamp, + committedOffset, + subscriptionTracker.consumer, + subscriptionTracker.consumerFlowControlStrategy instanceof MessageHandlingListener + ? (MessageHandlingListener) subscriptionTracker.consumerFlowControlStrategy + : c -> true + ), message); // FIXME set offset here as well, best effort to avoid duplicates? } else { @@ -648,14 +663,12 @@ private ClientSubscriptionsManager( "Consumers re-assignment after disconnection from %s", name)); } - clientListenerAdaptedConsumerFlowControlStrategy.handleShutdown(shutdownContext); }; MetadataListener metadataListener = (stream, code) -> { LOGGER.debug( "Received metadata notification for '{}', stream is likely to have become unavailable", stream); - localConsumerFlowControlStrategy.handleMetadata(stream, code); Set affectedSubscriptions; synchronized (this) { Set subscriptions = streamToStreamSubscriptions.remove(stream); @@ -704,7 +717,6 @@ private ClientSubscriptionsManager( }; ConsumerUpdateListener consumerUpdateListener = (client, subscriptionId, active) -> { - localConsumerFlowControlStrategy.handleConsumerUpdate(subscriptionId, active); OffsetSpecification result = null; SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); @@ -722,21 +734,27 @@ private ClientSubscriptionsManager( } return result; }; + Client.ChunkListener chunkListener = (client, subscriptionId, offset, messageCount, dataSize) -> { + SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); + if(subscriptionTracker == null) { + LOGGER.warn("Could not find stream subscription {} for chunk listener", subscriptionId); + return; + } + subscriptionTracker.consumerFlowControlStrategy.handleChunk(subscriptionId, offset, messageCount, dataSize); + }; String connectionName = connectionNamingStrategy.apply(ClientConnectionType.CONSUMER); ClientFactoryContext clientFactoryContext = ClientFactoryContext.fromParameters( clientParameters .clientProperty("connection_name", connectionName) - .chunkListener(clientListenerAdaptedConsumerFlowControlStrategy) + .chunkListener(chunkListener) .creditNotification(creditNotification) .messageListener(messageListener) .shutdownListener(shutdownListener) .metadataListener(metadataListener) .consumerUpdateListener(consumerUpdateListener)) .key(name); - Client localClient = clientFactory.client(clientFactoryContext); - this.client = localClient; - clientReference.set(localClient); + this.client = clientFactory.client(clientFactoryContext); LOGGER.debug("Created consumer connection '{}'", connectionName); } @@ -951,7 +969,7 @@ synchronized void add( subscriptionContext.offsetSpecification()); checkNotClosed(); - int initialCredits = this.consumerFlowControlStrategy.handleSubscribeReturningInitialCredits( + int initialCredits = subscriptionTracker.consumerFlowControlStrategy.handleSubscribeReturningInitialCredits( subscriptionId, subscriptionTracker.stream, subscriptionContext.offsetSpecification(), @@ -1030,7 +1048,6 @@ 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, @@ -1095,7 +1112,6 @@ 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( diff --git a/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java index c884da7e3c..1d5afbcba8 100644 --- a/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java +++ b/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java @@ -13,20 +13,21 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.impl.Utils.namedFunction; - import com.rabbitmq.stream.Consumer; import com.rabbitmq.stream.Message; import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; import com.rabbitmq.stream.impl.Utils.CompositeConsumerUpdateListener; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.concurrent.ConcurrentHashMap; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; + +import static com.rabbitmq.stream.impl.Utils.namedFunction; class SuperStreamConsumer implements Consumer { @@ -126,6 +127,11 @@ private ManualOffsetTrackingMessageHandler( public void handle(Context context, Message message) { Context ctx = new Context() { + @Override + public boolean markHandled(Context messageContext) { + return context.markHandled(messageContext); + } + @Override public long offset() { return context.offset(); @@ -145,7 +151,7 @@ public long committedChunkId() { public void storeOffset() { for (ConsumerState state : consumerStates) { if (ManualOffsetTrackingMessageHandler.this.consumerState == state) { - maybeStoreOffset(state, () -> context.storeOffset()); + maybeStoreOffset(state, context::storeOffset); } else if (state.offset != 0) { maybeStoreOffset(state, () -> state.consumer.store(state.offset)); } @@ -153,9 +159,7 @@ public void storeOffset() { } private void maybeStoreOffset(ConsumerState state, Runnable storeAction) { - if (state.consumer.isSac() && !state.consumer.sacActive()) { - // do nothing - } else { + if (!state.consumer.isSac() || state.consumer.sacActive()) { storeAction.run(); } } diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index 97396f2189..8be1de6ac2 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -1763,14 +1763,6 @@ void shouldCallConsumerFlowControlHandlers(Consumer co assertThat(subscriptionPropertiesArgumentCaptor.getAllValues()) .element(1) .isEqualTo(Collections.singletonMap("name", "consumer-name")); - when(client.unsubscribe(subscriptionIdCaptor.getValue())) - .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); - - closingRunnable.run(); - - verify(client, times(1)).unsubscribe(subscriptionIdCaptor.getValue()); - verify(mockedConsumerFlowControlStrategy, times(1)) - .handleUnsubscribe(subscriptionIdCaptor.getValue()); } Client.Broker leader() { diff --git a/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java index 8a7af3ece6..f8355ac9c3 100644 --- a/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java @@ -13,14 +13,16 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.impl.TestUtils.answer; -import static com.rabbitmq.stream.impl.TestUtils.latchAssert; -import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.Mockito.*; - import com.rabbitmq.stream.MessageHandler.Context; import com.rabbitmq.stream.impl.OffsetTrackingCoordinator.Registration; import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + import java.time.Duration; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; @@ -29,12 +31,11 @@ import java.util.function.Consumer; import java.util.function.LongConsumer; import java.util.stream.IntStream; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; + +import static com.rabbitmq.stream.impl.TestUtils.answer; +import static com.rabbitmq.stream.impl.TestUtils.latchAssert; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.*; public class OffsetTrackingCoordinatorTest { @@ -334,6 +335,11 @@ void manualShouldStoreIfRequestedStoredOffsetIsBehind() { Context context(long offset, Runnable action) { return new Context() { + @Override + public boolean markHandled(Context messageContext) { + return true; + } + @Override public long offset() { return offset; diff --git a/src/test/java/com/rabbitmq/stream/impl/SacClientTest.java b/src/test/java/com/rabbitmq/stream/impl/SacClientTest.java index b915375448..85db4f670c 100644 --- a/src/test/java/com/rabbitmq/stream/impl/SacClientTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/SacClientTest.java @@ -13,19 +13,6 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import static com.rabbitmq.stream.impl.TestUtils.BrokerVersion.RABBITMQ_3_11_14; -import static com.rabbitmq.stream.impl.TestUtils.ResponseConditions.ko; -import static com.rabbitmq.stream.impl.TestUtils.ResponseConditions.ok; -import static com.rabbitmq.stream.impl.TestUtils.ResponseConditions.responseCode; -import static com.rabbitmq.stream.impl.TestUtils.b; -import static com.rabbitmq.stream.impl.TestUtils.declareSuperStreamTopology; -import static com.rabbitmq.stream.impl.TestUtils.deleteSuperStreamTopology; -import static com.rabbitmq.stream.impl.TestUtils.latchAssert; -import static com.rabbitmq.stream.impl.TestUtils.streamName; -import static com.rabbitmq.stream.impl.TestUtils.waitAtMost; -import static java.util.stream.Collectors.toList; -import static org.assertj.core.api.Assertions.assertThat; - import com.rabbitmq.client.Connection; import com.rabbitmq.client.ConnectionFactory; import com.rabbitmq.stream.Constants; @@ -36,9 +23,11 @@ import com.rabbitmq.stream.impl.Client.CreditNotification; import com.rabbitmq.stream.impl.Client.MessageListener; import com.rabbitmq.stream.impl.Client.Response; -import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast; -import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast311Condition; -import com.rabbitmq.stream.impl.TestUtils.DisabledIfRabbitMqCtlNotSet; +import com.rabbitmq.stream.impl.TestUtils.*; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInfo; +import org.junit.jupiter.api.extension.ExtendWith; + import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.Collections; @@ -53,9 +42,12 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.stream.IntStream; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInfo; -import org.junit.jupiter.api.extension.ExtendWith; + +import static com.rabbitmq.stream.impl.TestUtils.BrokerVersion.RABBITMQ_3_11_14; +import static com.rabbitmq.stream.impl.TestUtils.ResponseConditions.*; +import static com.rabbitmq.stream.impl.TestUtils.*; +import static java.util.stream.Collectors.toList; +import static org.assertj.core.api.Assertions.assertThat; @ExtendWith({ TestUtils.StreamTestInfrastructureExtension.class, @@ -522,7 +514,7 @@ void superStreamRebalancingShouldWorkWhilePublishing(TestInfo info) throws Excep // we keep track of credit errors // with the amount of initial credit and the rebalancing, // the first subscriber is likely to have in-flight credit commands - // when it becomes inactive. The server should then sends some credit + // when it becomes inactive. The server should then send some credit // notifications to tell the client it's not supposed to ask for credits // for this subscription. CreditNotification creditNotification = @@ -572,7 +564,7 @@ void superStreamRebalancingShouldWorkWhilePublishing(TestInfo info) throws Excep waitAtMost( () -> - creditNotificationResponseCode.get() == Constants.RESPONSE_CODE_PRECONDITION_FAILED); + creditNotificationResponseCode.get() == Constants.RESPONSE_CODE_PRECONDITION_FAILED, () -> "Code was actually: " + creditNotificationResponseCode.get()); Response response = client1.unsubscribe(b(0)); assertThat(response).is(ok()); From 70486f4bb68fe4a37cdac951467813a7c505029c Mon Sep 17 00:00:00 2001 From: henry701 Date: Tue, 27 Jun 2023 21:26:41 -0300 Subject: [PATCH 11/14] Change MessageHandler.Context.markHandled to receive no parameters, and use it in one of the integrated tests to mark messages as handled --- .../java/com/rabbitmq/stream/MessageHandler.java | 12 +++++++++--- .../rabbitmq/stream/impl/ConsumersCoordinator.java | 4 ++-- .../rabbitmq/stream/impl/SuperStreamConsumer.java | 4 ++-- .../stream/impl/OffsetTrackingCoordinatorTest.java | 2 +- .../com/rabbitmq/stream/impl/StreamConsumerTest.java | 2 +- 5 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/main/java/com/rabbitmq/stream/MessageHandler.java b/src/main/java/com/rabbitmq/stream/MessageHandler.java index 1e21b412b1..907ba62d61 100644 --- a/src/main/java/com/rabbitmq/stream/MessageHandler.java +++ b/src/main/java/com/rabbitmq/stream/MessageHandler.java @@ -13,8 +13,6 @@ // info@rabbitmq.com. package com.rabbitmq.stream; -import com.rabbitmq.stream.flow.MessageHandlingListener; - /** * Callback API for inbound messages. * @@ -32,7 +30,7 @@ public interface MessageHandler { void handle(Context context, Message message); /** Information about the message. */ - interface Context extends MessageHandlingListener { + interface Context { /** * The offset of the message in the stream. @@ -86,5 +84,13 @@ interface Context extends MessageHandlingListener { * @see Consumer#store(long) */ Consumer consumer(); + + /** + * Marks this message as handled + * + * @return Whether the message was marked as handled (returning {@code true}) + * or was not found (either because it was already marked as handled, or wasn't tracked) + */ + boolean markHandled(); } } diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 4871440f00..b5bd339f4c 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -529,8 +529,8 @@ public Consumer consumer() { } @Override - public boolean markHandled(Context messageContext) { - return this.handledCallback.markHandled(messageContext); + public boolean markHandled() { + return this.handledCallback.markHandled(this); } } diff --git a/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java index 1d5afbcba8..014f6a33b8 100644 --- a/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java +++ b/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java @@ -128,8 +128,8 @@ public void handle(Context context, Message message) { Context ctx = new Context() { @Override - public boolean markHandled(Context messageContext) { - return context.markHandled(messageContext); + public boolean markHandled() { + return context.markHandled(); } @Override diff --git a/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java index f8355ac9c3..701ff883b2 100644 --- a/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java @@ -336,7 +336,7 @@ void manualShouldStoreIfRequestedStoredOffsetIsBehind() { Context context(long offset, Runnable action) { return new Context() { @Override - public boolean markHandled(Context messageContext) { + public boolean markHandled() { return true; } diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java index a78e320900..da80db7d75 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java @@ -251,7 +251,7 @@ void consumeWithAsyncConsumerFlowControl() throws Exception { .messageHandler( (context, message) -> { if(shouldInstaConsume.get()) { - if(!messageHandlingListener.markHandled(context)) { + if(!context.markHandled()) { unhandledOnInstaConsume.set(true); } } else { From a156eb0e03ac9b4c66f2e1f4c12452867cd0f3df Mon Sep 17 00:00:00 2001 From: henry701 Date: Tue, 27 Jun 2023 22:05:17 -0300 Subject: [PATCH 12/14] Remove MessageHandlingListenerConsumerBuilderAccessor --- .../com/rabbitmq/stream/ConsumerBuilder.java | 41 +++++++++++++------ ...ndlingListenerConsumerBuilderAccessor.java | 7 ---- .../stream/impl/StreamConsumerBuilder.java | 24 +++++------ ...ptionAsyncConsumerFlowControlStrategy.java | 24 +---------- .../stream/impl/StreamConsumerTest.java | 28 ++++--------- 5 files changed, 51 insertions(+), 73 deletions(-) delete mode 100644 src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java diff --git a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java index b18bb2e297..933ad569d9 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java @@ -16,8 +16,6 @@ import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; -import com.rabbitmq.stream.flow.MessageHandlingListener; -import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import java.time.Duration; @@ -66,22 +64,41 @@ public interface ConsumerBuilder { ConsumerBuilder messageHandler(MessageHandler messageHandler); /** - * Configure credit parameters for synchronous flow control. - * - * @param initial Credits to ask for with each new subscription - * @param onChunkDelivery Credits to ask for after a chunk is delivered + * Configure prefetching parameters for synchronous flow control. + * + *

+ * Treat the parameters as an abstract scale at the {@link Consumer} level. + *

+ * + * @param initialPrefetchLevel The initial level of message pre-fetching. + * This may me implemented as the credits to initially + * ask for with each new subscription, + * but do not depend strongly on this aspect. + * @param prefetchLevelAfterDelivery The level of message pre-fetching after the initial batch. + * This may be implemented as the credits to ask for after a chunk + * is delivered on each subscription, + * but do not depend strongly on this aspect. + * + * The recommended value is 1. + * Higher values may cause excessive over-fetching. + * * @return this {@link ConsumerBuilder} */ - ConsumerBuilder synchronousControlFlow(int initial, int onChunkDelivery); + ConsumerBuilder synchronousControlFlow(int initialPrefetchLevel, int prefetchLevelAfterDelivery); /** - * Configure credit parameters for asynchronous flow control. + * Configure prefetching parameters for asynchronous flow control. + * + *

+ * Treat the parameters as an abstract scale at the {@link Consumer} level. + *

* - * @param concurrencyLevel Maximum chunks to have in-processing at a given moment - * @return A {@link MessageHandlingListenerConsumerBuilderAccessor} for obtaining the {@link MessageHandlingListener} - * and navigating fluently back to this {@link ConsumerBuilder} + * @param prefetchLevel The desired level of message pre-fetching. + * This may be implemented as the maximum chunks to have in processing or pending arrival + * per subscription at a given moment, but do not depend strongly on this aspect. + * @return this {@link ConsumerBuilder} */ - MessageHandlingListenerConsumerBuilderAccessor asynchronousControlFlow(int concurrencyLevel); + ConsumerBuilder asynchronousControlFlow(int prefetchLevel); /** * Factory for the flow control strategy to be used when consuming messages. diff --git a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java b/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java deleted file mode 100644 index 6c14552e65..0000000000 --- a/src/main/java/com/rabbitmq/stream/flow/MessageHandlingListenerConsumerBuilderAccessor.java +++ /dev/null @@ -1,7 +0,0 @@ -package com.rabbitmq.stream.flow; - -import com.rabbitmq.stream.ConsumerBuilder; - -public interface MessageHandlingListenerConsumerBuilderAccessor extends ConsumerBuilder.ConsumerBuilderAccessor { - MessageHandlingListener messageHandlingListener(); -} \ No newline at end of file diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java index fb2133b71a..26a44e3dd8 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java @@ -23,7 +23,6 @@ import com.rabbitmq.stream.flow.ConsumerFlowControlStrategy; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilderFactory; -import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import com.rabbitmq.stream.impl.flow.MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy; import com.rabbitmq.stream.impl.flow.SynchronousConsumerFlowControlStrategy; @@ -151,32 +150,31 @@ public ConsumerBuilder noTrackingStrategy() { * Configure credit parameters for flow control. * Implies usage of a traditional {@link SynchronousConsumerFlowControlStrategy}. * - * @param initial Credits to ask for with each new subscription - * @param onChunkDelivery Credits to ask for after a chunk is delivered + * @param initialPrefetchLevel Credits to ask for with each new subscription + * @param prefetchLevelAfterDelivery Credits to ask for after a chunk is delivered * @return this {@link StreamConsumerBuilder} */ @Override - public StreamConsumerBuilder synchronousControlFlow(int initial, int onChunkDelivery) { - if (initial <= 0 || onChunkDelivery <= 0) { + public StreamConsumerBuilder synchronousControlFlow(int initialPrefetchLevel, int prefetchLevelAfterDelivery) { + if (initialPrefetchLevel <= 0 || prefetchLevelAfterDelivery <= 0) { throw new IllegalArgumentException("Credits must be positive"); } this.consumerFlowControlStrategyBuilder = SynchronousConsumerFlowControlStrategy .builder(this) - .initialCredits(initial) - .additionalCredits(onChunkDelivery); + .initialCredits(initialPrefetchLevel) + .additionalCredits(prefetchLevelAfterDelivery); return this; } @Override - public MessageHandlingListenerConsumerBuilderAccessor asynchronousControlFlow(int concurrencyLevel) { - if (concurrencyLevel <= 0) { + public ConsumerBuilder asynchronousControlFlow(int prefetchLevel) { + if (prefetchLevel <= 0) { throw new IllegalArgumentException("ConcurrencyLevel must be positive"); } - MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.Builder localBuilder = MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy + this.consumerFlowControlStrategyBuilder = MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy .builder(this) - .maximumInflightChunksPerSubscription(concurrencyLevel); - this.consumerFlowControlStrategyBuilder = localBuilder; - return localBuilder; + .maximumInflightChunksPerSubscription(prefetchLevel); + return this; } StreamConsumerBuilder lazyInit(boolean lazyInit) { diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java index e5b3c41408..33e40273f9 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java @@ -6,15 +6,11 @@ import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.CreditAsker; import com.rabbitmq.stream.flow.MessageHandlingListener; -import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import com.rabbitmq.stream.impl.ConsumerStatisticRecorder; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Collections; import java.util.Map; -import java.util.Set; -import java.util.WeakHashMap; import java.util.function.IntUnaryOperator; import java.util.function.Supplier; @@ -91,26 +87,22 @@ public static Builder builder(ConsumerBuilder consumerBuilder) { return new Builder(consumerBuilder); } - public static class Builder implements ConsumerFlowControlStrategyBuilder, MessageHandlingListenerConsumerBuilderAccessor { + public static class Builder implements ConsumerFlowControlStrategyBuilder, ConsumerBuilder.ConsumerBuilderAccessor { private final ConsumerBuilder consumerBuilder; private int maximumInflightChunksPerSubscription = 1; - private final Set instances = Collections.newSetFromMap(new WeakHashMap<>()); - public Builder(ConsumerBuilder consumerBuilder) { this.consumerBuilder = consumerBuilder; } @Override public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { - MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy built = new MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( + return new MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( creditAskerSupplier, this.maximumInflightChunksPerSubscription ); - instances.add(built); - return built; } @Override @@ -123,18 +115,6 @@ public Builder maximumInflightChunksPerSubscription(int maximumInflightChunksPer return this; } - @Override - public MessageHandlingListener messageHandlingListener() { - return this::messageHandlingMulticaster; - } - - private boolean messageHandlingMulticaster(MessageHandler.Context context) { - boolean changed = false; - for(MessageHandlingListener instance : instances) { - changed = changed || instance.markHandled(context); - } - return changed; - } } } diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java index da80db7d75..b321e18ad3 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java @@ -26,8 +26,6 @@ import com.rabbitmq.stream.OffsetSpecification; import com.rabbitmq.stream.Producer; import com.rabbitmq.stream.StreamDoesNotExistException; -import com.rabbitmq.stream.flow.MessageHandlingListener; -import com.rabbitmq.stream.flow.MessageHandlingListenerConsumerBuilderAccessor; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.MonitoringTestUtils.ConsumerInfo; import com.rabbitmq.stream.impl.flow.MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy; @@ -235,19 +233,15 @@ void consumeWithAsyncConsumerFlowControl() throws Exception { AtomicLong chunkTimestamp = new AtomicLong(); ConsumerBuilder consumerBuilder = environment.consumerBuilder().stream(stream) - .offset(OffsetSpecification.first()); - - MessageHandlingListenerConsumerBuilderAccessor messageHandlingListenerConsumerBuilderAccessor = consumerBuilder + .offset(OffsetSpecification.first()) .asynchronousControlFlow(5); - MessageHandlingListener messageHandlingListener = messageHandlingListenerConsumerBuilderAccessor.messageHandlingListener(); List messageContexts = new ArrayList<>(); AtomicBoolean shouldInstaConsume = new AtomicBoolean(false); AtomicBoolean unhandledOnInstaConsume = new AtomicBoolean(false); - consumerBuilder = messageHandlingListenerConsumerBuilderAccessor - .builder() + consumerBuilder = consumerBuilder .messageHandler( (context, message) -> { if(shouldInstaConsume.get()) { @@ -268,7 +262,7 @@ void consumeWithAsyncConsumerFlowControl() throws Exception { assertThat(chunkTimestamp.get()).isNotZero(); shouldInstaConsume.set(true); - boolean allMarkedHandled = messageContexts.parallelStream().allMatch(messageHandlingListener::markHandled); + boolean allMarkedHandled = messageContexts.parallelStream().allMatch(MessageHandler.Context::markHandled); assertThat(allMarkedHandled).isTrue(); assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); @@ -309,7 +303,6 @@ void consumeWithCustomAsyncConsumerFlowControl() throws Exception { MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.Builder flowControlStrategyBuilder = consumerBuilder .customFlowControlStrategy(MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy::builder) .maximumInflightChunksPerSubscription(5); - MessageHandlingListener messageHandlingListener = flowControlStrategyBuilder.messageHandlingListener(); List messageContexts = new ArrayList<>(); @@ -321,7 +314,7 @@ void consumeWithCustomAsyncConsumerFlowControl() throws Exception { .messageHandler( (context, message) -> { if(shouldInstaConsume.get()) { - if(!messageHandlingListener.markHandled(context)) { + if(!context.markHandled()) { unhandledOnInstaConsume.set(true); } } else { @@ -338,7 +331,7 @@ void consumeWithCustomAsyncConsumerFlowControl() throws Exception { assertThat(chunkTimestamp.get()).isNotZero(); shouldInstaConsume.set(true); - boolean allMarkedHandled = messageContexts.parallelStream().allMatch(messageHandlingListener::markHandled); + boolean allMarkedHandled = messageContexts.parallelStream().allMatch(MessageHandler.Context::markHandled); assertThat(allMarkedHandled).isTrue(); assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); @@ -721,12 +714,9 @@ void consumerWithAsyncFlowControlShouldKeepConsumingAfterDisruption( CountDownLatch consumeLatch = new CountDownLatch(messageCount); CountDownLatch consumeLatchSecondWave = new CountDownLatch(messageCount * 2); - ConsumerBuilder consumerBuilder = environment.consumerBuilder().stream(s); - - MessageHandlingListenerConsumerBuilderAccessor messageHandlingListenerConsumerBuilderAccessor = consumerBuilder + ConsumerBuilder consumerBuilder = environment.consumerBuilder() + .stream(s) .asynchronousControlFlow(5); - MessageHandlingListener messageHandlingListener = messageHandlingListenerConsumerBuilderAccessor.messageHandlingListener(); - consumerBuilder = messageHandlingListenerConsumerBuilderAccessor.builder(); List messageContexts = new ArrayList<>(); @@ -740,7 +730,7 @@ void consumerWithAsyncFlowControlShouldKeepConsumingAfterDisruption( .messageHandler( (context, message) -> { if(shouldInstaConsume.get()) { - if(!messageHandlingListener.markHandled(context)) { + if(!context.markHandled()) { unhandledOnInstaConsume.set(true); } } else { @@ -759,7 +749,7 @@ void consumerWithAsyncFlowControlShouldKeepConsumingAfterDisruption( assertThat(consumer.isOpen()).isTrue(); shouldInstaConsume.set(true); - boolean allMarkedHandled = messageContexts.parallelStream().allMatch(messageHandlingListener::markHandled); + boolean allMarkedHandled = messageContexts.parallelStream().allMatch(MessageHandler.Context::markHandled); assertThat(allMarkedHandled).isTrue(); assertThat(consumeLatch.await(20, TimeUnit.SECONDS)).isTrue(); From 09f3cb7e929aa17195c5f9b1821100ff1654360f Mon Sep 17 00:00:00 2001 From: henry701 Date: Sat, 1 Jul 2023 13:28:12 -0300 Subject: [PATCH 13/14] Change lifecycle of ConsumerFlowControlStrategy to be 1-1 with SubscriptionTracker, removing subscription and stream name parameters --- .../stream/CallbackStreamDataHandler.java | 39 +-- .../AbstractConsumerFlowControlStrategy.java | 45 +-- .../flow/ConsumerFlowControlStrategy.java | 14 - .../ConsumerFlowControlStrategyBuilder.java | 7 +- .../com/rabbitmq/stream/flow/CreditAsker.java | 7 +- .../java/com/rabbitmq/stream/impl/Client.java | 54 +++- ...lientCallbackStreamDataHandlerAdapter.java | 134 -------- .../impl/ConsumerStatisticRecorder.java | 290 +++++++----------- .../stream/impl/ConsumersCoordinator.java | 11 +- ...cRecordingConsumerFlowControlStrategy.java | 49 ++- ...ptionAsyncConsumerFlowControlStrategy.java | 34 +- ...ynchronousConsumerFlowControlStrategy.java | 27 +- .../stream/impl/ConsumersCoordinatorTest.java | 9 +- 13 files changed, 248 insertions(+), 472 deletions(-) delete mode 100644 src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java diff --git a/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java b/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java index e043a1ef1f..e548946f82 100644 --- a/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java +++ b/src/main/java/com/rabbitmq/stream/CallbackStreamDataHandler.java @@ -1,56 +1,40 @@ package com.rabbitmq.stream; -import java.util.Map; - /** - * Exposes callbacks to handle events from a particular Stream connection, + * Exposes callbacks to handle events from a particular Stream subscription, * with specific names for methods and no connection-oriented parameter. */ public interface CallbackStreamDataHandler { - default void handlePublishConfirm(byte publisherId, long publishingId) { - // No-op by default - } - - default void handlePublishError(byte publisherId, long publishingId, short errorCode) { + default void handleChunk(long offset, long messageCount, long dataSize) { // No-op by default } - default void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { + default void handleMessage(long offset, long chunkTimestamp, long committedChunkId, Message message) { // No-op by default } - default void handleMessage(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { + default void handleCreditNotification(short responseCode) { // No-op by default } - default void handleCreditNotification(byte subscriptionId, short responseCode) { + default void handleConsumerUpdate(boolean active) { // No-op by default } - default void handleConsumerUpdate(byte subscriptionId, boolean active) { - // No-op by default - } - - default void handleMetadata(String stream, short code) { + default void handleMetadata(short code) { // No-op by default } /** * Callback for handling a stream subscription. * - * @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 * @param isInitialSubscription Whether this subscription is an initial subscription * or a recovery for an existing subscription */ default void handleSubscribe( - byte subscriptionId, - String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties, boolean isInitialSubscription ) { // No-op by default @@ -58,10 +42,15 @@ default void handleSubscribe( /** * 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 + default void handleUnsubscribe() { + if(this instanceof AutoCloseable) { + try { + ((AutoCloseable) this).close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + } } } diff --git a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java index 75ccdddd6e..b810ed6904 100644 --- a/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/AbstractConsumerFlowControlStrategy.java @@ -1,44 +1,47 @@ package com.rabbitmq.stream.flow; import java.util.Objects; -import java.util.function.Supplier; /** - * Abstract class for Consumer Flow Control Strategies which keeps a cached lazily-initialized - * {@link CreditAsker} ready for retrieval by its inheritors. + * Abstract class for Consumer Flow Control Strategies which keeps a + * {@link CreditAsker creditAsker} field ready for retrieval by its inheritors. */ public abstract class AbstractConsumerFlowControlStrategy implements ConsumerFlowControlStrategy { - private final Supplier creditAskerSupplier; - private volatile CreditAsker lastRetrievedCreditAsker; + private final String identifier; + private final CreditAsker creditAsker; - protected AbstractConsumerFlowControlStrategy(Supplier creditAskerSupplier) { - this.creditAskerSupplier = Objects.requireNonNull(creditAskerSupplier, "creditAskerSupplier"); + protected AbstractConsumerFlowControlStrategy(String identifier, CreditAsker creditAsker) { + this.identifier = identifier; + this.creditAsker = Objects.requireNonNull(creditAsker, "creditAsker"); } - protected CreditAsker mandatoryCreditAsker() { - CreditAsker localSupplied = creditAskerSupplier.get(); - if(localSupplied == null) { - throw new IllegalStateException("Requested CreditAsker but it's not yet available! Supplier: " + this.creditAskerSupplier); - } - this.lastRetrievedCreditAsker = localSupplied; - return localSupplied; + public CreditAsker getCreditAsker() { + return creditAsker; } - public CreditAsker getLastRetrievedCreditAsker() { - return lastRetrievedCreditAsker; + public String getIdentifier() { + return identifier; } - public Supplier getCreditAskerSupplier() { - return creditAskerSupplier; + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + AbstractConsumerFlowControlStrategy that = (AbstractConsumerFlowControlStrategy) o; + return Objects.equals(identifier, that.identifier) && Objects.equals(creditAsker, that.creditAsker); + } + + @Override + public int hashCode() { + return Objects.hash(identifier, creditAsker); } @Override public String toString() { return "AbstractConsumerFlowControlStrategy{" + - "creditAskerSupplier=" + creditAskerSupplier + - ", lastRetrievedCreditAsker=" + lastRetrievedCreditAsker + + "identifier='" + identifier + '\'' + + ", creditAsker=" + creditAsker + '}'; } - } diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java index 953d7b780d..33113294cf 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategy.java @@ -3,8 +3,6 @@ import com.rabbitmq.stream.CallbackStreamDataHandler; import com.rabbitmq.stream.OffsetSpecification; -import java.util.Map; - /** * A built and configured flow control strategy for consumers. * Implementations may freely implement reactions to the various client callbacks. @@ -21,34 +19,22 @@ public interface ConsumerFlowControlStrategy extends CallbackStreamDataHandler { * Either this variant or {@link CallbackStreamDataHandler#handleSubscribe} should be called, NOT both. *

* - * @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 * @param isInitialSubscription Whether this subscription is an initial subscription * or a recovery for an existing subscription * @return The initial credits that should be granted to this new subscription */ int handleSubscribeReturningInitialCredits( - byte subscriptionId, - String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties, boolean isInitialSubscription ); @Override default void handleSubscribe( - byte subscriptionId, - String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties, boolean isInitialSubscription) { handleSubscribeReturningInitialCredits( - subscriptionId, - stream, offsetSpecification, - subscriptionProperties, isInitialSubscription ); } diff --git a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java index 3b5c4444b3..53ead20daf 100644 --- a/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java +++ b/src/main/java/com/rabbitmq/stream/flow/ConsumerFlowControlStrategyBuilder.java @@ -2,8 +2,6 @@ import com.rabbitmq.stream.ConsumerBuilder; -import java.util.function.Supplier; - /** * Fluent builder for a {@link ConsumerFlowControlStrategyBuilderFactory}. * One instance of this is set per {@link com.rabbitmq.stream.Consumer}. @@ -15,8 +13,9 @@ public interface ConsumerFlowControlStrategyBuilder} for retrieving the instance (which may be lazily initialized). + * @param identifier A {@link String} to uniquely identify the built instance and/or its subscription. + * @param creditAsker {@link CreditAsker} for asking for credits. * @return {@link T} the built {@link ConsumerFlowControlStrategy} */ - T build(Supplier creditAskerSupplier); + T build(String identifier, CreditAsker creditAsker); } diff --git a/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java b/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java index 0a906597ad..7437053b83 100644 --- a/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java +++ b/src/main/java/com/rabbitmq/stream/flow/CreditAsker.java @@ -5,10 +5,9 @@ public interface CreditAsker { /** * Asks for credits for a given subscription. - * @param subscriptionId the subscription ID - * @param credit how many credits to ask for - * @throws IllegalArgumentException if credits are below 0 or above {@link Short#MAX_VALUE} + * @param credits How many credits to ask for + * @throws IllegalArgumentException If credits are below 0 or above {@link Short#MAX_VALUE} */ - void credit(byte subscriptionId, int credit); + void credit(int credits); } diff --git a/src/main/java/com/rabbitmq/stream/impl/Client.java b/src/main/java/com/rabbitmq/stream/impl/Client.java index ffc91ed963..d4bb42b93b 100644 --- a/src/main/java/com/rabbitmq/stream/impl/Client.java +++ b/src/main/java/com/rabbitmq/stream/impl/Client.java @@ -13,25 +13,47 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; -import com.rabbitmq.stream.*; +import com.rabbitmq.stream.AuthenticationFailureException; +import com.rabbitmq.stream.ByteCapacity; +import com.rabbitmq.stream.ChunkChecksum; +import com.rabbitmq.stream.Codec; import com.rabbitmq.stream.Codec.EncodedMessage; +import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.Environment; +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.MessageBuilder; +import com.rabbitmq.stream.OffsetSpecification; +import com.rabbitmq.stream.Producer; import com.rabbitmq.stream.StreamCreator.LeaderLocator; +import com.rabbitmq.stream.StreamException; import com.rabbitmq.stream.compression.Compression; import com.rabbitmq.stream.compression.CompressionCodec; import com.rabbitmq.stream.compression.CompressionCodecFactory; -import com.rabbitmq.stream.flow.CreditAsker; import com.rabbitmq.stream.impl.Client.ShutdownContext.ShutdownReason; import com.rabbitmq.stream.impl.ServerFrameHandler.FrameHandler; import com.rabbitmq.stream.impl.ServerFrameHandler.FrameHandlerInfo; -import com.rabbitmq.stream.impl.Utils.*; import com.rabbitmq.stream.metrics.MetricsCollector; import com.rabbitmq.stream.metrics.NoOpMetricsCollector; -import com.rabbitmq.stream.sasl.*; +import com.rabbitmq.stream.sasl.CredentialsProvider; +import com.rabbitmq.stream.sasl.DefaultSaslConfiguration; +import com.rabbitmq.stream.sasl.DefaultUsernamePasswordCredentialsProvider; +import com.rabbitmq.stream.sasl.SaslConfiguration; +import com.rabbitmq.stream.sasl.SaslMechanism; +import com.rabbitmq.stream.sasl.UsernamePasswordCredentialsProvider; import io.netty.bootstrap.Bootstrap; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufAllocator; import io.netty.buffer.ByteBufOutputStream; -import io.netty.channel.*; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInboundHandlerAdapter; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.ChannelOption; +import io.netty.channel.ChannelOutboundHandlerAdapter; +import io.netty.channel.ChannelPromise; +import io.netty.channel.ConnectTimeoutException; +import io.netty.channel.EventLoopGroup; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioSocketChannel; @@ -57,8 +79,23 @@ import java.net.SocketAddress; import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.util.*; -import java.util.concurrent.*; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -83,7 +120,7 @@ *

People wanting very fine control over their interaction with the broker can use {@link Client} * but at their own risk. */ -public class Client implements CreditAsker, AutoCloseable { +public class Client implements AutoCloseable { public static final int DEFAULT_PORT = 5552; public static final int DEFAULT_TLS_PORT = 5551; @@ -1004,7 +1041,6 @@ public MessageBuilder messageBuilder() { return this.codec.messageBuilder(); } - @Override public void credit(byte subscriptionId, int credit) { if (credit < 0 || credit > Short.MAX_VALUE) { throw new IllegalArgumentException("Credit value must be between 0 and " + Short.MAX_VALUE); diff --git a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java b/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java deleted file mode 100644 index 6f7268e9fb..0000000000 --- a/src/main/java/com/rabbitmq/stream/impl/ClientCallbackStreamDataHandlerAdapter.java +++ /dev/null @@ -1,134 +0,0 @@ -package com.rabbitmq.stream.impl; - -import com.rabbitmq.stream.CallbackStreamDataHandler; -import com.rabbitmq.stream.Message; -import com.rabbitmq.stream.OffsetSpecification; - -import java.util.Map; - -public class ClientCallbackStreamDataHandlerAdapter implements - CallbackStreamDataHandler, - Client.PublishConfirmListener, - Client.PublishErrorListener, - Client.ChunkListener, - Client.MessageListener, - Client.CreditNotification, - Client.ConsumerUpdateListener, - Client.ShutdownListener, - Client.MetadataListener { - - private final CallbackStreamDataHandler callbackStreamDataHandler; - - public ClientCallbackStreamDataHandlerAdapter(CallbackStreamDataHandler callbackStreamDataHandler) { - this.callbackStreamDataHandler = callbackStreamDataHandler; - } - - @Override - public void handle(byte publisherId, long publishingId) { - this.handlePublishConfirm(publisherId, publishingId); - } - - @Override - public void handlePublishConfirm(byte publisherId, long publishingId) { - this.callbackStreamDataHandler.handlePublishConfirm(publisherId, publishingId); - } - - @Override - public void handle(byte publisherId, long publishingId, short errorCode) { - this.handlePublishError(publisherId, publishingId, errorCode); - } - - @Override - public void handlePublishError(byte publisherId, long publishingId, short errorCode) { - this.callbackStreamDataHandler.handlePublishError(publisherId, publishingId, errorCode); - } - - @Override - public void handle(Client client, byte subscriptionId, long offset, long messageCount, long dataSize) { - this.handleChunk(subscriptionId, offset, messageCount, dataSize); - } - - @Override - public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - this.callbackStreamDataHandler.handleChunk(subscriptionId, offset, messageCount, dataSize); - } - - @Override - public void handle(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { - this.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); - } - - @Override - public void handleMessage(byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message) { - this.callbackStreamDataHandler.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); - } - - @Override - public void handle(byte subscriptionId, short responseCode) { - this.handleCreditNotification(subscriptionId, responseCode); - } - - @Override - public void handleCreditNotification(byte subscriptionId, short responseCode) { - this.callbackStreamDataHandler.handleCreditNotification(subscriptionId, responseCode); - } - - @Override - public OffsetSpecification handle(Client client, byte subscriptionId, boolean active) { - this.handleConsumerUpdate(subscriptionId, active); - return null; - } - - @Override - public void handleConsumerUpdate(byte subscriptionId, boolean active) { - this.callbackStreamDataHandler.handleConsumerUpdate(subscriptionId, active); - } - - @Override - public void handle(Client.ShutdownContext shutdownContext) { - this.handleShutdown(shutdownContext); - } - - public void handleShutdown(Client.ShutdownContext shutdownContext) { - if(callbackStreamDataHandler instanceof AutoCloseable) { - try { - ((AutoCloseable) callbackStreamDataHandler).close(); - } catch (Exception e) { - throw new RuntimeException(e); - } - } - } - - @Override - public void handle(String stream, short code) { - this.handleMetadata(stream, code); - } - - @Override - public void handleMetadata(String stream, short code) { - this.callbackStreamDataHandler.handleMetadata(stream, code); - } - - @Override - public void handleSubscribe( - byte subscriptionId, - String stream, - OffsetSpecification offsetSpecification, - Map subscriptionProperties, - boolean isInitialSubscription - ) { - this.callbackStreamDataHandler.handleSubscribe( - subscriptionId, - stream, - offsetSpecification, - subscriptionProperties, - isInitialSubscription - ); - } - - @Override - public void handleUnsubscribe(byte subscriptionId) { - this.callbackStreamDataHandler.handleUnsubscribe(subscriptionId); - } - -} diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java index d12e53379b..aac9466367 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumerStatisticRecorder.java @@ -14,80 +14,53 @@ import java.util.Iterator; import java.util.Map; import java.util.NavigableMap; -import java.util.Set; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; public class ConsumerStatisticRecorder implements CallbackStreamDataHandler, MessageHandlingListener { private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerStatisticRecorder.class); - private final ConcurrentMap> streamNameToSubscriptionIdMap = new ConcurrentHashMap<>(); - private final ConcurrentMap subscriptionStatisticsMap = new ConcurrentHashMap<>(); + private final String identifier; + private final AtomicReference subscriptionStatistics = new AtomicReference<>(); + + public ConsumerStatisticRecorder(String identifier) { + this.identifier = identifier; + } @Override public void handleSubscribe( - byte subscriptionId, - String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties, boolean isInitialSubscription ) { - this.streamNameToSubscriptionIdMap.compute( - stream, - (k, v) -> { - if(v == null) { - v = Collections.newSetFromMap(new ConcurrentHashMap<>()); - } - boolean isNewElement = v.add(subscriptionId); - if(isInitialSubscription && !isNewElement) { - LOGGER.warn( - "handleSubscribe called for stream that already had same associated subscription! " + - "subscriptionId={} stream={} offsetSpecification={}", - subscriptionId, - stream, - offsetSpecification - ); - } - return v; - } - ); - this.subscriptionStatisticsMap.compute( - subscriptionId, - (k, v) -> { - if(v != null && isInitialSubscription) { - LOGGER.warn( - "handleSubscribe called for subscription that already exists! " + - "subscriptionId={} stream={} offsetSpecification={}", - subscriptionId, - stream, - offsetSpecification - ); - } - // Only overwrite if is a de-facto initial subscription - if(v == null) { - return new SubscriptionStatistics( - subscriptionId, - stream, - offsetSpecification, - subscriptionProperties - ); - } - v.offsetSpecification = offsetSpecification; - v.pendingChunks.set(0); - v.subscriptionProperties = subscriptionProperties; - cleanupOldTrackingData(v); - return v; - } - ); + SubscriptionStatistics localSubscriptionStatistics = this.subscriptionStatistics.get(); + if(localSubscriptionStatistics == null) { + this.subscriptionStatistics.set(new SubscriptionStatistics(offsetSpecification)); + return; + } + if(isInitialSubscription) { + LOGGER.warn( + "handleSubscribe called for stream that already had same associated subscription! " + + "identifier={} offsetSpecification={}", + this.identifier, + offsetSpecification + ); + } + localSubscriptionStatistics.offsetSpecification = offsetSpecification; + localSubscriptionStatistics.pendingChunks.set(0); + cleanupOldTrackingData(localSubscriptionStatistics); } - private static void cleanupOldTrackingData(SubscriptionStatistics subscriptionStatistics) { + private void cleanupOldTrackingData(SubscriptionStatistics subscriptionStatistics) { if (!subscriptionStatistics.offsetSpecification.isOffset()) { - LOGGER.debug("Can't cleanup old tracking data: offsetSpecification is not an offset! {}", subscriptionStatistics.offsetSpecification); + LOGGER.debug("Can't cleanup old tracking data: offsetSpecification is not an offset! " + + "identifier={} offsetSpecification={}", + this.identifier, + subscriptionStatistics.offsetSpecification + ); return; } // Mark messages before the initial offset as handled @@ -113,89 +86,67 @@ private static void cleanupOldTrackingData(SubscriptionStatistics subscriptionSt } @Override - public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - this.subscriptionStatisticsMap.compute( - subscriptionId, - (k, v) -> { - if(v == null) { - LOGGER.warn( - "handleChunk called for subscription that does not exist! subscriptionId={} offset={}", - subscriptionId, - offset - ); - return null; - } - v.pendingChunks.decrementAndGet(); - v.unprocessedChunksByOffset.put(offset, new ChunkStatistics(offset, messageCount, dataSize)); - return v; - } - ); + public void handleChunk(long offset, long messageCount, long dataSize) { + SubscriptionStatistics localSubscriptionStatistics = this.subscriptionStatistics.get(); + if(localSubscriptionStatistics == null) { + LOGGER.warn( + "handleChunk called for subscription that does not exist! " + + "identifier={} offset={} messageCount={} dataSize={}", + this.identifier, + offset, + messageCount, + dataSize + ); + return; + } + localSubscriptionStatistics.pendingChunks.decrementAndGet(); + localSubscriptionStatistics.unprocessedChunksByOffset.put(offset, new ChunkStatistics(offset, messageCount, dataSize)); } @Override public void handleMessage( - byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message ) { - this.subscriptionStatisticsMap.compute( - subscriptionId, - (k, v) -> { - if(v == null) { - LOGGER.warn( - "handleMessage called for subscription that does not exist! subscriptionId={} offset={}", - subscriptionId, - offset - ); - return null; - } - NavigableMap subHeadMapByOffset = v.unprocessedChunksByOffset.headMap(offset, true); - Map.Entry lastOffsetToChunkEntry = subHeadMapByOffset.lastEntry(); - if(lastOffsetToChunkEntry == null) { - LOGGER.warn( - "handleMessage called but chunk was not found! subscriptionId={} offset={}", - subscriptionId, - offset - ); - return v; - } - ChunkStatistics chunkStatistics = lastOffsetToChunkEntry.getValue(); - chunkStatistics.unprocessedMessagesByOffset.put(offset, message); - return v; - } - ); + SubscriptionStatistics localSubscriptionStatistics = this.subscriptionStatistics.get(); + if(localSubscriptionStatistics == null) { + LOGGER.warn( + "handleMessage called for subscription that does not exist! " + + "identifier={} offset={} chunkTimestamp={} committedChunkId={}", + this.identifier, + offset, + chunkTimestamp, + committedChunkId + ); + return; + } + NavigableMap subHeadMapByOffset = localSubscriptionStatistics.unprocessedChunksByOffset.headMap(offset, true); + Map.Entry lastOffsetToChunkEntry = subHeadMapByOffset.lastEntry(); + if(lastOffsetToChunkEntry == null) { + LOGGER.warn( + "handleMessage called but chunk was not found! " + + "identifier={} offset={} chunkTimestamp={} committedChunkId={}", + this.identifier, + offset, + chunkTimestamp, + committedChunkId + ); + return; + } + ChunkStatistics chunkStatistics = lastOffsetToChunkEntry.getValue(); + chunkStatistics.unprocessedMessagesByOffset.put(offset, message); } @Override - public void handleUnsubscribe(byte subscriptionId) { - ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStatistics = this.subscriptionStatisticsMap.remove(subscriptionId); - if(subscriptionStatistics == null) { + public void handleUnsubscribe() { + ConsumerStatisticRecorder.SubscriptionStatistics localSubscriptionStatistics = this.subscriptionStatistics.getAndSet(null); + if(localSubscriptionStatistics == null) { LOGGER.warn( - "handleUnsubscribe called for subscriptionId that does not exist! subscriptionId={}", - subscriptionId + "handleUnsubscribe called for subscriptionId that does not exist! Identifier: {}", identifier ); - return; } - this.streamNameToSubscriptionIdMap.compute(subscriptionStatistics.stream, (k, v) -> { - if(v == null) { - LOGGER.warn( - "handleUnsubscribe called and stream name '{}' did not contain subscriptions!", - subscriptionStatistics.stream - ); - return null; - } - boolean removed = v.remove(subscriptionId); - if(!removed) { - LOGGER.warn( - "handleUnsubscribe called and stream name '{}' did not contain subscriptionId {}!", - subscriptionStatistics.stream, - subscriptionId - ); - } - return v.isEmpty() ? null : v; - }); } /** @@ -232,7 +183,10 @@ public boolean markHandled(AggregatedMessageStatistics aggregatedMessageStatisti long initialOffset = aggregatedMessageStatistics.subscriptionStatistics.offsetSpecification.getOffset(); // Old tracked message, should already be handled, probably a late acknowledgment of a defunct connection. if(aggregatedMessageStatistics.offset < initialOffset) { - LOGGER.debug("Old message registered as consumed. Message Offset: {}, Start Offset: {}", aggregatedMessageStatistics.offset, initialOffset); + LOGGER.debug("Old message registered as consumed. Identifier={} Message Offset: {}, Start Offset: {}", + this.identifier, + aggregatedMessageStatistics.offset, + initialOffset); return true; } } @@ -250,36 +204,19 @@ public boolean markHandled(AggregatedMessageStatistics aggregatedMessageStatisti return true; } - public AggregatedMessageStatistics retrieveStatistics(String stream, long offset) { - Set possibleSubscriptionIds = this.streamNameToSubscriptionIdMap.get(stream); - AggregatedMessageStatistics entry = null; - for (Byte subscriptionId : possibleSubscriptionIds) { - entry = retrieveStatistics(subscriptionId, offset); - if (entry == null) { - continue; - } - // We have all the info we need, we found the specific chunk. Stop right here - if (entry.chunkHeadMap != null && entry.chunkStatistics != null) { - return entry; - } - } - // Return the next-best result, because we might find the subscription but not the message - return entry; - } - - public AggregatedMessageStatistics retrieveStatistics(byte subscriptionId, long offset) { - SubscriptionStatistics subscriptionStatistics = this.subscriptionStatisticsMap.get(subscriptionId); - if (subscriptionStatistics == null) { + public AggregatedMessageStatistics retrieveStatistics(long offset) { + SubscriptionStatistics localSubscriptionStatistics = this.subscriptionStatistics.get(); + if (localSubscriptionStatistics == null) { return null; } - NavigableMap chunkStatisticsHeadMap = subscriptionStatistics.unprocessedChunksByOffset.headMap(offset, true); + NavigableMap chunkStatisticsHeadMap = localSubscriptionStatistics.unprocessedChunksByOffset.headMap(offset, true); Map.Entry messageEntry = chunkStatisticsHeadMap.lastEntry(); ChunkStatistics chunkStatistics = messageEntry == null ? null : messageEntry.getValue(); - return new AggregatedMessageStatistics(offset, subscriptionStatistics, chunkStatisticsHeadMap, chunkStatistics, messageEntry); + return new AggregatedMessageStatistics(offset, localSubscriptionStatistics, chunkStatisticsHeadMap, chunkStatistics, messageEntry); } public AggregatedMessageStatistics retrieveStatistics(MessageHandler.Context messageContext) { - return retrieveStatistics(messageContext.stream(), messageContext.offset()); + return retrieveStatistics(messageContext.offset()); } public static class AggregatedMessageStatistics { @@ -331,42 +268,18 @@ public long getOffset() { public static class SubscriptionStatistics { - private final byte subscriptionId; - private final String stream; private final AtomicInteger pendingChunks = new AtomicInteger(0); private OffsetSpecification offsetSpecification; - private Map subscriptionProperties; private final NavigableMap unprocessedChunksByOffset; - public SubscriptionStatistics( - byte subscriptionId, - String stream, - OffsetSpecification offsetSpecification, - Map subscriptionProperties - ) { - this(subscriptionId, stream, offsetSpecification, subscriptionProperties, new ConcurrentSkipListMap<>()); - } - - public SubscriptionStatistics( - byte subscriptionId, - String stream, - OffsetSpecification offsetSpecification, - Map subscriptionProperties, - NavigableMap unprocessedChunksByOffset - ) { - this.subscriptionId = subscriptionId; - this.stream = stream; - this.offsetSpecification = offsetSpecification; - this.subscriptionProperties = subscriptionProperties; - this.unprocessedChunksByOffset = unprocessedChunksByOffset; - } - - public byte getSubscriptionId() { - return subscriptionId; + public SubscriptionStatistics(OffsetSpecification offsetSpecification) { + this(offsetSpecification, new ConcurrentSkipListMap<>()); } - public String getStream() { - return stream; + public SubscriptionStatistics(OffsetSpecification offsetSpecification, + NavigableMap unprocessedChunksByOffset) { + this.offsetSpecification = offsetSpecification; + this.unprocessedChunksByOffset = unprocessedChunksByOffset; } public AtomicInteger getPendingChunks() { @@ -377,10 +290,6 @@ public OffsetSpecification getOffsetSpecification() { return offsetSpecification; } - public Map getSubscriptionProperties() { - return Collections.unmodifiableMap(subscriptionProperties); - } - public NavigableMap getUnprocessedChunksByOffset() { return Collections.unmodifiableNavigableMap(unprocessedChunksByOffset); } @@ -390,7 +299,7 @@ public NavigableMap getUnprocessedChunksByOffset() { public static class ChunkStatistics { private final long offset; - private AtomicLong processedMessages = new AtomicLong(); + private final AtomicLong processedMessages = new AtomicLong(); private final long messageCount; private final long dataSize; private final Map unprocessedMessagesByOffset; @@ -427,12 +336,19 @@ public boolean isDone() { } } - public Map> getStreamNameToSubscriptionIdMap() { - return Collections.unmodifiableMap(streamNameToSubscriptionIdMap); + public String getIdentifier() { + return identifier; } - public Map getSubscriptionStatisticsMap() { - return Collections.unmodifiableMap(subscriptionStatisticsMap); + public SubscriptionStatistics getSubscriptionStatistics() { + return subscriptionStatistics.get(); } + @Override + public String toString() { + return "ConsumerStatisticRecorder{" + + "identifier='" + identifier + '\'' + + ", subscriptionStatistics=" + subscriptionStatistics.get() + + '}'; + } } diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index b5bd339f4c..48263ce26b 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -418,7 +418,10 @@ private SubscriptionTracker( this.trackingClosingCallback = trackingClosingCallback; this.messageHandler = messageHandler; this.clientReference = new AtomicReference<>(); - this.consumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build(this.clientReference::get); + this.consumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build( + String.format("[stream=%s, subscriptionId=%s]", stream, String.valueOf(subscriptionIdInClient)), + credits -> this.clientReference.get().credit(this.subscriptionIdInClient, credits) + ); if (this.offsetTrackingReference == null) { this.subscriptionProperties = subscriptionProperties; } else { @@ -585,7 +588,6 @@ private ClientSubscriptionsManager( subscriptionTracker.offset = offset; subscriptionTracker.hasReceivedSomething = true; subscriptionTracker.consumerFlowControlStrategy.handleMessage( - subscriptionId, offset, chunkTimestamp, committedOffset, @@ -740,7 +742,7 @@ private ClientSubscriptionsManager( LOGGER.warn("Could not find stream subscription {} for chunk listener", subscriptionId); return; } - subscriptionTracker.consumerFlowControlStrategy.handleChunk(subscriptionId, offset, messageCount, dataSize); + subscriptionTracker.consumerFlowControlStrategy.handleChunk(offset, messageCount, dataSize); }; String connectionName = connectionNamingStrategy.apply(ClientConnectionType.CONSUMER); ClientFactoryContext clientFactoryContext = @@ -970,10 +972,7 @@ synchronized void add( checkNotClosed(); int initialCredits = subscriptionTracker.consumerFlowControlStrategy.handleSubscribeReturningInitialCredits( - subscriptionId, - subscriptionTracker.stream, subscriptionContext.offsetSpecification(), - subscriptionTracker.subscriptionProperties, isInitialSubscription ); final byte finalSubscriptionId = subscriptionId; diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java index 17089b2198..8c48e71e79 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/AbstractStatisticRecordingConsumerFlowControlStrategy.java @@ -11,10 +11,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.IntUnaryOperator; -import java.util.function.Supplier; /** * Abstract class that calls an instance of {@link ConsumerStatisticRecorder} and exposes it to child implementations @@ -26,10 +24,11 @@ public abstract class AbstractStatisticRecordingConsumerFlowControlStrategy private static final Logger LOGGER = LoggerFactory.getLogger(AbstractStatisticRecordingConsumerFlowControlStrategy.class); - protected final ConsumerStatisticRecorder consumerStatisticRecorder = new ConsumerStatisticRecorder(); + protected final ConsumerStatisticRecorder consumerStatisticRecorder; - protected AbstractStatisticRecordingConsumerFlowControlStrategy(Supplier creditAskerSupplier) { - super(creditAskerSupplier); + protected AbstractStatisticRecordingConsumerFlowControlStrategy(String identifier, CreditAsker creditAsker) { + super(identifier, creditAsker); + this.consumerStatisticRecorder = new ConsumerStatisticRecorder(identifier); } /** @@ -41,26 +40,19 @@ protected AbstractStatisticRecordingConsumerFlowControlStrategy(Supplier subscriptionProperties, boolean isInitialSubscription ) { this.consumerStatisticRecorder.handleSubscribe( - subscriptionId, - stream, offsetSpecification, - subscriptionProperties, isInitialSubscription ); } @Override - public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - super.handleChunk(subscriptionId, offset, messageCount, dataSize); + public void handleChunk(long offset, long messageCount, long dataSize) { + super.handleChunk(offset, messageCount, dataSize); this.consumerStatisticRecorder.handleChunk( - subscriptionId, offset, messageCount, dataSize @@ -69,15 +61,13 @@ public void handleChunk(byte subscriptionId, long offset, long messageCount, lon @Override public void handleMessage( - byte subscriptionId, long offset, long chunkTimestamp, long committedChunkId, Message message ) { - super.handleMessage(subscriptionId, offset, chunkTimestamp, committedChunkId, message); + super.handleMessage(offset, chunkTimestamp, committedChunkId, message); this.consumerStatisticRecorder.handleMessage( - subscriptionId, offset, chunkTimestamp, committedChunkId, @@ -86,24 +76,23 @@ public void handleMessage( } @Override - public void handleCreditNotification(byte subscriptionId, short responseCode) { - super.handleCreditNotification(subscriptionId, responseCode); - this.consumerStatisticRecorder.handleCreditNotification(subscriptionId, responseCode); + public void handleCreditNotification(short responseCode) { + super.handleCreditNotification(responseCode); + this.consumerStatisticRecorder.handleCreditNotification(responseCode); } @Override - public void handleUnsubscribe(byte subscriptionId) { - super.handleUnsubscribe(subscriptionId); - this.consumerStatisticRecorder.handleUnsubscribe(subscriptionId); + public void handleUnsubscribe() { + super.handleUnsubscribe(); + this.consumerStatisticRecorder.handleUnsubscribe(); } - protected int registerCredits(byte subscriptionId, IntUnaryOperator askedToAsk, boolean askForCredits) { + protected int registerCredits(IntUnaryOperator askedToAsk, boolean askForCredits) { AtomicInteger outerCreditsToAsk = new AtomicInteger(); ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStatistics = this.consumerStatisticRecorder - .getSubscriptionStatisticsMap() - .get(subscriptionId); + .getSubscriptionStatistics(); if(subscriptionStatistics == null) { - LOGGER.warn("Lost subscription {}, returning no credits. askForCredits={}", subscriptionId, askForCredits); + LOGGER.warn("Lost subscription, returning no credits. askForCredits={}", askForCredits); return 0; } subscriptionStatistics.getPendingChunks().updateAndGet(credits -> { @@ -113,10 +102,10 @@ protected int registerCredits(byte subscriptionId, IntUnaryOperator askedToAsk, }); int finalCreditsToAsk = outerCreditsToAsk.get(); if(askForCredits && finalCreditsToAsk > 0) { - LOGGER.debug("Asking for {} credits for subscriptionId {}", finalCreditsToAsk, subscriptionId); - mandatoryCreditAsker().credit(subscriptionId, finalCreditsToAsk); + LOGGER.debug("Asking for {} credits", finalCreditsToAsk); + getCreditAsker().credit(finalCreditsToAsk); } - LOGGER.debug("Returning {} credits for subscriptionId {} with askForCredits={}", finalCreditsToAsk, subscriptionId, askForCredits); + LOGGER.debug("Returning {} credits with askForCredits={}", finalCreditsToAsk, askForCredits); return finalCreditsToAsk; } diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java index 33e40273f9..c16929a314 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy.java @@ -10,9 +10,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.Map; import java.util.function.IntUnaryOperator; -import java.util.function.Supplier; /** * A flow control strategy that enforces a maximum amount of Inflight chunks per registered subscription. @@ -25,10 +23,11 @@ public class MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy extend private final int maximumSimultaneousChunksPerSubscription; public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( - Supplier creditAskerSupplier, + String identifier, + CreditAsker creditAsker, int maximumSimultaneousChunksPerSubscription ) { - super(creditAskerSupplier); + super(identifier, creditAsker); if(maximumSimultaneousChunksPerSubscription <= 0) { throw new IllegalArgumentException( "maximumSimultaneousChunksPerSubscription must be greater than 0. Was: " + maximumSimultaneousChunksPerSubscription @@ -39,43 +38,35 @@ public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( @Override public int handleSubscribeReturningInitialCredits( - byte subscriptionId, - String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties, boolean isInitialSubscription) { this.handleSubscribe( - subscriptionId, - stream, offsetSpecification, - subscriptionProperties, isInitialSubscription ); - return registerCredits(subscriptionId, getCreditAsker(subscriptionId), false); + return registerCredits(getCreditRegistererFunction(), false); } @Override protected void afterMarkHandledStateChanged( MessageHandler.Context messageContext, ConsumerStatisticRecorder.AggregatedMessageStatistics messageStatistics) { - byte subscriptionId = messageStatistics.getSubscriptionStatistics().getSubscriptionId(); - registerCredits(subscriptionId, getCreditAsker(subscriptionId), true); + registerCredits(getCreditRegistererFunction(), true); } - private IntUnaryOperator getCreditAsker(byte subscriptionId) { + private IntUnaryOperator getCreditRegistererFunction() { return pendingChunks -> { - int inProcessingChunks = extractInProcessingChunks(subscriptionId); + int inProcessingChunks = extractInProcessingChunks(); return Math.max(0, this.maximumSimultaneousChunksPerSubscription - (pendingChunks + inProcessingChunks)); }; } - private int extractInProcessingChunks(byte subscriptionId) { + private int extractInProcessingChunks() { int inProcessingChunks; ConsumerStatisticRecorder.SubscriptionStatistics subscriptionStats = this.consumerStatisticRecorder - .getSubscriptionStatisticsMap() - .get(subscriptionId); + .getSubscriptionStatistics(); if(subscriptionStats == null) { - LOGGER.warn("Subscription data not found while calculating credits to ask! subscriptionId: {}", subscriptionId); + LOGGER.warn("Subscription data not found while calculating credits to ask! Identifier: {}", this.getIdentifier()); inProcessingChunks = 0; } else { inProcessingChunks = subscriptionStats.getUnprocessedChunksByOffset().size(); @@ -98,9 +89,10 @@ public Builder(ConsumerBuilder consumerBuilder) { } @Override - public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { + public MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy build(String identifier, CreditAsker creditAsker) { return new MaximumChunksPerSubscriptionAsyncConsumerFlowControlStrategy( - creditAskerSupplier, + identifier, + creditAsker, this.maximumInflightChunksPerSubscription ); } diff --git a/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java b/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java index f652f6de99..f03849f53a 100644 --- a/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java +++ b/src/main/java/com/rabbitmq/stream/impl/flow/SynchronousConsumerFlowControlStrategy.java @@ -6,9 +6,6 @@ import com.rabbitmq.stream.flow.ConsumerFlowControlStrategyBuilder; import com.rabbitmq.stream.flow.CreditAsker; -import java.util.Map; -import java.util.function.Supplier; - /** * The default flow control strategy. * Requests a set amount of credits after each chunk arrives. @@ -20,26 +17,27 @@ public class SynchronousConsumerFlowControlStrategy extends AbstractConsumerFlow private final int initialCredits; private final int additionalCredits; - public SynchronousConsumerFlowControlStrategy(Supplier creditAskerSupplier, int initialCredits, int additionalCredits) { - super(creditAskerSupplier); + public SynchronousConsumerFlowControlStrategy( + String identifier, + CreditAsker creditAsker, + int initialCredits, + int additionalCredits) { + super(identifier, creditAsker); this.initialCredits = initialCredits; this.additionalCredits = additionalCredits; } @Override public int handleSubscribeReturningInitialCredits( - byte subscriptionId, - String stream, OffsetSpecification offsetSpecification, - Map subscriptionProperties, boolean isInitialSubscription ) { return this.initialCredits; } @Override - public void handleChunk(byte subscriptionId, long offset, long messageCount, long dataSize) { - mandatoryCreditAsker().credit(subscriptionId, this.additionalCredits); + public void handleChunk(long offset, long messageCount, long dataSize) { + getCreditAsker().credit(this.additionalCredits); } public static SynchronousConsumerFlowControlStrategy.Builder builder(ConsumerBuilder consumerBuilder) { @@ -59,8 +57,13 @@ public Builder(ConsumerBuilder consumerBuilder) { } @Override - public SynchronousConsumerFlowControlStrategy build(Supplier creditAskerSupplier) { - return new SynchronousConsumerFlowControlStrategy(creditAskerSupplier, this.initialCredits, this.additionalCredits); + public SynchronousConsumerFlowControlStrategy build(String identifier, CreditAsker creditAsker) { + return new SynchronousConsumerFlowControlStrategy( + identifier, + creditAsker, + this.initialCredits, + this.additionalCredits + ); } @Override diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index 8be1de6ac2..9b9837d048 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -1699,11 +1699,11 @@ void shouldCallConsumerFlowControlHandlers(Consumer co int numberOfInitialCreditsOnSubscribe = 7; - when(mockedConsumerFlowControlStrategy.handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(), anyMap(), anyBoolean())) + when(mockedConsumerFlowControlStrategy.handleSubscribeReturningInitialCredits(any(), anyBoolean())) .thenReturn(numberOfInitialCreditsOnSubscribe); ConsumerFlowControlStrategyBuilder mockedConsumerFlowControlStrategyBuilder = Mockito.mock(ConsumerFlowControlStrategyBuilder.class); - when(mockedConsumerFlowControlStrategyBuilder.build(any())).thenReturn(mockedConsumerFlowControlStrategy); + when(mockedConsumerFlowControlStrategyBuilder.build(any(), any())).thenReturn(mockedConsumerFlowControlStrategy); Runnable closingRunnable = coordinator.subscribe( @@ -1721,7 +1721,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co verify(client, times(1)) .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), eq(numberOfInitialCreditsOnSubscribe), anyMap()); verify(mockedConsumerFlowControlStrategy, times(1)) - .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap(), anyBoolean()); + .handleSubscribeReturningInitialCredits(any(OffsetSpecification.class), anyBoolean()); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(0) .isEqualTo(OffsetSpecification.next()); @@ -1739,7 +1739,6 @@ void shouldCallConsumerFlowControlHandlers(Consumer co message); verify(mockedConsumerFlowControlStrategy).handleMessage( - subscriptionIdCaptor.getValue(), lastReceivedOffset, 0, 0, @@ -1754,7 +1753,7 @@ void shouldCallConsumerFlowControlHandlers(Consumer co .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); verify(mockedConsumerFlowControlStrategy, times(2)) - .handleSubscribeReturningInitialCredits(anyByte(), anyString(), any(OffsetSpecification.class), anyMap(), anyBoolean()); + .handleSubscribeReturningInitialCredits(any(OffsetSpecification.class), anyBoolean()); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(1) From 17835b0e3b60b9b287cfc56f999c6d41d6b63114 Mon Sep 17 00:00:00 2001 From: henry701 Date: Tue, 4 Jul 2023 19:11:32 -0300 Subject: [PATCH 14/14] Add specific error for when client is not initialized --- .../rabbitmq/stream/impl/ConsumersCoordinator.java | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 48263ce26b..0f7f603dc6 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -418,9 +418,18 @@ private SubscriptionTracker( this.trackingClosingCallback = trackingClosingCallback; this.messageHandler = messageHandler; this.clientReference = new AtomicReference<>(); + String identifier = String.format("[stream=%s, subscriptionId=%s]", stream, String.valueOf(subscriptionIdInClient)); this.consumerFlowControlStrategy = consumerFlowControlStrategyBuilder.build( - String.format("[stream=%s, subscriptionId=%s]", stream, String.valueOf(subscriptionIdInClient)), - credits -> this.clientReference.get().credit(this.subscriptionIdInClient, credits) + identifier, + credits -> { + Client retrievedClient = this.clientReference.get(); + if(retrievedClient == null) { + LOGGER.error("Client is not initialized, cannot ask for credits! " + + "Asked for {} credits. Identifier: {}", credits, identifier); + return; + } + retrievedClient.credit(this.subscriptionIdInClient, credits); + } ); if (this.offsetTrackingReference == null) { this.subscriptionProperties = subscriptionProperties;