diff --git a/.github/workflows/test-pr.yml b/.github/workflows/test-pr.yml index ef52661fdb..a3fd4163c9 100644 --- a/.github/workflows/test-pr.yml +++ b/.github/workflows/test-pr.yml @@ -30,6 +30,8 @@ jobs: cache: 'maven' - name: Start broker run: ci/start-broker.sh + env: + RABBITMQ_IMAGE: pivotalrabbitmq/rabbitmq:pr-17185-otp28 - name: Test run: | ./mvnw test -Drabbitmqctl.bin=DOCKER:rabbitmq \ diff --git a/src/docs/asciidoc/api.adoc b/src/docs/asciidoc/api.adoc index eccce6c09d..057699356d 100644 --- a/src/docs/asciidoc/api.adoc +++ b/src/docs/asciidoc/api.adoc @@ -888,6 +888,7 @@ Useful when using an external store for offset tracking. |`flow#initialCredits` |Number of credits when the subscription is created. Increase for higher throughput at the expense of memory usage. +Accepts a `ByteCapacity` instead of an `int` to use a <> rather than a number of chunks (requires broker support). |10 |`flow#strategy` @@ -1184,6 +1185,46 @@ No calling it will stop the dispatching of messages. Whether the method is idempotent depends on the flow strategy implementation. Apart from the default one, the implementations the library provides does not make `processed()` idempotent. +[[byte-based-flow-control]] +===== Byte-Based Flow Control + +WARNING: Byte-based flow control requires a broker new enough to support the corresponding +version of the `Subscribe` and `Credit` commands. +The client fails fast with an exception when creating a consumer with a byte-based strategy +against a broker that does not support it, it does not fall back to chunk-based flow control. + +The credit unit for the strategies described above is the chunk: 1 credit lets the broker send 1 +more chunk. Chunks can vary widely in size, so a chunk-based window does not directly translate +into a memory bound on the consumer side. `creditOnChunkArrival`, `creditWhenHalfMessagesProcessed`, +and `creditOnProcessedMessageCount` all accept a `ByteCapacity` window instead of a number of +chunks, so the client can express the same arrival-versus-feedback trade-off in terms of bytes: + +.Setting a byte-based consumer flow control strategy +[source,java,indent=0] +-------- +include::{test-examples}/ConsumerUsage.java[tag=flow-control-byte-based] +-------- +<1> Set a byte-based flow control strategy with a 512 KB window +<2> Make sure to call `Context#processed()` + +* `creditOnChunkArrival(ByteCapacity)` grants the whole window back as soon as a chunk arrives, so +the number of bytes in flight is bounded, but the number of bytes received and not yet processed +is unbounded if processing falls behind. +* `creditWhenHalfMessagesProcessed(ByteCapacity)` and `creditOnProcessedMessageCount(ByteCapacity, +double)` grant credit as messages are processed, so the number of bytes received and not yet +processed stays bounded by roughly the window plus the size of the largest chunk received so far. +This is the feedback property that distinguishes them from `creditOnChunkArrival`. +* the broker may deliver a chunk larger than the window: a consumer never stalls just because its +window is smaller than the next chunk, it only ever falls behind by that one chunk. +* a byte-based strategy must eventually grant credit for every byte the broker charged for a +chunk (`Context#chunkByteCount()`), whether on the chunk's arrival or on the processing of its +messages. Granting less would make the subscription's credit drift down and eventually stall it. +This only matters for a custom `ConsumerFlowStrategy`, the strategies described above already do +it. +* the broker resumes sending only once about half the window has been granted back as credit, so +size the window to roughly twice the amount of data the application is willing to hold in memory +at once. + [[single-active-consumer]] ==== Single Active Consumer diff --git a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java index 9f30457c4f..6f7347a76b 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerBuilder.java @@ -1,4 +1,4 @@ -// Copyright (c) 2020-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2020-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -274,6 +274,23 @@ interface FlowConfiguration { */ FlowConfiguration initialCredits(int initialCredits); + /** + * The initial credits for the subscription, as a window in bytes. + * + *

