-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #649 from rabbitmq/dynamic-batch-publishing
Add dynamic-batch publishing option
- Loading branch information
Showing
21 changed files
with
1,126 additions
and
476 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
60 changes: 60 additions & 0 deletions
60
src/main/java/com/rabbitmq/stream/impl/ConcurrencyUtils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| // Copyright (c) 2024 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 java.lang.reflect.InvocationTargetException; | ||
| import java.util.Arrays; | ||
| import java.util.concurrent.Executors; | ||
| import java.util.concurrent.ThreadFactory; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| final class ConcurrencyUtils { | ||
|
|
||
| private static final Logger LOGGER = LoggerFactory.getLogger(ConcurrencyUtils.class); | ||
|
|
||
| private static final ThreadFactory THREAD_FACTORY; | ||
|
|
||
| static { | ||
| if (isJava21OrMore()) { | ||
| LOGGER.debug("Running Java 21 or more, using virtual threads"); | ||
| Class<?> builderClass = | ||
| Arrays.stream(Thread.class.getDeclaredClasses()) | ||
| .filter(c -> "Builder".equals(c.getSimpleName())) | ||
| .findFirst() | ||
| .get(); | ||
| // Reflection code is the same as: | ||
| // Thread.ofVirtual().factory(); | ||
| try { | ||
| Object builder = Thread.class.getDeclaredMethod("ofVirtual").invoke(null); | ||
| THREAD_FACTORY = (ThreadFactory) builderClass.getDeclaredMethod("factory").invoke(builder); | ||
| } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) { | ||
| throw new RuntimeException(e); | ||
| } | ||
| } else { | ||
| THREAD_FACTORY = Executors.defaultThreadFactory(); | ||
| } | ||
| } | ||
|
|
||
| private ConcurrencyUtils() {} | ||
|
|
||
| static ThreadFactory defaultThreadFactory() { | ||
| return THREAD_FACTORY; | ||
| } | ||
|
|
||
| private static boolean isJava21OrMore() { | ||
| return Utils.versionCompare(System.getProperty("java.version"), "21.0") >= 0; | ||
| } | ||
| } |
121 changes: 121 additions & 0 deletions
121
src/main/java/com/rabbitmq/stream/impl/DynamicBatch.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,121 @@ | ||
| // Copyright (c) 2024 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 java.lang.Math.max; | ||
| import static java.lang.Math.min; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.concurrent.BlockingQueue; | ||
| import java.util.concurrent.LinkedBlockingQueue; | ||
| import java.util.concurrent.TimeUnit; | ||
| import java.util.function.BiPredicate; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| final class DynamicBatch<T> implements AutoCloseable { | ||
|
|
||
| private static final Logger LOGGER = LoggerFactory.getLogger(DynamicBatch.class); | ||
| private static final int MIN_BATCH_SIZE = 32; | ||
| private static final int MAX_BATCH_SIZE = 8192; | ||
|
|
||
| private final BlockingQueue<T> requests = new LinkedBlockingQueue<>(); | ||
| private final BiPredicate<List<T>, Boolean> consumer; | ||
| private final int configuredBatchSize; | ||
| private final Thread thread; | ||
|
|
||
| DynamicBatch(BiPredicate<List<T>, Boolean> consumer, int batchSize) { | ||
| this.consumer = consumer; | ||
| this.configuredBatchSize = min(max(batchSize, MIN_BATCH_SIZE), MAX_BATCH_SIZE); | ||
| this.thread = ConcurrencyUtils.defaultThreadFactory().newThread(this::loop); | ||
| this.thread.start(); | ||
| } | ||
|
|
||
| void add(T item) { | ||
| try { | ||
| requests.put(item); | ||
| } catch (InterruptedException e) { | ||
| throw new RuntimeException(e); | ||
| } | ||
| } | ||
|
|
||
| private void loop() { | ||
| State<T> state = new State<>(); | ||
| state.batchSize = this.configuredBatchSize; | ||
| state.items = new ArrayList<>(state.batchSize); | ||
| state.retry = false; | ||
| Thread currentThread = Thread.currentThread(); | ||
| T item; | ||
| while (!currentThread.isInterrupted()) { | ||
| try { | ||
| item = this.requests.poll(100, TimeUnit.MILLISECONDS); | ||
| } catch (InterruptedException e) { | ||
| currentThread.interrupt(); | ||
| return; | ||
| } | ||
| if (item != null) { | ||
| state.items.add(item); | ||
| if (state.items.size() >= state.batchSize) { | ||
| this.maybeCompleteBatch(state, true); | ||
| } else { | ||
| item = this.requests.poll(); | ||
| if (item == null) { | ||
| this.maybeCompleteBatch(state, false); | ||
| } else { | ||
| state.items.add(item); | ||
| if (state.items.size() >= state.batchSize) { | ||
| this.maybeCompleteBatch(state, true); | ||
| } | ||
| } | ||
| } | ||
| } else { | ||
| this.maybeCompleteBatch(state, false); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private static final class State<T> { | ||
|
|
||
| int batchSize; | ||
| List<T> items; | ||
| boolean retry; | ||
| } | ||
|
|
||
| private void maybeCompleteBatch(State<T> state, boolean increaseIfCompleted) { | ||
| try { | ||
| boolean completed = this.consumer.test(state.items, state.retry); | ||
| if (completed) { | ||
| if (increaseIfCompleted) { | ||
| state.batchSize = min(state.batchSize * 2, MAX_BATCH_SIZE); | ||
| } else { | ||
| state.batchSize = max(state.batchSize / 2, MIN_BATCH_SIZE); | ||
| } | ||
| state.items = new ArrayList<>(state.batchSize); | ||
| state.retry = false; | ||
| } else { | ||
| state.retry = true; | ||
| } | ||
| } catch (Exception e) { | ||
| LOGGER.warn("Error during dynamic batch completion: {}", e.getMessage()); | ||
| state.retry = true; | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void close() { | ||
| this.thread.interrupt(); | ||
| } | ||
| } |
Oops, something went wrong.