This calls uses {@link ConsumerFlowStrategy#creditOnChunkArrival(ByteCapacity)}. + * + *

The broker may exceed the window by up to one chunk: a chunk larger than the window is + * delivered anyway, so that a consumer with a small window still makes forward progress. + * + *

Requires a broker supporting {@code Subscribe} version 2. + * + * @param initialCredits the initial credit window, in bytes + * @return this configuration instance + * @since 1.11.0 + * @see ConsumerFlowStrategy#creditOnChunkArrival(ByteCapacity) + */ + FlowConfiguration initialCredits(ByteCapacity initialCredits); + /** * Flow strategy to use * @@ -283,9 +300,12 @@ interface FlowConfiguration { * @see ConsumerFlowStrategy * @see ConsumerFlowStrategy#creditOnChunkArrival() * @see ConsumerFlowStrategy#creditOnChunkArrival(int) + * @see ConsumerFlowStrategy#creditOnChunkArrival(ByteCapacity) * @see ConsumerFlowStrategy#creditWhenHalfMessagesProcessed() * @see ConsumerFlowStrategy#creditWhenHalfMessagesProcessed(int) + * @see ConsumerFlowStrategy#creditWhenHalfMessagesProcessed(ByteCapacity) * @see ConsumerFlowStrategy#creditOnProcessedMessageCount(int, double) + * @see ConsumerFlowStrategy#creditOnProcessedMessageCount(ByteCapacity, double) */ FlowConfiguration strategy(ConsumerFlowStrategy strategy); diff --git a/src/main/java/com/rabbitmq/stream/ConsumerFlowStrategy.java b/src/main/java/com/rabbitmq/stream/ConsumerFlowStrategy.java index 804dca3ad4..545da8c570 100644 --- a/src/main/java/com/rabbitmq/stream/ConsumerFlowStrategy.java +++ b/src/main/java/com/rabbitmq/stream/ConsumerFlowStrategy.java @@ -1,4 +1,4 @@ -// Copyright (c) 2023-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2023-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -31,6 +31,17 @@ * ideal solution, it depends on the use cases and several parameters (processing time, network, * etc). * + *

Custom implementations must never make a credit release depend on the arrival of a + * later chunk. Releasing when a chunk arrives, or when its own messages are processed, is + * safe: the trigger has already happened by the time the credit is due. Waiting for a chunk that + * has not arrived yet is not safe, because the broker may be blocked precisely because that chunk + * was never sent. + * + *

Credit is usually expressed in chunks, but a strategy can express it in bytes instead, see + * {@link CreditUnit}. Byte-based credit still follows a chunk granularity: a chunk is delivered in + * full even if it exceeds the outstanding byte credit, so a consumer never stalls just because its + * window is smaller than the next chunk. + * *

This is an experimental API, subject to change. * * @since 0.12.0 @@ -61,6 +72,27 @@ public interface ConsumerFlowStrategy { */ MessageProcessedCallback start(Context context); + /** The unit a subscription's credit is expressed in. */ + enum CreditUnit { + /** 1 credit lets the broker send 1 more chunk, whatever its size. */ + CHUNK, + /** + * Credit is a number of bytes; the client must eventually grant back every byte the broker + * charged for a chunk. Requires broker support, see {@link + * com.rabbitmq.stream.ConsumerBuilder.FlowConfiguration#initialCredits(ByteCapacity)}. + */ + BYTE + } + + /** + * The unit this strategy's credit is expressed in. + * + *

Defaults to {@link CreditUnit#CHUNK}. + * + * @return the credit unit + */ + CreditUnit unit(); + /** Chunk context. */ interface Context { @@ -70,6 +102,14 @@ interface Context { *

{@link ConsumerFlowStrategy} implementation should always provide 1 credit for a given * chunk. * + *

credits counts chunks in both units. For a byte-based subscription, the + * client grants the bytes of the corresponding chunks as credit, not the raw credits + * value. + * + *

Implementations must never call this method for a chunk based on the arrival of a + * later chunk, only on the arrival of the chunk itself or the processing of its own + * messages, see {@link ConsumerFlowStrategy}. + * * @param credits the number of credits provided, usually 1 */ void credits(int credits); @@ -87,6 +127,15 @@ interface Context { * @return offset of the first message in the chunk (chunk ID) */ long chunkId(); + + /** + * The cost the broker charged for the chunk, in bytes. + * + *

This is what a byte-based subscription must eventually grant back as credit. + * + * @return the chunk cost, in bytes + */ + long chunkByteCount(); } /** Behavior for {@link MessageHandler.Context#processed()} calls. */ @@ -134,6 +183,22 @@ static ConsumerFlowStrategy creditOnChunkArrival(int initialCredits) { return new CreditOnChunkArrivalConsumerFlowStrategy(initialCredits); } + /** + * Strategy that provides a byte window as initial credits and a credit on each new chunk. + * + *

Calls to {@link MessageHandler.Context#processed()} are ignored. + * + *

Requires a broker supporting {@code Subscribe} version 2. + * + * @param window initial credit window, in bytes + * @return flow strategy + * @see com.rabbitmq.stream.ConsumerBuilder.FlowConfiguration#initialCredits(ByteCapacity) + */ + static ConsumerFlowStrategy creditOnChunkArrival(ByteCapacity window) { + return new CreditOnChunkArrivalConsumerFlowStrategy( + windowToInitialCredits(window), CreditUnit.BYTE); + } + /** * Strategy that provides 10 initial credits and a credit when half of the chunk messages are * processed. @@ -162,6 +227,23 @@ static ConsumerFlowStrategy creditWhenHalfMessagesProcessed(int initialCredits) return creditOnProcessedMessageCount(initialCredits, 0.5); } + /** + * Strategy that provides a byte window as initial credits and a credit when half of the chunk + * messages are processed. + * + *

Make sure to call {@link MessageHandler.Context#processed()} on every message when using + * this strategy, otherwise the broker may stop sending messages to the consumer. + * + *

Requires a broker supporting {@code Subscribe} version 2. + * + * @param window initial credit window, in bytes + * @return flow strategy + * @see com.rabbitmq.stream.ConsumerBuilder.FlowConfiguration#initialCredits(ByteCapacity) + */ + static ConsumerFlowStrategy creditWhenHalfMessagesProcessed(ByteCapacity window) { + return creditOnProcessedMessageCount(window, 0.5); + } + /** * Strategy that provides the specified number of initial credits and a credit when the specified * ratio of the chunk messages are processed. @@ -176,6 +258,36 @@ static ConsumerFlowStrategy creditOnProcessedMessageCount(int initialCredits, do return new MessageCountConsumerFlowStrategy(initialCredits, ratio); } + /** + * Strategy that provides a byte window as initial credits and a credit when the specified ratio + * of the chunk messages are processed. + * + *

Make sure to call {@link MessageHandler.Context#processed()} on every message when using + * this strategy, otherwise the broker may stop sending messages to the consumer. + * + *

Requires a broker supporting {@code Subscribe} version 2. + * + * @param window initial credit window, in bytes + * @param ratio ratio of messages to process before providing credits + * @return flow strategy + * @see com.rabbitmq.stream.ConsumerBuilder.FlowConfiguration#initialCredits(ByteCapacity) + */ + static ConsumerFlowStrategy creditOnProcessedMessageCount(ByteCapacity window, double ratio) { + return new MessageCountConsumerFlowStrategy( + windowToInitialCredits(window), ratio, CreditUnit.BYTE); + } + + private static int windowToInitialCredits(ByteCapacity window) { + if (window == null || window.compareTo(ByteCapacity.B(0)) <= 0) { + throw new IllegalArgumentException("The window must be positive"); + } + if (window.compareTo(ByteCapacity.B(Integer.MAX_VALUE)) > 0) { + throw new IllegalArgumentException( + "The window must be at most " + Integer.MAX_VALUE + " bytes"); + } + return (int) window.toBytes(); + } + /** * Strategy that provides the specified number of initial credits and n credits every * n chunks. @@ -189,6 +301,14 @@ static ConsumerFlowStrategy creditOnProcessedMessageCount(int initialCredits, do * *

Calls to {@link MessageHandler.Context#processed()} are ignored. * + *

This strategy has no byte-based variant and always uses {@link CreditUnit#CHUNK}: it + * releases credit once n chunks have arrived, which makes the release of the last + * n - 1 chunks depend on the arrival of a chunk that has not happened yet, breaking + * the contract described in {@link ConsumerFlowStrategy}. In chunk mode the residual is bounded + * by n chunks, which the constructor check below keeps under control; in byte mode + * the residual would be an unbounded number of bytes, since chunk sizes are not known when the + * consumer is built. + * * @param initialCredits number of initial credits * @param n number of chunks and number of credits * @return flow strategy @@ -234,6 +354,11 @@ public int initialCredits() { return this.initialCredits; } + @Override + public CreditUnit unit() { + return CreditUnit.CHUNK; + } + @Override public MessageProcessedCallback start(Context context) { if (chunkCount.incrementAndGet() % n == 0) { @@ -253,9 +378,15 @@ final class CreditOnChunkArrivalConsumerFlowStrategy implements ConsumerFlowStra private static final MessageProcessedCallback CALLBACK = v -> {}; private final int initialCredits; + private final CreditUnit unit; private CreditOnChunkArrivalConsumerFlowStrategy(int initialCredits) { + this(initialCredits, CreditUnit.CHUNK); + } + + CreditOnChunkArrivalConsumerFlowStrategy(int initialCredits, CreditUnit unit) { this.initialCredits = initialCredits; + this.unit = unit; } @Override @@ -263,6 +394,11 @@ public int initialCredits() { return this.initialCredits; } + @Override + public CreditUnit unit() { + return this.unit; + } + @Override public MessageProcessedCallback start(Context context) { context.credits(1); @@ -281,10 +417,16 @@ final class MessageCountConsumerFlowStrategy implements ConsumerFlowStrategy { private final int initialCredits; private final double ratio; + private final CreditUnit unit; private MessageCountConsumerFlowStrategy(int initialCredits, double ratio) { + this(initialCredits, ratio, CreditUnit.CHUNK); + } + + MessageCountConsumerFlowStrategy(int initialCredits, double ratio, CreditUnit unit) { this.initialCredits = initialCredits; this.ratio = ratio; + this.unit = unit; } @Override @@ -292,6 +434,11 @@ public int initialCredits() { return this.initialCredits; } + @Override + public CreditUnit unit() { + return this.unit; + } + @Override public MessageProcessedCallback start(Context context) { long l = (long) (context.messageCount() * ratio); diff --git a/src/main/java/com/rabbitmq/stream/impl/Client.java b/src/main/java/com/rabbitmq/stream/impl/Client.java index ee42892bcb..eab1f26d6b 100644 --- a/src/main/java/com/rabbitmq/stream/impl/Client.java +++ b/src/main/java/com/rabbitmq/stream/impl/Client.java @@ -64,6 +64,7 @@ import com.rabbitmq.stream.Codec; import com.rabbitmq.stream.Codec.EncodedMessage; import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; import com.rabbitmq.stream.Environment; import com.rabbitmq.stream.Message; import com.rabbitmq.stream.MessageBuilder; @@ -254,6 +255,7 @@ public long applyAsLong(Object value) { private final AtomicReference shutdownReason = new AtomicReference<>(); private final Runnable streamStatsCommandVersionsCheck; private final boolean filteringSupported; + private final boolean byteCreditSupported; private final Runnable superStreamManagementCommandVersionsCheck; private final Runnable resolveOffsetSpecCommandVersionsCheck; private final Registration credentialsRegistration; @@ -495,13 +497,14 @@ public void initChannel(SocketChannel ch) { // until now) this.channel .pipeline() - .replace( - NETTY_HANDLER_FRAME_DECODER, NETTY_HANDLER_FRAME_DECODER, frameDecoder()); + .replace(NETTY_HANDLER_FRAME_DECODER, NETTY_HANDLER_FRAME_DECODER, frameDecoder()); Set supportedCommands = maybeExchangeCommandVersions(); AtomicBoolean streamStatsSupported = new AtomicBoolean(false); AtomicBoolean filteringSupportedReference = new AtomicBoolean(false); AtomicBoolean superStreamManagementSupported = new AtomicBoolean(false); AtomicBoolean resolveOffsetSpecSupported = new AtomicBoolean(false); + AtomicBoolean subscribeVersion2Supported = new AtomicBoolean(false); + AtomicBoolean creditVersion2Supported = new AtomicBoolean(false); supportedCommands.forEach( c -> { if (c.getKey() == COMMAND_STREAM_STATS) { @@ -516,6 +519,12 @@ public void initChannel(SocketChannel ch) { if (c.getKey() == COMMAND_RESOLVE_OFFSET_SPEC) { resolveOffsetSpecSupported.set(true); } + if (c.getKey() == COMMAND_SUBSCRIBE && c.getMaxVersion() >= VERSION_2) { + subscribeVersion2Supported.set(true); + } + if (c.getKey() == COMMAND_CREDIT && c.getMaxVersion() >= VERSION_2) { + creditVersion2Supported.set(true); + } }); this.streamStatsCommandVersionsCheck = streamStatsSupported.get() @@ -525,6 +534,7 @@ public void initChannel(SocketChannel ch) { "QueryStreamInfo is available only on RabbitMQ 3.11 or more."); }; this.filteringSupported = filteringSupportedReference.get(); + this.byteCreditSupported = subscribeVersion2Supported.get() && creditVersion2Supported.get(); this.superStreamManagementCommandVersionsCheck = superStreamManagementSupported.get() ? () -> {} @@ -1206,18 +1216,38 @@ public MessageBuilder messageBuilder() { } 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); - } - int length = 2 + 2 + 1 + 2; + this.credit(subscriptionId, credit, CreditUnit.CHUNK); + } - ByteBuf bb = allocate(length + 4); - bb.writeInt(length); - bb.writeShort(encodeRequestCode(COMMAND_CREDIT)); - bb.writeShort(VERSION_1); - bb.writeByte(subscriptionId); - bb.writeShort((short) credit); - channel.writeAndFlush(bb, channel.voidPromise()); + public void credit(byte subscriptionId, int credit, CreditUnit unit) { + if (unit == CreditUnit.BYTE) { + if (credit < 0) { + throw new IllegalArgumentException( + "Credit value must be between 0 and " + Integer.MAX_VALUE); + } + int length = 2 + 2 + 1 + 4; + + ByteBuf bb = allocate(length + 4); + bb.writeInt(length); + bb.writeShort(encodeRequestCode(COMMAND_CREDIT)); + bb.writeShort(VERSION_2); + bb.writeByte(subscriptionId); + bb.writeInt(credit); + channel.writeAndFlush(bb, channel.voidPromise()); + } else { + if (credit < 0 || credit > Short.MAX_VALUE) { + throw new IllegalArgumentException("Credit value must be between 0 and " + Short.MAX_VALUE); + } + int length = 2 + 2 + 1 + 2; + + ByteBuf bb = allocate(length + 4); + bb.writeInt(length); + bb.writeShort(encodeRequestCode(COMMAND_CREDIT)); + bb.writeShort(VERSION_1); + bb.writeByte(subscriptionId); + bb.writeShort((short) credit); + channel.writeAndFlush(bb, channel.voidPromise()); + } } /** @@ -1259,10 +1289,52 @@ public Response subscribe( OffsetSpecification offsetSpecification, int initialCredits, Map properties) { - if (initialCredits < 0 || initialCredits > Short.MAX_VALUE) { - throw new IllegalArgumentException("Credit value must be between 0 and " + Short.MAX_VALUE); + return this.subscribe( + subscriptionId, stream, offsetSpecification, initialCredits, properties, CreditUnit.CHUNK); + } + + /** + * Subscribe to receive messages from a stream, with credit expressed in the given unit. + * + *

Note the offset is an unsigned long. Longs are signed in Java, but unsigned longs can be + * used as long as some care is taken for some operations. See the unsigned* static + * methods in {@link Long}. + * + * @param subscriptionId identifier to correlate inbound messages to this subscription + * @param stream the stream to consume from + * @param offsetSpecification the specification of the offset to consume from + * @param initialCredits the initial number of credits, in {@code unit} + * @param properties some optional properties to describe the subscription + * @param unit the unit {@code initialCredits} is expressed in + * @return the subscription confirmation + * @throws UnsupportedOperationException if {@code unit} is {@link CreditUnit#BYTE} and the broker + * does not support {@code Subscribe} version 2 + */ + public Response subscribe( + byte subscriptionId, + String stream, + OffsetSpecification offsetSpecification, + int initialCredits, + Map properties, + CreditUnit unit) { + if (unit == CreditUnit.BYTE && !this.byteCreditSupported()) { + throw new UnsupportedOperationException( + "Byte-based consumer credit requires a broker supporting Subscribe version 2 and " + + "Credit version 2"); + } + if (unit == CreditUnit.BYTE) { + if (initialCredits < 0) { + throw new IllegalArgumentException( + "Credit value must be between 0 and " + Integer.MAX_VALUE); + } + } else { + if (initialCredits < 0 || initialCredits > Short.MAX_VALUE) { + throw new IllegalArgumentException("Credit value must be between 0 and " + Short.MAX_VALUE); + } } - int length = 2 + 2 + 4 + 1 + 2 + stringByteSize(stream) + 2 + 2; // misses the offset + int creditFieldSize = unit == CreditUnit.BYTE ? 4 : 2; + int length = + 2 + 2 + 4 + 1 + 2 + stringByteSize(stream) + 2 + creditFieldSize; // misses the offset if (offsetSpecification.isOffset() || offsetSpecification.isTimestamp()) { length += 8; } @@ -1276,7 +1348,7 @@ public Response subscribe( Response response = sendRpc( COMMAND_SUBSCRIBE, - VERSION_1, + unit == CreditUnit.BYTE ? VERSION_2 : VERSION_1, length, bb -> { bb.writeByte(subscriptionId); @@ -1285,7 +1357,11 @@ public Response subscribe( if (offsetSpecification.isOffset() || offsetSpecification.isTimestamp()) { bb.writeLong(offsetSpecification.getOffset()); } - bb.writeShort(initialCredits); + if (unit == CreditUnit.BYTE) { + bb.writeInt(initialCredits); + } else { + bb.writeShort(initialCredits); + } if (properties != null && !properties.isEmpty()) { writeMap(bb, properties); } @@ -1599,6 +1675,10 @@ public boolean filteringSupported() { return this.filteringSupported; } + public boolean byteCreditSupported() { + return this.byteCreditSupported; + } + public List route(String routingKey, String superStream) { if (routingKey == null || superStream == null) { throw new IllegalArgumentException("routing key and stream must not be null"); @@ -1794,10 +1874,18 @@ public interface ChunkListener { * @param offset the first offset in the chunk * @param messageCount the total number of messages in the chunk * @param dataSize the size in bytes of the data in the chunk + * @param chunkByteCount the number of bytes of the chunk carried in the {@code Deliver} frame, + * that is the cost the broker charged for it; this is what a byte-based subscription must + * grant back as credit * @return a "chunk context" instance that'll be passed in to the {@link MessageListener} */ Object handle( - Client client, byte subscriptionId, long offset, long messageCount, long dataSize); + Client client, + byte subscriptionId, + long offset, + long messageCount, + long dataSize, + long chunkByteCount); } public interface MessageListener { @@ -2402,7 +2490,7 @@ public static class ClientParameters { private PublishConfirmListener publishConfirmListener = NO_OP_PUBLISH_CONFIRM_LISTENER; private PublishErrorListener publishErrorListener = NO_OP_PUBLISH_ERROR_LISTENER; private ChunkListener chunkListener = - (client, correlationId, offset, messageCount, dataSize) -> null; + (client, correlationId, offset, messageCount, dataSize, chunkByteCount) -> null; private MessageListener messageListener = (correlationId, offset, chunkTimestamp, committedOffset, chunkContext, message) -> {}; private MessageIgnoredListener messageIgnoredListener = diff --git a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java index 70084b8ce3..d1b0611e51 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java +++ b/src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java @@ -34,6 +34,7 @@ import com.rabbitmq.stream.Constants; import com.rabbitmq.stream.Consumer; import com.rabbitmq.stream.ConsumerFlowStrategy; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; import com.rabbitmq.stream.MessageHandler; import com.rabbitmq.stream.MessageHandler.Context; import com.rabbitmq.stream.OffsetSpecification; @@ -453,6 +454,7 @@ private static class SubscriptionTracker { private final AtomicReference state = new AtomicReference<>(SubscriptionState.OPENING); private final ConsumerFlowStrategy flowStrategy; + private final CreditAccountant creditAccountant; private final Lock subscriptionTrackerLock = new ReentrantLock(); private SubscriptionTracker( @@ -475,6 +477,10 @@ private SubscriptionTracker( this.trackingClosingCallback = trackingClosingCallback; this.messageHandler = messageHandler; this.flowStrategy = flowStrategy; + this.creditAccountant = + flowStrategy.unit() == CreditUnit.BYTE + ? new ByteCreditAccountant() + : ChunkCreditAccountant.INSTANCE; if (this.offsetTrackingReference == null) { this.subscriptionProperties = subscriptionProperties; } else { @@ -659,15 +665,22 @@ private ClientSubscriptionsManager( AtomicBoolean clientInitializedInManager = new AtomicBoolean(false); ChunkListener chunkListener = - (client, subscriptionId, offset, messageCount, dataSize) -> { + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); ConsumerFlowStrategy.MessageProcessedCallback processCallback; if (subscriptionTracker != null && subscriptionTracker.consumer.isOpen()) { + subscriptionTracker.creditAccountant.chunkArrived( + client, subscriptionId, chunkByteCount); processCallback = subscriptionTracker.flowStrategy.start( new DefaultConsumerFlowStrategyContext( - subscriptionId, client, messageCount, offset)); + subscriptionId, + client, + messageCount, + offset, + chunkByteCount, + subscriptionTracker.creditAccountant)); } else { LOGGER.debug( "Could not find stream subscription {} or subscription closing, not providing credits", @@ -682,11 +695,22 @@ private ClientSubscriptionsManager( SubscriptionTracker subscriptionTracker = subscriptionTrackers.get(subscriptionId & 0xFF); String stream = subscriptionTracker == null ? "?" : subscriptionTracker.stream; - LOGGER.debug( - "Received credit notification for subscription {} (stream '{}'): {}", - subscriptionId & 0xFF, - stream, - Utils.formatConstant(responseCode)); + if (responseCode == Constants.RESPONSE_CODE_PRECONDITION_FAILED) { + // a unit mismatch between the subscription and the credit frame, necessarily a + // client bug; the credit was dropped, so the subscription is short of credit for + // good + LOGGER.warn( + "Received credit notification for subscription {} (stream '{}'): {}", + subscriptionId & 0xFF, + stream, + Utils.formatConstant(responseCode)); + } else { + LOGGER.debug( + "Received credit notification for subscription {} (stream '{}'): {}", + subscriptionId & 0xFF, + stream, + Utils.formatConstant(responseCode)); + } }; MessageListener messageListener = @@ -1091,6 +1115,14 @@ void add( } checkNotClosed(); + if (tracker.flowStrategy.unit() == CreditUnit.BYTE && !this.client.byteCreditSupported()) { + // must not be an IllegalStateException: addToManager treats that as "this manager + // cannot take the subscription" and loops looking for another one, which would spin + // forever on a node that will never support Subscribe/Credit version 2 + throw new StreamException( + "Byte-based consumer credit requires a broker supporting Subscribe version 2 " + + "and Credit version 2"); + } byte subscriptionId = (byte) pickSlot(this.subscriptionTrackers, this.consumerIndexSequence); @@ -1161,6 +1193,10 @@ void add( subscriptionContext.offsetSpecification()); checkNotClosed(); + int initialCredits = tracker.flowStrategy.initialCredits(); + // resetting on every subscription, including recovery, keeps the mirror correct + // after a reconnection or a stream move + tracker.creditAccountant.reset(initialCredits); Client.Response subscribeResponse = Utils.callAndMaybeRetry( () -> @@ -1168,8 +1204,9 @@ void add( subscriptionId, tracker.stream, subscriptionContext.offsetSpecification(), - tracker.flowStrategy.initialCredits(), - tracker.subscriptionProperties), + initialCredits, + tracker.subscriptionProperties, + tracker.flowStrategy.unit()), RETRY_ON_TIMEOUT, "Subscribe request for consumer %d on stream '%s'", tracker.consumer.id(), @@ -1438,19 +1475,38 @@ private static class DefaultConsumerFlowStrategyContext implements ConsumerFlowS private final Client client; private final long messageCount; private final long chunkId; + private final long chunkByteCount; + private final CreditAccountant creditAccountant; + // guards against releasing this chunk's credit more than once, no matter how many times + // the strategy calls credits(...) for it + private final AtomicBoolean released = new AtomicBoolean(false); private DefaultConsumerFlowStrategyContext( - byte subscriptionId, Client client, long messageCount, long chunkId) { + byte subscriptionId, + Client client, + long messageCount, + long chunkId, + long chunkByteCount, + CreditAccountant creditAccountant) { this.subscriptionId = subscriptionId; this.client = client; this.messageCount = messageCount; this.chunkId = chunkId; + this.chunkByteCount = chunkByteCount; + this.creditAccountant = creditAccountant; } @Override public void credits(int credits) { + if (!this.released.compareAndSet(false, true)) { + LOGGER.debug( + "Credit already released for subscription {}, chunk {}, ignoring extra call", + subscriptionId, + chunkId); + return; + } try { - client.credit(subscriptionId, credits); + this.creditAccountant.release(client, subscriptionId, credits, chunkByteCount); } catch (Exception e) { LOGGER.info( "Error while providing {} credit(s) to subscription {}: {}", @@ -1469,6 +1525,169 @@ public long messageCount() { public long chunkId() { return this.chunkId; } + + @Override + public long chunkByteCount() { + return this.chunkByteCount; + } + } + + /** + * Grants credit for a subscription. + * + *

{@link #chunkArrived(Client, byte, long)} is called once per chunk, before the chunk's + * {@link ConsumerFlowStrategy} context is created. {@link #release(Client, byte, int, long)} is + * called at most once per chunk, from that chunk's flow strategy context. + */ + interface CreditAccountant { + + /** + * Called before subscription. + * + * @param initialCredits + */ + void reset(int initialCredits); + + /** + * Called on chunk arrival. + * + * @param client + * @param subscriptionId + * @param chunkByteCount + */ + void chunkArrived(Client client, byte subscriptionId, long chunkByteCount); + + /** + * Called when the flow strategy provide credits via the default context. + * + * @param client + * @param subscriptionId + * @param chunks + * @param chunkByteCount + */ + void release(Client client, byte subscriptionId, int chunks, long chunkByteCount); + } + + /** Chunk-based credit: a pass-through to {@link Client#credit(byte, int)}. */ + static final class ChunkCreditAccountant implements CreditAccountant { + + static final CreditAccountant INSTANCE = new ChunkCreditAccountant(); + + @Override + public void reset(int initialCredits) {} + + @Override + public void chunkArrived(Client client, byte subscriptionId, long chunkByteCount) {} + + @Override + public void release(Client client, byte subscriptionId, int chunks, long chunkByteCount) { + client.credit(subscriptionId, chunks); + } + } + + /** + * Byte-based credit: keeps an exact mirror of the broker-side credit for a subscription, and + * batches grants instead of sending one {@code Credit} frame per released chunk. + * + *

{@code credit = window + granted - received}, so it only ever decreases in {@link + * #chunkArrived(Client, byte, long)} and only ever increases when a grant is flushed. A grant is + * flushed once {@code credit} drops to {@code flushThreshold}, three quarters of the window, + * deliberately above the broker's {@code send_limit} (half the window, see {@code + * rabbit_stream_reader:send_chunks/6}): the client provably owes nothing by the time the broker + * can become blocked, so no timer or further delivery is needed to get the grant out. + */ + static final class ByteCreditAccountant implements CreditAccountant { + + // chunks arrive on the connection dispatching thread, processed() can be called from any + // application thread + private final Lock lock = new ReentrantLock(); + private long window; + private long flushThreshold; + private long pending; + private long credit; + + @Override + public void reset(int initialCredits) { + lock( + this.lock, + () -> { + this.window = initialCredits; + this.flushThreshold = this.window - this.window / 4; + this.credit = this.window; + this.pending = 0; + }); + } + + @Override + public void chunkArrived(Client client, byte subscriptionId, long chunkByteCount) { + long toGrant; + this.lock.lock(); + try { + this.credit -= chunkByteCount; + toGrant = maybeFlushLocked(); + } finally { + this.lock.unlock(); + } + grant(client, subscriptionId, toGrant); + } + + @Override + public void release(Client client, byte subscriptionId, int chunks, long chunkByteCount) { + if (chunks != 1) { + LOGGER.debug( + "Byte-based credit release called with {} chunk(s) instead of 1, " + + "ignoring the chunk count", + chunks); + } + long toGrant; + this.lock.lock(); + try { + this.pending += chunkByteCount; + toGrant = maybeFlushLocked(); + } finally { + this.lock.unlock(); + } + grant(client, subscriptionId, toGrant); + } + + // must be called with the lock held + private long maybeFlushLocked() { + if (this.pending > 0 && this.credit <= this.flushThreshold) { + long toGrant = this.pending; + this.credit += this.pending; + this.pending = 0; + return toGrant; + } + return 0; + } + + // the Credit frame is written outside the lock, concurrent grants are additive so their + // order does not matter + private static void grant(Client client, byte subscriptionId, long credit) { + if (credit > 0) { + client.credit(subscriptionId, (int) credit, CreditUnit.BYTE); + } + } + + // for tests + long credit() { + this.lock.lock(); + try { + return this.credit; + } finally { + this.lock.unlock(); + } + } + + // for tests + long pending() { + this.lock.lock(); + try { + return this.pending; + } finally { + this.lock.unlock(); + } + } } static int pickSlot(List list, AtomicInteger sequence) { diff --git a/src/main/java/com/rabbitmq/stream/impl/ServerFrameHandler.java b/src/main/java/com/rabbitmq/stream/impl/ServerFrameHandler.java index 3b1cea6b9a..ec10dcf342 100644 --- a/src/main/java/com/rabbitmq/stream/impl/ServerFrameHandler.java +++ b/src/main/java/com/rabbitmq/stream/impl/ServerFrameHandler.java @@ -103,6 +103,13 @@ class ServerFrameHandler { static final int INITIAL_DECOMPRESSION_BUFFER_SIZE = 64 * 1024; // transfer buffer size private static final int DECOMPRESSION_TRANSFER_BUFFER_SIZE = 1024; + // the chunk cost is the contract with the broker, not derived from the chunk header: it is + // the frame size minus the command id, version and subscription ID (5 bytes) + static final int CHUNK_COST_DELIVER_V1_OFFSET = 5; + // the chunk cost is the contract with the broker, not derived from the chunk header: it is + // the frame size minus the command id, version, subscription ID, and committed chunk ID (13 + // bytes) + static final int CHUNK_COST_DELIVER_V2_OFFSET = 13; private static final FrameHandler[][] HANDLERS; @@ -324,7 +331,8 @@ private abstract static class BaseFrameHandler implements FrameHandler { public void handle(Client client, int frameSize, ChannelHandlerContext ctx, ByteBuf message) { int readerIndexBefore = message.readerIndex(); try { - int read = doHandle(client, ctx, message) + 4; // already read the command id and version + // already read the command id and version + int read = doHandle(client, ctx, message, frameSize) + 4; if (read != frameSize) { LOGGER.warn("Read {} bytes in frame, expecting {}", read, frameSize); } @@ -341,7 +349,15 @@ public void handle(Client client, int frameSize, ChannelHandlerContext ctx, Byte } } - abstract int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message); + int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) { + throw new UnsupportedOperationException(); + } + + // overridden by handlers that need the frame size (e.g. to compute a byte cost); defaults to + // the frame-size-agnostic variant + int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message, int frameSize) { + return doHandle(client, ctx, message); + } protected void logMissingOutstandingRequest(int correlationId) { LOGGER.warn( @@ -444,7 +460,8 @@ static int handleDeliverVersion1( MessageIgnoredListener messageIgnoredListener, Codec codec, ChunkChecksum chunkChecksum, - MetricsCollector metricsCollector) { + MetricsCollector metricsCollector, + long chunkByteCount) { return handleDeliver( message, client, @@ -457,8 +474,8 @@ static int handleDeliverVersion1( metricsCollector, message.readByte(), // subscription ID 0, // last committed offset - 1 // byte read count - ); + 1, // byte read count + chunkByteCount); } static int handleDeliver( @@ -473,7 +490,8 @@ static int handleDeliver( MetricsCollector metricsCollector, byte subscriptionId, long committedOffset, - int read) { + int read, + long chunkByteCount) { /* %% << %% Magic=5:4/unsigned, @@ -520,7 +538,8 @@ static int handleDeliver( read += 4; Object chunkContext = - chunkListener.handle(client, subscriptionId, offset, numRecords, dataLength); + chunkListener.handle( + client, subscriptionId, offset, numRecords, dataLength, chunkByteCount); long offsetLimit = client.extractInitialSubscriptionOffset(subscriptionId); @@ -722,7 +741,7 @@ static int handleDeliver( } @Override - int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) { + int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message, int frameSize) { return handleDeliverVersion1( message, client, @@ -732,7 +751,8 @@ int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) { client.messageIgnoredListener, client.codec, client.chunkChecksum, - client.metricsCollector); + client.metricsCollector, + frameSize - CHUNK_COST_DELIVER_V1_OFFSET); } } @@ -744,7 +764,7 @@ public boolean isInitiatedByServer() { } @Override - int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) { + int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message, int frameSize) { return DeliverVersion1FrameHandler.handleDeliver( message, client, @@ -757,8 +777,8 @@ int doHandle(Client client, ChannelHandlerContext ctx, ByteBuf message) { client.metricsCollector, message.readByte(), // subscription ID message.readLong(), // committed chunk ID, unsigned long - 9 // byte read count, 1 + 9 - ); + 9, // byte read count, 1 + 9 + frameSize - CHUNK_COST_DELIVER_V2_OFFSET); } } diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java index 28f8ecc16e..f3b0352c75 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java @@ -98,6 +98,12 @@ final class StreamConsumer extends ResourceBase implements Consumer { "Filtering is not supported by the broker " + "(requires RabbitMQ 3.13+ and stream_filtering feature flag activated"); } + if (flowStrategy.unit() == ConsumerFlowStrategy.CreditUnit.BYTE + && !environment.byteCreditSupported()) { + throw new IllegalArgumentException( + "Byte-based consumer credit is not supported by the broker " + + "(requires a broker supporting Subscribe version 2 and Credit version 2)"); + } this.id = ID_SEQUENCE.getAndIncrement(); Runnable trackingClosingCallback; try { diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java index 3b318c142a..9f5a4654e7 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamConsumerBuilder.java @@ -17,6 +17,7 @@ import static com.rabbitmq.stream.impl.Utils.SUBSCRIPTION_PROPERTY_FILTER_PREFIX; import static com.rabbitmq.stream.impl.Utils.SUBSCRIPTION_PROPERTY_MATCH_UNFILTERED; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.Consumer; import com.rabbitmq.stream.ConsumerBuilder; import com.rabbitmq.stream.ConsumerFlowStrategy; @@ -476,6 +477,12 @@ public FlowConfiguration initialCredits(int initialCredits) { return this; } + @Override + public FlowConfiguration initialCredits(ByteCapacity initialCredits) { + this.strategy = ConsumerFlowStrategy.creditOnChunkArrival(initialCredits); + return this; + } + @Override public FlowConfiguration strategy(ConsumerFlowStrategy strategy) { this.strategy = strategy; diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java index f56cee8d6e..0ca9bbbf3d 100644 --- a/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java +++ b/src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java @@ -914,6 +914,10 @@ boolean filteringSupported() { return this.locatorOperation(Client::filteringSupported); } + boolean byteCreditSupported() { + return this.locatorOperation(Client::byteCreditSupported); + } + Clock clock() { return this.clock; } diff --git a/src/test/java/com/rabbitmq/stream/docs/ConsumerUsage.java b/src/test/java/com/rabbitmq/stream/docs/ConsumerUsage.java index e6ac4cc9b0..c4b2abaede 100644 --- a/src/test/java/com/rabbitmq/stream/docs/ConsumerUsage.java +++ b/src/test/java/com/rabbitmq/stream/docs/ConsumerUsage.java @@ -16,6 +16,7 @@ import java.time.Duration; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.Consumer; import com.rabbitmq.stream.ConsumerFlowStrategy; import com.rabbitmq.stream.Environment; @@ -173,6 +174,22 @@ void flowControl() { // end::flow-control[] } + void flowControlByteBased() { + Environment environment = Environment.builder().build(); + // tag::flow-control-byte-based[] + Consumer consumer = environment.consumerBuilder() + .stream("my-stream") + .flow() + .strategy(ConsumerFlowStrategy.creditWhenHalfMessagesProcessed(ByteCapacity.kB(512))) // <1> + .builder() + .messageHandler((context, message) -> { + // message handling code (possibly asynchronous)... + context.processed(); // <2> + }) + .build(); + // end::flow-control-byte-based[] + } + void enablingSingleActiveConsumer() { Environment environment = Environment.builder().build(); // tag::enabling-single-active-consumer[] diff --git a/src/test/java/com/rabbitmq/stream/impl/ByteCreditAccountantTest.java b/src/test/java/com/rabbitmq/stream/impl/ByteCreditAccountantTest.java new file mode 100644 index 0000000000..30d70ec36d --- /dev/null +++ b/src/test/java/com/rabbitmq/stream/impl/ByteCreditAccountantTest.java @@ -0,0 +1,135 @@ +// Copyright (c) 2026 Broadcom. All Rights Reserved. +// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. +// +// This software, the RabbitMQ Stream Java client library, is dual-licensed under the +// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL"). +// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL, +// please see LICENSE-APACHE2. +// +// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND, +// either express or implied. See the LICENSE file for specific language governing +// rights and limitations of this software. +// +// If you have any questions regarding licensing, please contact us at +// info@rabbitmq.com. +package com.rabbitmq.stream.impl; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeast; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; +import com.rabbitmq.stream.impl.ConsumersCoordinator.ByteCreditAccountant; +import java.util.Random; +import org.junit.jupiter.api.RepeatedTest; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +public class ByteCreditAccountantTest { + + static final byte SUBSCRIPTION_ID = 0; + + @Test + void releaseGrantsOnlyItsOwnChunkBytes() { + Client client = mock(Client.class); + ByteCreditAccountant accountant = new ByteCreditAccountant(); + accountant.reset(2_000_000); + accountant.chunkArrived(client, SUBSCRIPTION_ID, 1024); + accountant.chunkArrived(client, SUBSCRIPTION_ID, 1_000_000); + + accountant.release(client, SUBSCRIPTION_ID, 1, 1024); + + verify(client, times(1)).credit(SUBSCRIPTION_ID, 1024, CreditUnit.BYTE); + verify(client, never()).credit(eq(SUBSCRIPTION_ID), eq(1_000_000), eq(CreditUnit.BYTE)); + } + + @Test + void grantsAreDeferredUntilCreditDropsToFlushThresholdThenBatched() { + Client client = mock(Client.class); + ByteCreditAccountant accountant = new ByteCreditAccountant(); + accountant.reset(1000); // flush threshold is 750 + + accountant.chunkArrived(client, SUBSCRIPTION_ID, 100); // credit = 900 + accountant.release(client, SUBSCRIPTION_ID, 1, 100); // pending = 100 + accountant.chunkArrived(client, SUBSCRIPTION_ID, 100); // credit = 800 + accountant.release(client, SUBSCRIPTION_ID, 1, 100); // pending = 200 + verify(client, never()).credit(eq(SUBSCRIPTION_ID), anyInt(), eq(CreditUnit.BYTE)); + + accountant.chunkArrived(client, SUBSCRIPTION_ID, 100); // credit = 700 <= threshold, flushes + verify(client, times(1)).credit(SUBSCRIPTION_ID, 200, CreditUnit.BYTE); + } + + @Test + void chunkLargerThanWindowDrivesCreditNegativeAndReleaseBringsItBackToTheWindow() { + Client client = mock(Client.class); + ByteCreditAccountant accountant = new ByteCreditAccountant(); + accountant.reset(1000); + + accountant.chunkArrived(client, SUBSCRIPTION_ID, 1500); + assertThat(accountant.credit()).isEqualTo(-500); + verify(client, never()).credit(eq(SUBSCRIPTION_ID), anyInt(), eq(CreditUnit.BYTE)); + + accountant.release(client, SUBSCRIPTION_ID, 1, 1500); + verify(client, times(1)).credit(SUBSCRIPTION_ID, 1500, CreditUnit.BYTE); + assertThat(accountant.credit()).isEqualTo(1000); + } + + @Test + void resetDiscardsPriorState() { + Client client = mock(Client.class); + ByteCreditAccountant accountant = new ByteCreditAccountant(); + accountant.reset(1000); + accountant.chunkArrived(client, SUBSCRIPTION_ID, 100); + accountant.release(client, SUBSCRIPTION_ID, 1, 200); + assertThat(accountant.pending()).isNotZero(); + + accountant.reset(500); + + assertThat(accountant.credit()).isEqualTo(500); + assertThat(accountant.pending()).isZero(); + } + + @RepeatedTest(20) + void invariantsHoldOverRandomSequencesOfArrivalsAndReleases() { + Client client = mock(Client.class); + ByteCreditAccountant accountant = new ByteCreditAccountant(); + int window = 1000; + accountant.reset(window); + + Random random = new Random(); + long totalReceived = 0; + long totalReleased = 0; + + for (int i = 0; i < 200; i++) { + long chunkCost = 1 + random.nextInt(300); + accountant.chunkArrived(client, SUBSCRIPTION_ID, chunkCost); + totalReceived += chunkCost; + + // releases do not always happen right away: some chunks stay unreleased for a while, + // exercising sequences where receiving runs ahead of releasing + if (random.nextBoolean()) { + accountant.release(client, SUBSCRIPTION_ID, 1, chunkCost); + totalReleased += chunkCost; + } + + assertThat(accountant.pending() > 0 && accountant.credit() <= window / 2) + .describedAs("no-deadlock post-condition violated at iteration %d", i) + .isFalse(); + } + // force a final flush of whatever is still pending, to reconcile totals + accountant.chunkArrived(client, SUBSCRIPTION_ID, window * 10L); + + ArgumentCaptor creditCaptor = ArgumentCaptor.forClass(Integer.class); + verify(client, atLeast(0)) + .credit(eq(SUBSCRIPTION_ID), creditCaptor.capture(), eq(CreditUnit.BYTE)); + long totalGranted = creditCaptor.getAllValues().stream().mapToLong(Integer::longValue).sum(); + + assertThat(totalGranted).isEqualTo(totalReleased); + assertThat(totalGranted).isLessThanOrEqualTo(totalReceived); + } +} diff --git a/src/test/java/com/rabbitmq/stream/impl/ClientTest.java b/src/test/java/com/rabbitmq/stream/impl/ClientTest.java index bf5cf43343..ad37da8804 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ClientTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ClientTest.java @@ -26,12 +26,15 @@ import static java.util.concurrent.TimeUnit.SECONDS; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; import com.rabbitmq.client.Channel; import com.rabbitmq.client.Connection; import com.rabbitmq.client.ConnectionFactory; import com.rabbitmq.stream.Codec; import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; import com.rabbitmq.stream.Message; import com.rabbitmq.stream.MessageBuilder; import com.rabbitmq.stream.OffsetSpecification; @@ -49,6 +52,7 @@ import com.rabbitmq.stream.impl.ServerFrameHandler.FrameHandlerInfo; import com.rabbitmq.stream.impl.TestUtils.BrokerVersion; import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast; +import com.rabbitmq.stream.impl.TestUtils.DisabledIfByteCreditNotSupported; import com.rabbitmq.stream.impl.TestUtils.DisabledIfFilteringNotSupported; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufAllocator; @@ -502,7 +506,7 @@ void consume() throws Exception { AtomicInteger receivedCorrelationId = new AtomicInteger(); Client.ChunkListener chunkListener = - (client, corr, offset, messageCountInChunk, dataSize) -> { + (client, corr, offset, messageCountInChunk, dataSize, chunkByteCount) -> { receivedCorrelationId.set(corr); client.credit(correlationId, 1); return null; @@ -539,7 +543,7 @@ void publishAndConsume(boolean directBuffer) throws Exception { CountDownLatch consumedLatch = new CountDownLatch(publishCount); Client.ChunkListener chunkListener = - (client, correlationId, offset, messageCount, dataSize) -> { + (client, correlationId, offset, messageCount, dataSize, chunkByteCount) -> { if (consumedLatch.getCount() != 0) { client.credit(correlationId, 1); } @@ -1344,4 +1348,261 @@ void testSubscriptionTrackerState() { assertThat(offset).isEqualTo(500); assertThat(tracker.hasOffsets()).isFalse(); } + + private void publishMessagesOneByOne(int count, byte[] body) throws Exception { + AtomicReference confirmLatch = new AtomicReference<>(new CountDownLatch(1)); + Client publisher = + cf.get( + new ClientParameters() + .publishConfirmListener( + (publisherId, publishingId) -> confirmLatch.get().countDown())); + publisher.declarePublisher(b(1), null, stream); + for (int i = 0; i < count; i++) { + confirmLatch.set(new CountDownLatch(1)); + publisher.publish( + b(1), Collections.singletonList(publisher.messageBuilder().addData(body).build())); + assertThat(confirmLatch.get().await(10, SECONDS)).isTrue(); + } + } + + @Test + @BrokerVersionAtLeast(BrokerVersion.RABBITMQ_3_11_0) + void byteCreditSupportedReflectsExchangeCommandVersions() { + Client client = cf.get(); + List infos = client.exchangeCommandVersions(); + boolean subscribeVersion2Supported = + infos.stream() + .anyMatch( + info -> + info.getKey() == Constants.COMMAND_SUBSCRIBE + && info.getMaxVersion() >= Constants.VERSION_2); + boolean creditVersion2Supported = + infos.stream() + .anyMatch( + info -> + info.getKey() == Constants.COMMAND_CREDIT + && info.getMaxVersion() >= Constants.VERSION_2); + assertThat(client.byteCreditSupported()) + .isEqualTo(subscribeVersion2Supported && creditVersion2Supported); + } + + @Test + void subscribeWithByteCreditFailsOnUnsupportingBroker() { + Client client = spy(cf.get()); + when(client.byteCreditSupported()).thenReturn(false); + assertThatThrownBy( + () -> + client.subscribe( + b(1), + stream, + OffsetSpecification.first(), + 1, + Collections.emptyMap(), + CreditUnit.BYTE)) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + @DisabledIfByteCreditNotSupported + void subscribeWithByteCreditDeliversUpToTheWindow() throws Exception { + publishMessagesOneByOne(10, new byte[20]); + + AtomicInteger deliveredChunkCount = new AtomicInteger(0); + CountDownLatch firstChunkLatch = new CountDownLatch(1); + CountDownLatch secondChunkLatch = new CountDownLatch(2); + Client consumer = + cf.get( + new ClientParameters() + .chunkListener( + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + deliveredChunkCount.incrementAndGet(); + firstChunkLatch.countDown(); + secondChunkLatch.countDown(); + return null; + })); + Response response = + consumer.subscribe( + b(1), stream, OffsetSpecification.first(), 1, Collections.emptyMap(), CreditUnit.BYTE); + assertThat(response).is(ok()); + + assertThat(latchAssert(firstChunkLatch)).completes(); + assertThat(latchAssert(secondChunkLatch)).doesNotComplete(2); + assertThat(deliveredChunkCount.get()).isEqualTo(1); + } + + @Test + @DisabledIfByteCreditNotSupported + void chunkCostIsFrameSizeBased() throws Exception { + publishMessagesOneByOne(10, new byte[20]); + + AtomicInteger deliveredChunkCount = new AtomicInteger(0); + AtomicLong firstChunkByteCount = new AtomicLong(-1); + CountDownLatch firstChunkLatch = new CountDownLatch(1); + CountDownLatch secondChunkLatch = new CountDownLatch(1); + Client consumer = + cf.get( + new ClientParameters() + .chunkListener( + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + if (deliveredChunkCount.incrementAndGet() == 1) { + firstChunkByteCount.set(chunkByteCount); + firstChunkLatch.countDown(); + } else { + secondChunkLatch.countDown(); + } + return null; + })); + Response response = + consumer.subscribe( + b(1), stream, OffsetSpecification.first(), 1, Collections.emptyMap(), CreditUnit.BYTE); + assertThat(response).is(ok()); + assertThat(latchAssert(firstChunkLatch)).completes(); + + consumer.credit(b(1), (int) (firstChunkByteCount.get() - 1), CreditUnit.BYTE); + assertThat(latchAssert(secondChunkLatch)).doesNotComplete(2); + assertThat(deliveredChunkCount.get()).isEqualTo(1); + + consumer.credit(b(1), 1, CreditUnit.BYTE); + assertThat(latchAssert(secondChunkLatch)).completes(); + assertThat(deliveredChunkCount.get()).isEqualTo(2); + } + + @Test + @DisabledIfByteCreditNotSupported + void chunkLargerThanWindowIsStillDelivered() throws Exception { + publishMessagesOneByOne(1, new byte[10 * 1024]); + + CountDownLatch messageLatch = new CountDownLatch(1); + Client consumer = + cf.get( + new ClientParameters() + .chunkListener(TestUtils.creditBytes()) + .messageListener( + (subscriptionId, + offset, + chunkTimestamp, + committedChunkId, + chunkContext, + message) -> messageLatch.countDown())); + Response response = + consumer.subscribe( + b(1), + stream, + OffsetSpecification.first(), + 100, + Collections.emptyMap(), + CreditUnit.BYTE); + assertThat(response).is(ok()); + assertThat(latchAssert(messageLatch)).completes(); + } + + @Test + @DisabledIfByteCreditNotSupported + void byteCreditSpansManyChunks() throws Exception { + byte[] body = new byte[20]; + publishMessagesOneByOne(100, body); + + AtomicLong firstChunkByteCount = new AtomicLong(-1); + CountDownLatch probeLatch = new CountDownLatch(1); + Client probe = + cf.get( + new ClientParameters() + .chunkListener( + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + firstChunkByteCount.set(chunkByteCount); + probeLatch.countDown(); + return null; + })); + probe.subscribe( + b(1), stream, OffsetSpecification.first(), 1, Collections.emptyMap(), CreditUnit.BYTE); + assertThat(latchAssert(probeLatch)).completes(); + probe.close(); + + int window = (int) (3 * firstChunkByteCount.get() + 1); + + AtomicInteger deliveredChunkCount = new AtomicInteger(0); + CountDownLatch fourChunksLatch = new CountDownLatch(4); + CountDownLatch fifthChunkLatch = new CountDownLatch(5); + Client consumer = + cf.get( + new ClientParameters() + .chunkListener( + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + deliveredChunkCount.incrementAndGet(); + fourChunksLatch.countDown(); + fifthChunkLatch.countDown(); + return null; + })); + consumer.subscribe( + b(1), stream, OffsetSpecification.first(), window, Collections.emptyMap(), CreditUnit.BYTE); + assertThat(latchAssert(fourChunksLatch)).completes(); + assertThat(latchAssert(fifthChunkLatch)).doesNotComplete(2); + assertThat(deliveredChunkCount.get()).isEqualTo(4); + } + + @Test + @DisabledIfByteCreditNotSupported + void creditUnitMismatchIsRejectedForByteBasedSubscription() throws Exception { + publishMessagesOneByOne(1, new byte[20]); + + AtomicInteger deliveredChunkCount = new AtomicInteger(0); + CountDownLatch notificationLatch = new CountDownLatch(1); + AtomicReference responseCode = new AtomicReference<>(); + Client consumer = + cf.get( + new ClientParameters() + .chunkListener( + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + deliveredChunkCount.incrementAndGet(); + return null; + }) + .creditNotification( + (subscriptionId, code) -> { + responseCode.set(code); + notificationLatch.countDown(); + })); + Response response = + consumer.subscribe( + b(1), stream, OffsetSpecification.first(), 0, Collections.emptyMap(), CreditUnit.BYTE); + assertThat(response).is(ok()); + + // subscription is byte-based, sending a chunk-based (version 1) credit must be rejected + consumer.credit(b(1), 1); + + assertThat(latchAssert(notificationLatch)).completes(); + assertThat(responseCode.get()).isEqualTo(Constants.RESPONSE_CODE_PRECONDITION_FAILED); + assertThat(deliveredChunkCount.get()).isZero(); + } + + @Test + @DisabledIfByteCreditNotSupported + void creditUnitMismatchIsRejectedForChunkBasedSubscription() throws Exception { + publishMessagesOneByOne(1, new byte[20]); + + AtomicInteger deliveredChunkCount = new AtomicInteger(0); + CountDownLatch notificationLatch = new CountDownLatch(1); + AtomicReference responseCode = new AtomicReference<>(); + Client consumer = + cf.get( + new ClientParameters() + .chunkListener( + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + deliveredChunkCount.incrementAndGet(); + return null; + }) + .creditNotification( + (subscriptionId, code) -> { + responseCode.set(code); + notificationLatch.countDown(); + })); + Response response = consumer.subscribe(b(1), stream, OffsetSpecification.first(), 0); + assertThat(response).is(ok()); + + // subscription is chunk-based, sending a byte-based (version 2) credit must be rejected + consumer.credit(b(1), 1, CreditUnit.BYTE); + + assertThat(latchAssert(notificationLatch)).completes(); + assertThat(responseCode.get()).isEqualTo(Constants.RESPONSE_CODE_PRECONDITION_FAILED); + assertThat(deliveredChunkCount.get()).isZero(); + } } diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumerFlowStrategyTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumerFlowStrategyTest.java new file mode 100644 index 0000000000..af8ee6e271 --- /dev/null +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumerFlowStrategyTest.java @@ -0,0 +1,87 @@ +// Copyright (c) 2026 Broadcom. All Rights Reserved. +// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. +// +// This software, the RabbitMQ Stream Java client library, is dual-licensed under the +// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL"). +// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL, +// please see LICENSE-APACHE2. +// +// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND, +// either express or implied. See the LICENSE file for specific language governing +// rights and limitations of this software. +// +// If you have any questions regarding licensing, please contact us at +// info@rabbitmq.com. +package com.rabbitmq.stream.impl; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import com.rabbitmq.stream.ByteCapacity; +import com.rabbitmq.stream.ConsumerFlowStrategy; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; +import java.util.function.Function; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +public class ConsumerFlowStrategyTest { + + @ParameterizedTest + @MethodSource("byteCapacityFactories") + void byteCapacityFactoriesReturnByteUnitAndWindowAsInitialCredits( + Function factory) { + ConsumerFlowStrategy strategy = factory.apply(ByteCapacity.MB(2)); + assertThat(strategy.unit()).isEqualTo(CreditUnit.BYTE); + assertThat(strategy.initialCredits()).isEqualTo(2_000_000); + } + + @ParameterizedTest + @MethodSource("byteCapacityFactories") + void byteCapacityFactoriesAcceptWindowOfIntegerMaxValueBytes( + Function factory) { + ConsumerFlowStrategy strategy = factory.apply(ByteCapacity.B(Integer.MAX_VALUE)); + assertThat(strategy.initialCredits()).isEqualTo(Integer.MAX_VALUE); + } + + @ParameterizedTest + @MethodSource("byteCapacityFactories") + void byteCapacityFactoriesRejectNonPositiveWindow( + Function factory) { + assertThatThrownBy(() -> factory.apply(ByteCapacity.B(0))) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> factory.apply(ByteCapacity.B(-1))) + .isInstanceOf(IllegalArgumentException.class); + } + + @ParameterizedTest + @MethodSource("byteCapacityFactories") + void byteCapacityFactoriesRejectWindowLargerThanIntegerMaxValue( + Function factory) { + assertThatThrownBy(() -> factory.apply(ByteCapacity.B(Integer.MAX_VALUE + 1L))) + .isInstanceOf(IllegalArgumentException.class); + } + + @ParameterizedTest + @MethodSource("chunkBasedFactories") + void chunkBasedFactoriesUseChunkUnit(ConsumerFlowStrategy strategy) { + assertThat(strategy.unit()).isEqualTo(CreditUnit.CHUNK); + } + + static Stream> byteCapacityFactories() { + return Stream.of( + ConsumerFlowStrategy::creditOnChunkArrival, + ConsumerFlowStrategy::creditWhenHalfMessagesProcessed, + window -> ConsumerFlowStrategy.creditOnProcessedMessageCount(window, 0.5)); + } + + static Stream chunkBasedFactories() { + return Stream.of( + ConsumerFlowStrategy.creditOnChunkArrival(), + ConsumerFlowStrategy.creditOnChunkArrival(5), + ConsumerFlowStrategy.creditWhenHalfMessagesProcessed(), + ConsumerFlowStrategy.creditWhenHalfMessagesProcessed(5), + ConsumerFlowStrategy.creditOnProcessedMessageCount(5, 0.5), + ConsumerFlowStrategy.creditEveryNthChunk(10, 5)); + } +} diff --git a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java index 8258abab54..c1c678fece 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2020-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2020-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -36,7 +36,9 @@ import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -45,6 +47,7 @@ import com.rabbitmq.stream.Address; import com.rabbitmq.stream.BackOffDelayPolicy; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.Constants; import com.rabbitmq.stream.ConsumerFlowStrategy; import com.rabbitmq.stream.MessageHandler; @@ -109,6 +112,7 @@ public class ConsumersCoordinatorTest { volatile Client.MetadataListener metadataListener; volatile Client.MessageListener messageListener; volatile Client.MessageIgnoredListener messageIgnoredListener; + volatile Client.ChunkListener chunkListener; List messageListeners = new CopyOnWriteArrayList<>(); volatile Client.ShutdownListener shutdownListener; List shutdownListeners = @@ -168,6 +172,12 @@ public Client.ClientParameters shutdownListener( ConsumersCoordinatorTest.this.shutdownListeners.add(shutdownListener); return super.shutdownListener(shutdownListener); } + + @Override + public Client.ClientParameters chunkListener(Client.ChunkListener chunkListener) { + ConsumersCoordinatorTest.this.chunkListener = chunkListener; + return super.chunkListener(chunkListener); + } }; mocks = MockitoAnnotations.openMocks(this); StreamEnvironment.Locator l = new StreamEnvironment.Locator(-1, new Address("localhost", 5555)); @@ -227,7 +237,8 @@ void tearDown() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); when(client.serverAdvertisedHost()).thenReturn("foo").thenReturn(replica().get(0).getHost()); when(client.serverAdvertisedPort()).thenReturn(42).thenReturn(replica().get(0).getPort()); @@ -244,7 +255,13 @@ void tearDown() throws Exception { flowStrategy()); verify(clientFactory, times(2)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); } @Test @@ -269,7 +286,8 @@ void shouldGetExactNodeImmediatelyWithAdvertisedHostNameClientFactoryAndExactNod anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); when(client.serverAdvertisedHost()).thenReturn(replica().get(0).getHost()); when(client.serverAdvertisedPort()).thenReturn(replica().get(0).getPort()); @@ -286,7 +304,13 @@ void shouldGetExactNodeImmediatelyWithAdvertisedHostNameClientFactoryAndExactNod flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); } @Test @@ -311,7 +335,8 @@ void shouldAcceptCandidateNode() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); when(client.serverAdvertisedHost()).thenReturn("foo").thenReturn(replicas().get(1).getHost()); when(client.serverAdvertisedPort()).thenReturn(42).thenReturn(replicas().get(1).getPort()); @@ -328,7 +353,13 @@ void shouldAcceptCandidateNode() { flowStrategy()); verify(clientFactory, times(2)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); } @Test @@ -343,7 +374,8 @@ void shouldSubscribeWithEmptyPropertiesWithUnamedConsumer() { anyString(), any(OffsetSpecification.class), anyInt(), - subscriptionPropertiesArgumentCaptor.capture())) + subscriptionPropertiesArgumentCaptor.capture(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); coordinator.subscribe( @@ -358,7 +390,13 @@ void shouldSubscribeWithEmptyPropertiesWithUnamedConsumer() { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(subscriptionPropertiesArgumentCaptor.getValue()).isEmpty(); } @@ -410,7 +448,8 @@ void subscribePropagateExceptionWhenClientSubscriptionFails() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenThrow(new StreamException(exceptionMessage)); assertThatThrownBy( @@ -439,7 +478,8 @@ void subscribeShouldThrowStreamExceptionWhenClientSubscribeReturnsNull() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(null); assertThatThrownBy( @@ -560,7 +600,8 @@ void subscribeShouldSubscribeToStreamAndDispatchMessage_UnsubscribeShouldUnsubsc anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); AtomicInteger messageHandlerCalls = new AtomicInteger(); @@ -578,7 +619,13 @@ void subscribeShouldSubscribeToStreamAndDispatchMessage_UnsubscribeShouldUnsubsc flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -607,7 +654,8 @@ void shouldNotUnsubscribeIfClientIsClosed() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); Runnable closingRunnable = @@ -623,7 +671,13 @@ void shouldNotUnsubscribeIfClientIsClosed() { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); when(client.isOpen()).thenReturn(false); when(client.unsubscribe(subscriptionIdCaptor.getValue())) @@ -644,7 +698,8 @@ void subscribeShouldSubscribeToStreamAndDispatchMessageWithManySubscriptions() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); Map messageHandlerCalls = new ConcurrentHashMap<>(); @@ -668,7 +723,13 @@ void subscribeShouldSubscribeToStreamAndDispatchMessageWithManySubscriptions() { verify(clientFactory, times(1)).client(any()); verify(client, times(ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); Runnable messageToEachSubscription = () -> @@ -705,7 +766,8 @@ void ignoredMessageShouldTriggerMessageProcessing() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); AtomicInteger messageHandlerCalls = new AtomicInteger(); @@ -720,6 +782,11 @@ public int initialCredits() { return 10; } + @Override + public CreditUnit unit() { + return CreditUnit.CHUNK; + } + @Override public MessageProcessedCallback start(Context context) { return flowStrategyCallback; @@ -744,7 +811,13 @@ public MessageProcessedCallback start(Context context) { flowStrategy); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); Byte subId = subscriptionIdCaptor.getValue(); messageIgnoredListener.ignored(subId, 0, 0, 0, flowStrategyCallback); @@ -785,7 +858,8 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { subscriptionCount.incrementAndGet(); @@ -809,7 +883,13 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -833,7 +913,13 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { flowStrategy()); verify(client, times(1 + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -845,7 +931,13 @@ void shouldRedistributeConsumerIfConnectionIsLost() throws Exception { verify(consumer, times(1)).setSubscriptionClient(isNull()); verify(client, times(2 + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(1); messageListener.handle( @@ -884,7 +976,8 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { subscriptionCount.incrementAndGet(); @@ -918,7 +1011,13 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -929,7 +1028,13 @@ void shouldSkipRecoveryIfRecoveryIsAlreadyInProgress() throws Exception { verify(consumer, times(1)).setSubscriptionClient(isNull()); verify(client, times(1 + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); } @Test @@ -952,7 +1057,8 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { subscriptionCount.incrementAndGet(); @@ -973,7 +1079,13 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); coordinator.subscribe( consumerClosedAfterMetadataUpdate, @@ -987,7 +1099,13 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { flowStrategy()); verify(client, times(1 + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); firstMessageListener() @@ -1009,7 +1127,13 @@ void shouldRedistributeConsumerOnMetadataUpdate() throws Exception { // the second consumer does not re-subscribe because it returns it is not open waitAtMost(() -> subscriptionCount.get() == 2 + 1); verify(client, times(2 + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(1); // listener is per manager (connection), so it can have been disposed of, @@ -1057,7 +1181,8 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .then( invocation -> { subscriptionCount.incrementAndGet(); @@ -1078,7 +1203,13 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -1090,7 +1221,13 @@ void shouldRetryRedistributionIfMetadataIsNotUpdatedImmediately() throws Excepti waitAtMost(() -> subscriptionCount.get() == 2); verify(client, times(2)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(1); messageListener.handle( @@ -1127,7 +1264,8 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); AtomicInteger messageHandlerCalls = new AtomicInteger(); @@ -1143,7 +1281,13 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -1156,7 +1300,13 @@ void metadataUpdate_shouldCloseConsumerIfStreamIsDeleted() throws Exception { verify(consumer, times(1)).closeAfterStreamDeletion(); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); verify(client, times(0)).unsubscribe(anyByte()); assertThat(coordinator.managerCount()).isZero(); @@ -1180,7 +1330,8 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); AtomicInteger messageHandlerCalls = new AtomicInteger(); @@ -1196,7 +1347,13 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -1209,7 +1366,13 @@ void metadataUpdate_shouldCloseConsumerIfRetryTimeoutIsReached() throws Exceptio verify(consumer, times(1)).closeAfterStreamDeletion(); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); verify(client, times(0)).unsubscribe(anyByte()); assertThat(coordinator.managerCount()).isZero(); @@ -1229,7 +1392,8 @@ void shouldUseNewClientsForMoreThanMaxSubscriptionsAndCloseClientAfterUnsubscrip anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); when(client.isOpen()).thenReturn(true); @@ -1263,7 +1427,13 @@ void shouldUseNewClientsForMoreThanMaxSubscriptionsAndCloseClientAfterUnsubscrip verify(clientFactory, times(2)).client(any()); verify(client, times(subscriptionCount)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); when(client.unsubscribe(anyByte())).thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); @@ -1301,7 +1471,8 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); int extraSubscriptionCount = ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT / 5; @@ -1324,7 +1495,13 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E // the extra is allocated on another client from the same pool verify(clientFactory, times(2)).client(any()); verify(client, times(subscriptionCount)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); // let's kill the first client connection shutdownListeners @@ -1348,7 +1525,13 @@ void shouldRemoveClientSubscriptionManagerFromPoolAfterConnectionDies() throws E verify(clientFactory, times(2 + 1)).client(any()); verify(client, times(subscriptionCount + ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); } @Test @@ -1366,7 +1549,8 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); int extraSubscriptionCount = ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT / 5; @@ -1389,7 +1573,13 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t // the extra is allocated on another client from the same pool verify(clientFactory, times(2)).client(any()); verify(client, times(subscriptionCount)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); ConsumerCoordinatorInfo info = MonitoringTestUtils.extract(coordinator); assertThat(info.nodesConnected()); @@ -1420,7 +1610,13 @@ void shouldRemoveClientSubscriptionManagerFromPoolIfEmptyAfterMetadataUpdate() t // no more client creation verify(clientFactory, times(2)).client(any()); verify(client, times(subscriptionCount + ConsumersCoordinator.MAX_SUBSCRIPTIONS_PER_CLIENT + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); info = MonitoringTestUtils.extract(coordinator); assertThat(info.nodesConnected()).hasSize(1); @@ -1452,7 +1648,8 @@ void shouldRestartWhereItLeftOffAfterDisruption(Consumer configur anyString(), offsetSpecificationArgumentCaptor.capture(), anyInt(), - subscriptionPropertiesArgumentCaptor.capture())) + subscriptionPropertiesArgumentCaptor.capture(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); Runnable closingRunnable = @@ -1615,7 +1838,13 @@ void shouldUseStoredOffsetOnRecovery(Consumer configur flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(0) .isEqualTo(OffsetSpecification.next()); @@ -1636,7 +1865,13 @@ void shouldUseStoredOffsetOnRecovery(Consumer configur Thread.sleep(retryDelay.toMillis() * 5); verify(client, times(2)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(offsetSpecificationArgumentCaptor.getAllValues()) .element(1) @@ -1683,7 +1918,8 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( a -> { subscriptionCount.incrementAndGet(); @@ -1702,7 +1938,13 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); coordinator.subscribe( consumer, @@ -1716,7 +1958,13 @@ void shouldRetryAssignmentOnRecoveryTimeout() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1 + 1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); this.shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -1752,7 +2000,8 @@ void shouldRetryAssignmentOnRecoveryStreamNotAvailableFailure() throws Exception anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { subscriptionCount.incrementAndGet(); @@ -1781,7 +2030,13 @@ void shouldRetryAssignmentOnRecoveryStreamNotAvailableFailure() throws Exception flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); this.shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -1815,7 +2070,8 @@ void shouldRetryAssignmentOnRecoveryCandidateLookupFailure() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { // first subscription @@ -1846,7 +2102,13 @@ void shouldRetryAssignmentOnRecoveryCandidateLookupFailure() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); this.shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -1879,7 +2141,8 @@ void shouldRetryAssignmentOnRecoveryConnectionTimeout() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { subscriptionCount.incrementAndGet(); @@ -1898,7 +2161,13 @@ void shouldRetryAssignmentOnRecoveryConnectionTimeout() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); this.shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -1918,7 +2187,8 @@ void subscribeUnsubscribeInDifferentThreadsShouldNotDeadlock() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); when(client.unsubscribe(anyByte())).thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); @@ -1992,7 +2262,8 @@ void consumerShouldBeCreatedProperlyIfManagerClientIsRetried() { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenReturn(new Client.Response(Constants.RESPONSE_CODE_OK)); AtomicInteger messageHandlerCalls = new AtomicInteger(); @@ -2009,7 +2280,13 @@ void consumerShouldBeCreatedProperlyIfManagerClientIsRetried() { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -2050,7 +2327,8 @@ void shouldRetryUntilReplicaIsAvailableWhenForceReplicaIsOn() throws Exception { anyString(), any(OffsetSpecification.class), anyInt(), - anyMap())) + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) .thenAnswer( invocation -> { subscriptionCount.incrementAndGet(); @@ -2080,7 +2358,13 @@ void shouldRetryUntilReplicaIsAvailableWhenForceReplicaIsOn() throws Exception { flowStrategy()); verify(clientFactory, times(1)).client(any()); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(0); messageListener.handle( @@ -2093,7 +2377,13 @@ void shouldRetryUntilReplicaIsAvailableWhenForceReplicaIsOn() throws Exception { assertThat(messageHandlerCalls.get()).isEqualTo(1); verify(client, times(1)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); shutdownListener.handle( new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); @@ -2109,7 +2399,13 @@ void shouldRetryUntilReplicaIsAvailableWhenForceReplicaIsOn() throws Exception { verify(consumer, times(1)).setSubscriptionClient(isNull()); verify(client, times(2)) - .subscribe(anyByte(), anyString(), any(OffsetSpecification.class), anyInt(), anyMap()); + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class)); assertThat(messageHandlerCalls.get()).isEqualTo(1); messageListener.handle( @@ -2173,6 +2469,220 @@ void pickSlotTest() { assertThat(index).isEqualTo(5); } + @Test + void byteBasedStrategySubscribesWithByteCreditUnitAndWindowAsInitialCredits() { + when(locator.metadata("stream")).thenReturn(metadata(leader(), replica())); + when(clientFactory.client(any())).thenReturn(client); + when(client.byteCreditSupported()).thenReturn(true); + when(client.subscribe( + subscriptionIdCaptor.capture(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) + .thenReturn(responseOk()); + + coordinator.subscribe( + consumer, + "stream", + OffsetSpecification.first(), + null, + NO_OP_SUBSCRIPTION_LISTENER, + NO_OP_TRACKING_CLOSING_CALLBACK, + (offset, message) -> {}, + Collections.emptyMap(), + ConsumerFlowStrategy.creditOnChunkArrival(ByteCapacity.B(1000))); + + verify(client, times(1)) + .subscribe( + anyByte(), + anyString(), + any(OffsetSpecification.class), + eq(1000), + anyMap(), + eq(ConsumerFlowStrategy.CreditUnit.BYTE)); + } + + @Test + void byteBasedStrategyFailsWithStreamExceptionWhenBrokerDoesNotSupportByteCredit() { + when(locator.metadata("stream")).thenReturn(metadata(leader(), replica())); + when(clientFactory.client(any())).thenReturn(client); + when(client.byteCreditSupported()).thenReturn(false); + + assertThatThrownBy( + () -> + coordinator.subscribe( + consumer, + "stream", + OffsetSpecification.first(), + null, + NO_OP_SUBSCRIPTION_LISTENER, + NO_OP_TRACKING_CLOSING_CALLBACK, + (offset, message) -> {}, + Collections.emptyMap(), + ConsumerFlowStrategy.creditOnChunkArrival(ByteCapacity.kB(1)))) + .isInstanceOf(StreamException.class) + .isNotInstanceOf(IllegalStateException.class); + // must not loop looking for another manager on a node that will never support byte credit + verify(clientFactory, times(1)).client(any()); + } + + @Test + void byteBasedGrantsAreBatchedAcrossManyChunks() { + when(locator.metadata("stream")).thenReturn(metadata(leader(), replica())); + when(clientFactory.client(any())).thenReturn(client); + when(client.byteCreditSupported()).thenReturn(true); + when(consumer.isOpen()).thenReturn(true); + when(client.subscribe( + subscriptionIdCaptor.capture(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) + .thenReturn(responseOk()); + + coordinator.subscribe( + consumer, + "stream", + OffsetSpecification.first(), + null, + NO_OP_SUBSCRIPTION_LISTENER, + NO_OP_TRACKING_CLOSING_CALLBACK, + (offset, message) -> {}, + Collections.emptyMap(), + ConsumerFlowStrategy.creditOnChunkArrival(ByteCapacity.B(1000))); + + byte subId = subscriptionIdCaptor.getValue(); + int chunkCost = 100; + int chunkCount = 20; + for (int i = 0; i < chunkCount; i++) { + chunkListener.handle(client, subId, i, 1, chunkCost, chunkCost); + } + + ArgumentCaptor creditCaptor = ArgumentCaptor.forClass(Integer.class); + verify(client, atLeastOnce()) + .credit(eq(subId), creditCaptor.capture(), eq(ConsumerFlowStrategy.CreditUnit.BYTE)); + int totalGranted = creditCaptor.getAllValues().stream().mapToInt(Integer::intValue).sum(); + // batching: fewer Credit frames than chunks, and never more bytes granted than received + assertThat(creditCaptor.getAllValues()).hasSizeLessThan(chunkCount); + assertThat(totalGranted).isGreaterThan(0).isLessThanOrEqualTo(chunkCost * chunkCount); + } + + @Test + void creditIsGrantedOnceEvenIfCustomStrategyCallsCreditsTwiceForTheSameChunk() { + when(locator.metadata("stream")).thenReturn(metadata(leader(), replica())); + when(clientFactory.client(any())).thenReturn(client); + when(client.byteCreditSupported()).thenReturn(true); + when(consumer.isOpen()).thenReturn(true); + when(client.subscribe( + subscriptionIdCaptor.capture(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) + .thenReturn(responseOk()); + + ConsumerFlowStrategy flowStrategy = + new ConsumerFlowStrategy() { + @Override + public int initialCredits() { + return 1000; + } + + @Override + public CreditUnit unit() { + return CreditUnit.BYTE; + } + + @Override + public MessageProcessedCallback start(Context context) { + // a buggy custom strategy releasing the same chunk's credit twice + context.credits(1); + context.credits(1); + return messageContext -> {}; + } + }; + + coordinator.subscribe( + consumer, + "stream", + OffsetSpecification.first(), + null, + NO_OP_SUBSCRIPTION_LISTENER, + NO_OP_TRACKING_CLOSING_CALLBACK, + (offset, message) -> {}, + Collections.emptyMap(), + flowStrategy); + + byte subId = subscriptionIdCaptor.getValue(); + // window = 1000, flush threshold = 750; chunk1 (100) does not flush on its own, chunk2 (200) + // pushes credit to the threshold and flushes whatever is pending + chunkListener.handle(client, subId, 0, 1, 100, 100); + chunkListener.handle(client, subId, 1, 1, 200, 200); + + // 100, not 200: the second credits() call for chunk1 must not have doubled its release + verify(client, times(1)).credit(subId, 100, ConsumerFlowStrategy.CreditUnit.BYTE); + verify(client, never()).credit(eq(subId), eq(200), eq(ConsumerFlowStrategy.CreditUnit.BYTE)); + } + + @Test + void byteBasedAccountantIsResetOnRecovery() throws Exception { + scheduledExecutorService = createScheduledExecutorService(); + when(environment.scheduledExecutorService()).thenReturn(scheduledExecutorService); + when(environment.recoveryBackOffDelayPolicy()) + .thenReturn(BackOffDelayPolicy.fixed(Duration.ofMillis(50))); + when(consumer.isOpen()).thenReturn(true); + when(locator.metadata("stream")).thenReturn(metadata(null, replica())); + when(clientFactory.client(any())).thenReturn(client); + when(client.byteCreditSupported()).thenReturn(true); + + AtomicInteger subscriptionCount = new AtomicInteger(0); + when(client.subscribe( + subscriptionIdCaptor.capture(), + anyString(), + any(OffsetSpecification.class), + anyInt(), + anyMap(), + any(ConsumerFlowStrategy.CreditUnit.class))) + .thenAnswer( + invocation -> { + subscriptionCount.incrementAndGet(); + return responseOk(); + }); + + int window = 1000; + coordinator.subscribe( + consumer, + "stream", + OffsetSpecification.first(), + null, + NO_OP_SUBSCRIPTION_LISTENER, + NO_OP_TRACKING_CLOSING_CALLBACK, + (offset, message) -> {}, + Collections.emptyMap(), + ConsumerFlowStrategy.creditOnChunkArrival(ByteCapacity.B(window))); + + byte subId = subscriptionIdCaptor.getValue(); + // build up unflushed pending, below the flush threshold, without triggering a grant + chunkListener.handle(client, subId, 0, 1, 200, 200); + verify(client, never()).credit(eq(subId), anyInt(), eq(ConsumerFlowStrategy.CreditUnit.BYTE)); + + shutdownListener.handle( + new Client.ShutdownContext(Client.ShutdownContext.ShutdownReason.UNKNOWN)); + waitAtMost(() -> subscriptionCount.get() == 2); + + byte subId2 = + subscriptionIdCaptor.getAllValues().get(subscriptionIdCaptor.getAllValues().size() - 1); + // a chunk that exactly drains a freshly reset window: if the mirror had not been reset, the + // stale pending bytes from before the disruption would be granted as well + chunkListener.handle(client, subId2, 1, 1, window, window); + verify(client, times(1)).credit(subId2, window, ConsumerFlowStrategy.CreditUnit.BYTE); + verify(client, never()).credit(eq(subId2), eq(200), eq(ConsumerFlowStrategy.CreditUnit.BYTE)); + } + static Client.Broker leader() { return new Client.Broker("leader", -1); } diff --git a/src/test/java/com/rabbitmq/stream/impl/CreditEveryNthChunkConsumerFlowStrategyTest.java b/src/test/java/com/rabbitmq/stream/impl/CreditEveryNthChunkConsumerFlowStrategyTest.java index 7841944501..9d2bbaadce 100644 --- a/src/test/java/com/rabbitmq/stream/impl/CreditEveryNthChunkConsumerFlowStrategyTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/CreditEveryNthChunkConsumerFlowStrategyTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2025 Broadcom. All Rights Reserved. +// Copyright (c) 2025-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -21,7 +21,9 @@ import com.rabbitmq.stream.ConsumerFlowStrategy; import com.rabbitmq.stream.ConsumerFlowStrategy.Context; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; @@ -29,6 +31,11 @@ public class CreditEveryNthChunkConsumerFlowStrategyTest { AtomicInteger requestedCredits = new AtomicInteger(); + @Test + void unitIsAlwaysChunk() { + assertThat(build(10, 5).unit()).isEqualTo(CreditUnit.CHUNK); + } + @ParameterizedTest @CsvSource({"10,6", "1,1", "1,0", "10,0"}) void invalidArguments(int initialCredits, int limit) { @@ -71,6 +78,11 @@ public long messageCount() { public long chunkId() { throw new UnsupportedOperationException(); } + + @Override + public long chunkByteCount() { + throw new UnsupportedOperationException(); + } }; } } diff --git a/src/test/java/com/rabbitmq/stream/impl/DeliveryTest.java b/src/test/java/com/rabbitmq/stream/impl/DeliveryTest.java index dd56d2f7fd..08346e5ddf 100644 --- a/src/test/java/com/rabbitmq/stream/impl/DeliveryTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/DeliveryTest.java @@ -165,7 +165,7 @@ public TestConfig(long chunkOffset, long subscriptionOffset) { bb, client, ctx, - (client, subscriptionId, offset, messageCount, sizeOfData) -> { + (client, subscriptionId, offset, messageCount, sizeOfData, chunkByteCount) -> { assertThat(messageCount).isEqualTo(nbMessages); chunkCountInCallback.incrementAndGet(); return null; @@ -181,7 +181,8 @@ public TestConfig(long chunkOffset, long subscriptionOffset) { }, NO_OP_CODEC, ChunkChecksum.NO_OP, - NoOpMetricsCollector.SINGLETON); + NoOpMetricsCollector.SINGLETON, + dataSize); long expectedMessageCount = nbMessages - (subscriptionOffset - chunkOffset); long expectedFilteredMessageCount = nbMessages - expectedMessageCount; diff --git a/src/test/java/com/rabbitmq/stream/impl/MessageCountConsumerFlowStrategyTest.java b/src/test/java/com/rabbitmq/stream/impl/MessageCountConsumerFlowStrategyTest.java index bb4d8f3c12..12f9ce27e7 100644 --- a/src/test/java/com/rabbitmq/stream/impl/MessageCountConsumerFlowStrategyTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/MessageCountConsumerFlowStrategyTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2023-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2023-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -18,7 +18,9 @@ import static java.util.stream.LongStream.range; import static org.assertj.core.api.Assertions.assertThat; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.ConsumerFlowStrategy; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; @@ -54,6 +56,19 @@ void smallChunksAndSmallRatiosShouldCredit() { assertThat(requestedCredits).hasValue(1); } + @Test + void byteBasedStrategyBehavesLikeChunkBased() { + ConsumerFlowStrategy strategy = creditOnProcessedMessageCount(ByteCapacity.B(10), 0.5); + assertThat(strategy.unit()).isEqualTo(CreditUnit.BYTE); + assertThat(strategy.initialCredits()).isEqualTo(10); + long messageCount = 1000; + ConsumerFlowStrategy.MessageProcessedCallback callback = strategy.start(context(messageCount)); + range(0, messageCount / 2 - 1).forEach(ignored -> callback.processed(null)); + assertThat(requestedCredits).hasValue(0); + callback.processed(null); + assertThat(requestedCredits).hasValue(1); + } + ConsumerFlowStrategy build(double ratio) { return creditOnProcessedMessageCount(10, ratio); } @@ -75,6 +90,11 @@ public long messageCount() { public long chunkId() { return 0; } + + @Override + public long chunkByteCount() { + return 0; + } }; } } diff --git a/src/test/java/com/rabbitmq/stream/impl/OffsetTest.java b/src/test/java/com/rabbitmq/stream/impl/OffsetTest.java index 2de14230ae..9fd40a2dca 100644 --- a/src/test/java/com/rabbitmq/stream/impl/OffsetTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/OffsetTest.java @@ -162,7 +162,12 @@ void offsetTypeLastShouldReturnLastChunk( new Client.ClientParameters() .sslContext(sslContext) .chunkListener( - (client1, subscriptionId, offset12, messageCount1, dataSize) -> { + (client1, + subscriptionId, + offset12, + messageCount1, + dataSize, + chunkByteCount) -> { client1.credit(subscriptionId, 1); chunkOffset.compareAndSet(-1, offset12); chunkCount.incrementAndGet(); @@ -203,7 +208,12 @@ void amqpOffsetTypeLastShouldReturnLastChunk() throws Exception { cf.get( new Client.ClientParameters() .chunkListener( - (client1, subscriptionId, offset12, messageCount1, dataSize) -> { + (client1, + subscriptionId, + offset12, + messageCount1, + dataSize, + chunkByteCount) -> { client1.credit(subscriptionId, 1); chunkOffset.compareAndSet(-1, offset12); return null; diff --git a/src/test/java/com/rabbitmq/stream/impl/SacStreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/SacStreamConsumerTest.java index 538447f17e..59a01b7515 100644 --- a/src/test/java/com/rabbitmq/stream/impl/SacStreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/SacStreamConsumerTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2022-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2022-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -21,8 +21,10 @@ import static com.rabbitmq.stream.impl.TestUtils.waitAtMost; import static org.assertj.core.api.Assertions.assertThat; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.Cli; import com.rabbitmq.stream.Consumer; +import com.rabbitmq.stream.ConsumerFlowStrategy; import com.rabbitmq.stream.Environment; import com.rabbitmq.stream.EnvironmentBuilder; import com.rabbitmq.stream.MessageHandler; @@ -241,6 +243,30 @@ void externalTrackingSecondConsumerShouldTakeOverWhereTheFirstOneLeftOff() throw assertThat(cf.get().queryOffset(consumerName, stream).getOffset()).isZero(); } + @Test + @TestUtils.DisabledIfByteCreditNotSupported + void singleActiveConsumerWorksWithAByteBasedStrategy() throws Exception { + int messageCount = 10000; + AtomicInteger receivedMessageCount = new AtomicInteger(); + Consumer consumer = + environment.consumerBuilder().stream(stream) + .name("foo") + .singleActiveConsumer() + .flow() + .strategy(ConsumerFlowStrategy.creditOnChunkArrival(ByteCapacity.kB(64))) + .builder() + .messageHandler((context, message) -> receivedMessageCount.incrementAndGet()) + .offset(OffsetSpecification.first()) + .autoTrackingStrategy() + .builder() + .build(); + + publishAndWaitForConfirms(cf, messageCount, stream); + waitAtMost(() -> receivedMessageCount.get() == messageCount); + + consumer.close(); + } + public static Stream> activeConsumerShouldGetUpdateNotificationAfterDisruption() { return Stream.of( diff --git a/src/test/java/com/rabbitmq/stream/impl/ServerFrameHandlerTest.java b/src/test/java/com/rabbitmq/stream/impl/ServerFrameHandlerTest.java index 7ad3a56173..9f70d681f2 100644 --- a/src/test/java/com/rabbitmq/stream/impl/ServerFrameHandlerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/ServerFrameHandlerTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2020-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2020-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -14,14 +14,111 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; +import static com.rabbitmq.stream.impl.ServerFrameHandler.CHUNK_COST_DELIVER_V1_OFFSET; +import static com.rabbitmq.stream.impl.ServerFrameHandler.CHUNK_COST_DELIVER_V2_OFFSET; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.when; +import com.rabbitmq.stream.ChunkChecksum; +import com.rabbitmq.stream.Codec; import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.Message; +import com.rabbitmq.stream.MessageBuilder; +import com.rabbitmq.stream.impl.ServerFrameHandler.DeliverVersion1FrameHandler; +import com.rabbitmq.stream.impl.ServerFrameHandler.DeliverVersion2FrameHandler; import com.rabbitmq.stream.impl.ServerFrameHandler.FrameHandlerInfo; +import com.rabbitmq.stream.metrics.NoOpMetricsCollector; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.UnpooledByteBufAllocator; +import io.netty.channel.ChannelHandlerContext; +import java.lang.reflect.Field; +import java.util.concurrent.atomic.AtomicLong; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; public class ServerFrameHandlerTest { + static final byte SUBSCRIPTION_ID = 1; + + static final Codec NO_OP_CODEC = + new Codec() { + @Override + public EncodedMessage encode(Message message) { + return null; + } + + @Override + public Message decode(ByteBuf buf, int length) { + buf.skipBytes(length); + return null; + } + + @Override + public MessageBuilder messageBuilder() { + return null; + } + }; + + @Mock Client client; + AutoCloseable mocks; + ChannelHandlerContext ctx; + + @BeforeEach + void setUp() throws Exception { + mocks = MockitoAnnotations.openMocks(this); + ctx = Mockito.mock(ChannelHandlerContext.class); + when(ctx.alloc()).thenReturn(UnpooledByteBufAllocator.DEFAULT); + when(client.extractInitialSubscriptionOffset(SUBSCRIPTION_ID)).thenReturn(-1L); + setClientField("codec", NO_OP_CODEC); + setClientField("chunkChecksum", ChunkChecksum.NO_OP); + setClientField("metricsCollector", NoOpMetricsCollector.SINGLETON); + setClientField( + "messageListener", + (Client.MessageListener) + (subscriptionId, + offset, + chunkTimestamp, + committedChunkId, + chunkContext, + message) -> {}); + setClientField( + "messageIgnoredListener", + (Client.MessageIgnoredListener) + (subscriptionId, offset, chunkTimestamp, committedChunkId, chunkContext) -> {}); + } + + @AfterEach + void tearDown() throws Exception { + mocks.close(); + } + + private void setClientField(String name, Object value) throws Exception { + Field field = Client.class.getDeclaredField(name); + field.setAccessible(true); + field.set(client, value); + } + + private static ByteBuf chunkFrameBody(byte[] messageBody) { + ByteBuf bb = Utils.byteBufAllocator().buffer(256); + bb.writeByte(1) // magic and version + .writeByte(0) // chunk type, always 0 in our case + .writeShort(1) // num entries + .writeInt(1) // num records + .writeLong(System.currentTimeMillis()) + .writeLong(0) // epoch + .writeLong(0) // chunk offset + .writeInt(0) // CRC + .writeInt(4 + messageBody.length) // data size + .writeInt(0) // trailer size + .writeInt(0); // 4 reserved bytes + bb.writeInt(messageBody.length).writeBytes(messageBody); + return bb; + } + @Test void commandVersionsHasDeliver() { FrameHandlerInfo deliverInfo = @@ -34,4 +131,61 @@ void commandVersionsHasDeliver() { assertThat(deliverInfo.getMinVersion()).isEqualTo(Constants.VERSION_1); assertThat(deliverInfo.getMaxVersion()).isGreaterThanOrEqualTo(Constants.VERSION_2); } + + @Test + void deliverVersion1ChunkByteCountIsFrameSizeMinus5() throws Exception { + AtomicLong capturedChunkByteCount = new AtomicLong(-1); + setClientField( + "chunkListener", + (Client.ChunkListener) + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + capturedChunkByteCount.set(chunkByteCount); + return null; + }); + + ByteBuf bb = Utils.byteBufAllocator().buffer(256); + bb.writeShort(Utils.encodeRequestCode(Constants.COMMAND_DELIVER)) + .writeShort(Constants.VERSION_1) + .writeByte(SUBSCRIPTION_ID); + ByteBuf chunk = chunkFrameBody(new byte[20]); + bb.writeBytes(chunk); + chunk.release(); + + int frameSize = bb.readableBytes(); + bb.readShort(); // command key + bb.readShort(); // command version + + new DeliverVersion1FrameHandler().handle(client, frameSize, ctx, bb); + + assertThat(capturedChunkByteCount.get()).isEqualTo(frameSize - CHUNK_COST_DELIVER_V1_OFFSET); + } + + @Test + void deliverVersion2ChunkByteCountIsFrameSizeMinus13() throws Exception { + AtomicLong capturedChunkByteCount = new AtomicLong(-1); + setClientField( + "chunkListener", + (Client.ChunkListener) + (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + capturedChunkByteCount.set(chunkByteCount); + return null; + }); + + ByteBuf bb = Utils.byteBufAllocator().buffer(256); + bb.writeShort(Utils.encodeRequestCode(Constants.COMMAND_DELIVER)) + .writeShort(Constants.VERSION_2) + .writeByte(SUBSCRIPTION_ID) + .writeLong(0); // committed chunk id + ByteBuf chunk = chunkFrameBody(new byte[20]); + bb.writeBytes(chunk); + chunk.release(); + + int frameSize = bb.readableBytes(); + bb.readShort(); // command key + bb.readShort(); // command version + + new DeliverVersion2FrameHandler().handle(client, frameSize, ctx, bb); + + assertThat(capturedChunkByteCount.get()).isEqualTo(frameSize - CHUNK_COST_DELIVER_V2_OFFSET); + } } diff --git a/src/test/java/com/rabbitmq/stream/impl/StompInteroperabilityTest.java b/src/test/java/com/rabbitmq/stream/impl/StompInteroperabilityTest.java index b74135895c..154d8bb3e4 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StompInteroperabilityTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StompInteroperabilityTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2021-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2021-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -523,7 +523,12 @@ void offsetTypeLastShouldStartConsumingFromTheLastChunk() throws Exception { cf.get( new Client.ClientParameters() .chunkListener( - (client1, subscriptionId, offset12, messageCount1, dataSize) -> { + (client1, + subscriptionId, + offset12, + messageCount1, + dataSize, + chunkByteCount) -> { client1.credit(subscriptionId, 1); chunkOffset.compareAndSet(-1, offset12); return null; diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java index e3d79afd96..f243831aa7 100644 --- a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java @@ -15,6 +15,8 @@ package com.rabbitmq.stream.impl; import static com.rabbitmq.stream.ConsumerFlowStrategy.creditEveryNthChunk; +import static com.rabbitmq.stream.ConsumerFlowStrategy.creditOnChunkArrival; +import static com.rabbitmq.stream.ConsumerFlowStrategy.creditOnProcessedMessageCount; import static com.rabbitmq.stream.ConsumerFlowStrategy.creditWhenHalfMessagesProcessed; import static com.rabbitmq.stream.impl.Assertions.assertThat; import static com.rabbitmq.stream.impl.TestUtils.CountDownLatchConditions.completed; @@ -33,10 +35,12 @@ import com.rabbitmq.stream.Address; import com.rabbitmq.stream.BackOffDelayPolicy; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.Cli; import com.rabbitmq.stream.ConfirmationHandler; import com.rabbitmq.stream.Consumer; import com.rabbitmq.stream.ConsumerBuilder; +import com.rabbitmq.stream.ConsumerFlowStrategy; import com.rabbitmq.stream.Environment; import com.rabbitmq.stream.EnvironmentBuilder; import com.rabbitmq.stream.Message; @@ -53,6 +57,7 @@ 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.DisabledIfByteCreditNotSupported; import com.rabbitmq.stream.impl.TestUtils.DisabledIfRabbitMqCtlNotSet; import com.rabbitmq.stream.impl.TestUtils.Sync; import io.netty.channel.ChannelOption; @@ -316,6 +321,208 @@ void asynchronousProcessingWithFlowControl() { } } + @ParameterizedTest + @MethodSource + @DisabledIfByteCreditNotSupported + void consumeBacklogToCompletionWithByteBasedStrategies(ConsumerFlowStrategy strategy) + throws Exception { + int messageCount = 10_000; + publishAndWaitForConfirms(cf, messageCount, stream); + + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + Consumer consumer = + environment.consumerBuilder().stream(stream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(strategy) + .builder() + .messageHandler( + (context, message) -> { + // harmless no-op for the arrival-based strategy, which ignores it + context.processed(); + consumeLatch.countDown(); + }) + .build(); + + org.assertj.core.api.Assertions.assertThat(consumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + consumer.close(); + } + + static Stream consumeBacklogToCompletionWithByteBasedStrategies() { + return Stream.of( + creditOnChunkArrival(ByteCapacity.kB(64)), + creditWhenHalfMessagesProcessed(ByteCapacity.kB(64)), + creditOnProcessedMessageCount(ByteCapacity.kB(64), 0.3), + // window smaller than a single chunk: the broker still delivers it (overshoot rule) + creditOnChunkArrival(ByteCapacity.B(1))); + } + + @Test + @DisabledIfByteCreditNotSupported + void chunksLargerThanHalfWindowAreConsumedToCompletion() throws Exception { + int messageCount = 200; + int bodySize = 20_000; + // close to one chunk's size, so a chunk is always comfortably above half the window + ByteCapacity window = ByteCapacity.kB(20); + publishAndWaitForConfirms( + cf, builder -> builder.addData(new byte[bodySize]).build(), messageCount, stream); + + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + Consumer consumer = + environment.consumerBuilder().stream(stream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(creditWhenHalfMessagesProcessed(window)) + .builder() + .messageHandler( + (context, message) -> { + context.processed(); + consumeLatch.countDown(); + }) + .build(); + + org.assertj.core.api.Assertions.assertThat(consumeLatch.await(20, TimeUnit.SECONDS)).isTrue(); + + consumer.close(); + } + + @Test + @DisabledIfRabbitMqCtlNotSet + @DisabledIfByteCreditNotSupported + void byteBasedConsumerRecoversAfterConnectionKill() throws Exception { + int messageCount = 50_000; + publishAndWaitForConfirms(cf, messageCount, stream); + + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + AtomicInteger receivedMessageCount = new AtomicInteger(); + Consumer consumer = + environment.consumerBuilder().stream(stream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(creditOnChunkArrival(ByteCapacity.kB(4))) + .builder() + .messageHandler( + (context, message) -> { + receivedMessageCount.incrementAndGet(); + consumeLatch.countDown(); + }) + .build(); + + waitAtMost(() -> receivedMessageCount.get() > 0); + + Cli.killConnection("rabbitmq-stream-consumer-0"); + + // the accountant is reset on the reconnection, so the mirror is correct from a full window + latchAssert(consumeLatch).completes(recoveryInitialDelay.plusSeconds(10)); + + consumer.close(); + } + + @Test + @DisabledIfByteCreditNotSupported + void byteBasedConsumerWithProcessedNeverCalledStopsThenResumes() throws Exception { + int messageCount = 5_000; + int bodySize = 200; + ByteCapacity window = ByteCapacity.kB(20); + publishAndWaitForConfirms( + cf, builder -> builder.addData(new byte[bodySize]).build(), messageCount, stream); + + List heldContexts = synchronizedList(new ArrayList<>()); + AtomicInteger receivedMessageCount = new AtomicInteger(); + AtomicBoolean processingStarted = new AtomicBoolean(false); + Consumer consumer = + environment.consumerBuilder().stream(stream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(creditWhenHalfMessagesProcessed(window)) + .builder() + .messageHandler( + (context, message) -> { + receivedMessageCount.incrementAndGet(); + if (processingStarted.get()) { + context.processed(); + } else { + heldContexts.add(context); + } + }) + .build(); + + waitAtMost(() -> receivedMessageCount.get() > 0); + waitUntilStable(receivedMessageCount::get); + + org.assertj.core.api.Assertions.assertThat(receivedMessageCount.get()).isLessThan(messageCount); + + // future arrivals now process themselves, the held backlog is flushed once + processingStarted.set(true); + heldContexts.forEach(MessageHandler.Context::processed); + waitAtMost(() -> receivedMessageCount.get() == messageCount); + + consumer.close(); + } + + @Test + @DisabledIfByteCreditNotSupported + void feedbackPropertyBoundsUnprocessedBytesUnlikeArrivalBasedStrategy(TestInfo info) + throws Exception { + int messageCount = 5_000; + int bodySize = 200; + ByteCapacity window = ByteCapacity.kB(20); + + String processedStream = streamName(info) + "-processed"; + String arrivalStream = streamName(info) + "-arrival"; + try { + environment.streamCreator().stream(processedStream).create(); + environment.streamCreator().stream(arrivalStream).create(); + publishAndWaitForConfirms( + cf, + builder -> builder.addData(new byte[bodySize]).build(), + messageCount, + processedStream); + publishAndWaitForConfirms( + cf, builder -> builder.addData(new byte[bodySize]).build(), messageCount, arrivalStream); + + // creditWhenHalfMessagesProcessed: bytes received but not processed are bounded by the + // window, so consumption stalls well before the whole backlog is delivered + AtomicInteger processedStreamReceived = new AtomicInteger(); + Consumer processedStreamConsumer = + environment.consumerBuilder().stream(processedStream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(creditWhenHalfMessagesProcessed(window)) + .builder() + // processed() deliberately never called, simulating a stuck application + .messageHandler((context, message) -> processedStreamReceived.incrementAndGet()) + .build(); + + waitAtMost(() -> processedStreamReceived.get() > 0); + waitUntilStable(processedStreamReceived::get); + org.assertj.core.api.Assertions.assertThat(processedStreamReceived.get()) + .isLessThan(messageCount); + processedStreamConsumer.close(); + + // creditOnChunkArrival: credit is granted on arrival, not on processing, so the whole + // backlog is delivered regardless of processing speed + CountDownLatch arrivalConsumeLatch = new CountDownLatch(messageCount); + Consumer arrivalStreamConsumer = + environment.consumerBuilder().stream(arrivalStream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(creditOnChunkArrival(window)) + .builder() + // processed() is ignored by this strategy, never called here either + .messageHandler((context, message) -> arrivalConsumeLatch.countDown()) + .build(); + + org.assertj.core.api.Assertions.assertThat(arrivalConsumeLatch.await(10, TimeUnit.SECONDS)) + .isTrue(); + arrivalStreamConsumer.close(); + } finally { + environment.deleteStream(processedStream); + environment.deleteStream(arrivalStream); + } + } + @Test void closeOnCondition() throws Exception { int messageCount = 50_000; diff --git a/src/test/java/com/rabbitmq/stream/impl/SubEntryDecompressionTest.java b/src/test/java/com/rabbitmq/stream/impl/SubEntryDecompressionTest.java index 680e3eb2b9..0e6ea8ac45 100644 --- a/src/test/java/com/rabbitmq/stream/impl/SubEntryDecompressionTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/SubEntryDecompressionTest.java @@ -267,13 +267,14 @@ private void deliver(ByteBuf bb, ChannelHandlerContext ctx, AtomicInteger messag bb, client, ctx, - (client, subscriptionId, offset, count, sizeOfData) -> null, + (client, subscriptionId, offset, count, sizeOfData, chunkByteCount) -> null, (subscriptionId, offset, chunkTimestamp, committedChunkId, chunkContext, message) -> messageCount.incrementAndGet(), (subscriptionId, offset, chunkTimestamp, committedChunkId, chunkContext) -> {}, NO_OP_CODEC, ChunkChecksum.NO_OP, - NoOpMetricsCollector.SINGLETON); + NoOpMetricsCollector.SINGLETON, + 0); } @Test @@ -389,7 +390,7 @@ public MessageBuilder messageBuilder() { bb, client, ctx, - (client, subscriptionId, offset, count, sizeOfData) -> null, + (client, subscriptionId, offset, count, sizeOfData, chunkByteCount) -> null, (subscriptionId, offset, chunkTimestamp, @@ -399,7 +400,8 @@ public MessageBuilder messageBuilder() { (subscriptionId, offset, chunkTimestamp, committedChunkId, chunkContext) -> {}, corruptCodec, ChunkChecksum.NO_OP, - NoOpMetricsCollector.SINGLETON)) + NoOpMetricsCollector.SINGLETON, + 0)) .isInstanceOf(StreamException.class); assertThat(allocator.requests).hasSize(1); diff --git a/src/test/java/com/rabbitmq/stream/impl/SuperStreamConsumerTest.java b/src/test/java/com/rabbitmq/stream/impl/SuperStreamConsumerTest.java index 0b72393184..449ac23afd 100644 --- a/src/test/java/com/rabbitmq/stream/impl/SuperStreamConsumerTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/SuperStreamConsumerTest.java @@ -1,4 +1,4 @@ -// Copyright (c) 2021-2025 Broadcom. All Rights Reserved. +// Copyright (c) 2021-2026 Broadcom. All Rights Reserved. // The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. // // This software, the RabbitMQ Stream Java client library, is dual-licensed under the @@ -14,6 +14,7 @@ // info@rabbitmq.com. package com.rabbitmq.stream.impl; +import static com.rabbitmq.stream.ConsumerFlowStrategy.creditOnChunkArrival; import static com.rabbitmq.stream.impl.TestUtils.BrokerVersion.RABBITMQ_3_11_11; import static com.rabbitmq.stream.impl.TestUtils.b; import static com.rabbitmq.stream.impl.TestUtils.declareSuperStreamTopology; @@ -24,13 +25,16 @@ import static java.lang.String.format; import static org.assertj.core.api.Assertions.assertThat; +import com.rabbitmq.stream.ByteCapacity; import com.rabbitmq.stream.Consumer; +import com.rabbitmq.stream.ConsumerFlowStrategy; import com.rabbitmq.stream.Environment; import com.rabbitmq.stream.EnvironmentBuilder; import com.rabbitmq.stream.OffsetSpecification; import com.rabbitmq.stream.impl.Client.ClientParameters; import com.rabbitmq.stream.impl.Client.QueryOffsetResponse; import com.rabbitmq.stream.impl.TestUtils.BrokerVersionAtLeast; +import com.rabbitmq.stream.impl.TestUtils.DisabledIfByteCreditNotSupported; import io.netty.channel.EventLoopGroup; import java.nio.charset.StandardCharsets; import java.time.Duration; @@ -329,6 +333,32 @@ void autoOffsetTrackingShouldStoreOffsetZero() { })); } + @Test + @DisabledIfByteCreditNotSupported + void sharedByteBasedStrategyInstanceIsAppliedIndependentlyToEachPartition() { + declareSuperStreamTopology(configurationClient, superStream, partitionCount); + Client client = cf.get(); + List partitions = client.partitions(superStream); + int messageCount = 10000 * partitionCount; + publishToPartitions(cf, partitions, messageCount); + // StreamConsumerBuilder#duplicate copies this very instance for each partition subscription; + // per-subscription accounting must live in ConsumersCoordinator, not in the strategy + ConsumerFlowStrategy sharedStrategy = creditOnChunkArrival(ByteCapacity.kB(64)); + CountDownLatch consumeLatch = new CountDownLatch(messageCount); + Consumer consumer = + environment + .consumerBuilder() + .superStream(superStream) + .offset(OffsetSpecification.first()) + .flow() + .strategy(sharedStrategy) + .builder() + .messageHandler((context, message) -> consumeLatch.countDown()) + .build(); + latchAssert(consumeLatch).completes(); + consumer.close(); + } + @Test @BrokerVersionAtLeast(RABBITMQ_3_11_11) void rebalancedPartitionShouldGetMessagesWhenItComesBackToOriginalConsumerInstance() diff --git a/src/test/java/com/rabbitmq/stream/impl/TestUtils.java b/src/test/java/com/rabbitmq/stream/impl/TestUtils.java index cfe03988d3..ffd3aaa288 100644 --- a/src/test/java/com/rabbitmq/stream/impl/TestUtils.java +++ b/src/test/java/com/rabbitmq/stream/impl/TestUtils.java @@ -34,6 +34,7 @@ import com.rabbitmq.stream.Cli; import com.rabbitmq.stream.Codec; import com.rabbitmq.stream.Constants; +import com.rabbitmq.stream.ConsumerFlowStrategy.CreditUnit; import com.rabbitmq.stream.Message; import com.rabbitmq.stream.MessageBuilder; import com.rabbitmq.stream.StreamException; @@ -543,6 +544,12 @@ static boolean atLeastVersion(String expectedVersion, String currentVersion) { @ExtendWith(DisabledIfFilteringNotSupportedCondition.class) @interface DisabledIfFilteringNotSupported {} + @Target({ElementType.TYPE, ElementType.METHOD}) + @Retention(RetentionPolicy.RUNTIME) + @Documented + @ExtendWith(DisabledIfByteCreditNotSupportedCondition.class) + @interface DisabledIfByteCreditNotSupported {} + @Target({ElementType.TYPE, ElementType.METHOD}) @Retention(RetentionPolicy.RUNTIME) @Documented @@ -707,6 +714,7 @@ public void beforeEach(ExtensionContext context) throws Exception { Client.Response response = client.create(stream); assertThat(response.isOk()).isTrue(); store(context.getRoot()).put("filteringSupported", client.filteringSupported()); + store(context.getRoot()).put("byteCreditSupported", client.byteCreditSupported()); client.close(); store(context).put("testMethodStream", stream); } @@ -881,6 +889,30 @@ public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext con } } + static class DisabledIfByteCreditNotSupportedCondition implements ExecutionCondition { + + @Override + public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext context) { + Boolean byteCreditSupported = + StreamTestInfrastructureExtension.store(context) + .get("byteCreditSupported", Boolean.class); + if (byteCreditSupported == null) { + EventLoopGroup eventLoop = StreamTestInfrastructureExtension.eventLoopGroup(context); + try (Client client = new Client(new ClientParameters().eventLoopGroup(eventLoop))) { + byteCreditSupported = client.byteCreditSupported(); + StreamTestInfrastructureExtension.store(context) + .put("byteCreditSupported", byteCreditSupported); + } + } + + if (byteCreditSupported) { + return ConditionEvaluationResult.enabled("byte credit is supported"); + } else { + return ConditionEvaluationResult.disabled("byte credit is not supported"); + } + } + } + static class DisabledIfRabbitMqCtlNotSetCondition implements ExecutionCondition { @Override @@ -1170,12 +1202,19 @@ public String toString() { } static Client.ChunkListener credit() { - return (client, subscriptionId, offset, messageCount, dataSize) -> { + return (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { client.credit(subscriptionId, 1); return null; }; } + static Client.ChunkListener creditBytes() { + return (client, subscriptionId, offset, messageCount, dataSize, chunkByteCount) -> { + client.credit(subscriptionId, (int) chunkByteCount, CreditUnit.BYTE); + return null; + }; + } + static void waitUntilStable(LongSupplier value) { int sameValueCount = 0; Duration timeout = Duration.ofSeconds(10); diff --git a/src/test/java/com/rabbitmq/stream/impl/TlsTest.java b/src/test/java/com/rabbitmq/stream/impl/TlsTest.java index 45460e9cf3..ad4314b3dd 100644 --- a/src/test/java/com/rabbitmq/stream/impl/TlsTest.java +++ b/src/test/java/com/rabbitmq/stream/impl/TlsTest.java @@ -153,7 +153,7 @@ void publishAndConsumeWithUnverifiedConnection() { CountDownLatch consumedLatch = new CountDownLatch(publishCount); Client.ChunkListener chunkListener = - (client, correlationId, offset, messageCount, dataSize) -> { + (client, correlationId, offset, messageCount, dataSize, chunkByteCount) -> { if (consumedLatch.getCount() != 0) { client.credit(correlationId, 1); }