mvallim · mvallim · Aug 8, 2026 · Aug 8, 2026 · Aug 8, 2026
@@ -105,7 +105,16 @@ repositories {
 | **`linger`**          | **int**     | refers to the time to wait before sending messages out to SNS.                 |
 | **`maxBatchSize`**    | **int**     | refers to the maximum amount of data to be collected before sending the batch. |
 
-**NOTICE**: the buffer of message store in memory is calculate using **`maximumPoolSize`** * **`maxBatchSize`** huge values demand huge memory.
+> [!NOTE]
+> The buffer of message store in memory is calculated using **`maximumPoolSize`** * **`maxBatchSize`**; huge values demand huge memory.
+>
+> **Note on effective capacity:** the default queue implementation (`RingBufferBlockingQueue`) internally
+> rounds its capacity up to the next power of two, to allow fast bitwise index calculation. This means the
+> *actual* allocated capacity may be up to ~2x the value computed above — e.g. `maximumPoolSize=10` and
+> `maxBatchSize=10` yields a requested capacity of 100, but the queue actually allocates 128 slots. If you
+> need to budget memory precisely, use `RingBufferBlockingQueue#remainingCapacity()` (or a `LinkedBlockingQueue`
+> via the [Custom `BlockingQueue`](#custom-blockingqueue) option below, which does not round up) rather than
+> relying on the `maximumPoolSize * maxBatchSize` formula as an exact figure.
 
 #### Custom `BlockingQueue`
 

@@ -34,6 +34,13 @@
  * blocking {@link #put(Object)} and {@link #take()} operations. Other
  * {@link BlockingQueue} methods throw {@link UnsupportedOperationException}.
  *
+ * <p><b>Capacity rounding:</b> the requested capacity is rounded up to the next
+ * power of two (see {@link #RingBufferBlockingQueue(int)}), so that slot indices
+ * can be computed with a bitwise AND ({@code sequence & indexMask}) instead of a
+ * modulo operation. As a result, the <i>effective</i> capacity — reported by
+ * {@link #remainingCapacity()} plus {@link #size()} — may be up to twice the
+ * value that was requested.
+ *
  * @param <E> the type of elements held in this queue
  */
 @SuppressWarnings({ "unchecked", "java:S3078", "java:S1948" })
@@ -83,8 +90,12 @@ public class RingBufferBlockingQueue<E> extends AbstractQueue<E> implements Bloc
   /**
    * Creates a ring buffer with the specified capacity.
    *
-   * @param capacity the maximum number of elements the queue can hold; must be
-   *                 positive
+   * @param capacity the requested capacity; must be positive. This value is
+   *                 rounded up to the next power of two internally — the queue's
+   *                 actual capacity (see {@link #remainingCapacity()}) may end up
+   *                 up to 2x larger than the value passed here. This trade-off
+   *                 enables the ring buffer to use a bitwise mask for index
+   *                 calculation instead of a modulo operation on every put/take.
    * @throws IllegalArgumentException if {@code capacity <= 0}
    */
   public RingBufferBlockingQueue(final int capacity) {

@@ -23,11 +23,14 @@
 import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiFunction;
 import java.util.function.UnaryOperator;
@@ -263,7 +266,7 @@ private boolean maxBatchSizeReached(final BlockingQueue<RequestEntry<E>> request
    * @return true if the request can be added
    */
   private boolean canAddToBatch(final int batchSizeBytes, final int requestEntriesSize, final RequestEntry<E> request) {
-    return (batchSizeBytes < AbstractAmazonSnsConsumer.BATCH_SIZE_BYTES_THRESHOLD)
+    return (batchSizeBytes < BATCH_SIZE_BYTES_THRESHOLD)
       && (requestEntriesSize < topicProperty.getMaxBatchSize())
       && Objects.nonNull(request);
   }
@@ -275,7 +278,7 @@ private boolean canAddToBatch(final int batchSizeBytes, final int requestEntries
    * @return true if the batch is still within the size limit
    */
   private boolean canAddPayload(final int batchSizeBytes) {
-    return batchSizeBytes <= AbstractAmazonSnsConsumer.BATCH_SIZE_BYTES_THRESHOLD;
+    return batchSizeBytes <= BATCH_SIZE_BYTES_THRESHOLD;
   }
 
   /**
@@ -349,5 +352,44 @@ public CompletableFuture<Void> await() {
     });
   }
 
+  /**
+   * Returns a {@link CompletableFuture} that completes once all pending requests have
+   * been processed (i.e., both the pending requests map and the topic requests queue are empty),
+   * bounded by the given timeout.
+   * <p>
+   * Internally reuses {@link #await()} and waits on it via {@link CompletableFuture#get(long, TimeUnit)}
+   * on a separate thread, so the calling thread is never blocked directly. If the timeout elapses
+   * before all pending requests are drained, the returned future completes exceptionally with a
+   * {@link CompletionException} wrapping a {@link java.util.concurrent.TimeoutException}.
+   * <p>
+   * Note that the underlying drain triggered by {@link #await()} is not cancelled when the timeout
+   * elapses; it keeps running in the background until the pending requests and topic requests queue
+   * are actually empty.
+   *
+   * @param timeout the maximum time to wait for all pending requests to be processed
+   * @return a future that completes when all requests are drained, or completes exceptionally
+   *         if {@code timeout} elapses first
+   * @throws NullPointerException if {@code timeout} is {@code null}
+   */
+  @Override
+  public CompletableFuture<Void> await(final Duration timeout) {
+    Objects.requireNonNull(timeout, "timeout cannot be null");
+
+    final CompletableFuture<Void> pending = await();
+
+    return CompletableFuture.runAsync(() -> {
+      try {
+        pending.get(timeout.toNanos(), TimeUnit.NANOSECONDS);
+      } catch (final InterruptedException ex) {
+        Thread.currentThread().interrupt();
+        throw new CompletionException(ex);
+      } catch (final ExecutionException ex) {
+        throw new CompletionException(ex.getCause());
+      } catch (final TimeoutException ex) {
+        throw new CompletionException(ex);
+      }
+    });
+  }
+
 }
 // @formatter:on
@@ -16,10 +16,18 @@
 
 package com.amazon.sns.messaging.lib.core;
 
+import java.util.List;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.amazon.sns.messaging.lib.concurrent.ThreadFactoryProvider;
 import com.amazon.sns.messaging.lib.model.RequestEntry;
 import com.amazon.sns.messaging.lib.model.ResponseFailEntry;
 import com.amazon.sns.messaging.lib.model.ResponseSuccessEntry;
@@ -39,6 +47,9 @@
 @RequiredArgsConstructor(access = AccessLevel.PROTECTED)
 abstract class AbstractAmazonSnsProducer<E> implements AmazonSnsProducer<E> {
 
+  /** Class logger. */
+  private static final Logger LOGGER = LoggerFactory.getLogger(AbstractAmazonSnsProducer.class);
+
   /** The producer lifecycle state, initially {@link State#RUNNING}. */
   private final AtomicReference<State> state = new AtomicReference<>(State.RUNNING);
 
@@ -48,6 +59,8 @@ abstract class AbstractAmazonSnsProducer<E> implements AmazonSnsProducer<E> {
   /** The blocking queue for buffering requests before batch processing. */
   private final BlockingQueue<RequestEntry<E>> topicRequests;
 
+  private final ExecutorService callbackExecutor = Executors.newCachedThreadPool(ThreadFactoryProvider.getThreadFactory());
+
   /**
    * Sends a request entry by enqueuing it for batch processing.
    *
@@ -59,7 +72,7 @@ public ListenableFuture<ResponseSuccessEntry, ResponseFailEntry> send(final Requ
     if (State.RUNNING.equals(state.get())) {
       return enqueueRequest(requestEntry);
     } else {
-      final ListenableFutureImpl listenableFutureImpl = new ListenableFutureImpl();
+      final ListenableFutureImpl listenableFutureImpl = new ListenableFutureImpl(Runnable::run);
 
       listenableFutureImpl.fail(ResponseFailEntry.builder()
         .withCode("000")
@@ -78,8 +91,24 @@ public ListenableFuture<ResponseSuccessEntry, ResponseFailEntry> send(final Requ
    * accepted once shutdown.
    */
   @Override
-  public void shutdown() {
+  public void shutdown(final Runnable runnable) {
     state.compareAndSet(State.RUNNING, State.SHUTDOWN);
+
+    runnable.run();
+
+    try {
+      LOGGER.warn("Shutdown producer {}", getClass().getSimpleName());
+
+      callbackExecutor.shutdown();
+      if (!callbackExecutor.awaitTermination(60, TimeUnit.SECONDS)) {
+        LOGGER.warn("Producer executor service did not terminate in the specified time.");
+        final List<Runnable> droppedTasks = callbackExecutor.shutdownNow();
+        LOGGER.warn("Producer executor service was abruptly shut down. {} tasks will not be executed.", droppedTasks.size());
+      }
+    } catch (final InterruptedException ex) {
+      LOGGER.error(ex.getMessage(), ex);
+      Thread.currentThread().interrupt();
+    }
   }
 
   /**
@@ -92,7 +121,7 @@ public void shutdown() {
   @SneakyThrows
   private ListenableFuture<ResponseSuccessEntry, ResponseFailEntry> enqueueRequest(final RequestEntry<E> requestEntry) {
     try {
-      final ListenableFuture<ResponseSuccessEntry, ResponseFailEntry> trackPendingRequest = new ListenableFutureImpl();
+      final ListenableFuture<ResponseSuccessEntry, ResponseFailEntry> trackPendingRequest = new ListenableFutureImpl(callbackExecutor);
       pendingRequests.put(requestEntry.getId(), trackPendingRequest);
       topicRequests.put(requestEntry);
       return trackPendingRequest;

@@ -16,6 +16,7 @@
 
 package com.amazon.sns.messaging.lib.core;
 
+import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
@@ -72,8 +73,7 @@ public ListenableFuture<ResponseSuccessEntry, ResponseFailEntry> send(final Requ
    * Shuts down both the producer and consumer gracefully.
    */
   public void shutdown() {
-    amazonSnsProducer.shutdown();
-    amazonSnsConsumer.shutdown();
+    amazonSnsProducer.shutdown(amazonSnsConsumer::shutdown);
   }
 
   /**
@@ -85,6 +85,20 @@ public CompletableFuture<Void> await() {
     return amazonSnsConsumer.await();
   }
 
+  /**
+   * Returns a future that completes once all pending requests are drained and processed,
+   * bounded by the given timeout.
+   *
+   * @param timeout the maximum time to wait for all pending requests to be processed
+   * @return a {@link CompletableFuture} that completes when the consumer has finished, or
+   *         completes exceptionally with a {@link java.util.concurrent.TimeoutException}
+   *         if {@code timeout} elapses first
+   * @throws NullPointerException if {@code timeout} is {@code null}
+   */
+  public CompletableFuture<Void> await(final Duration timeout) {
+    return amazonSnsConsumer.await(timeout);
+  }
+
   /**
    * Creates an {@link AmazonSnsThreadPoolExecutor} configured for the given topic property.
    * For FIFO topics, a single-threaded pool is used to guarantee order.

@@ -16,6 +16,7 @@
 
 package com.amazon.sns.messaging.lib.core;
 
+import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -66,4 +67,18 @@ public interface AmazonSnsConsumer<R, O> {
    */
   public CompletableFuture<Void> await();
 
+  /**
+   * Returns a {@link CompletableFuture} that completes once all pending requests
+   * have been processed, bounded by the given timeout. If the timeout elapses
+   * before all pending requests are drained, the returned future completes
+   * exceptionally with a {@link java.util.concurrent.TimeoutException}.
+   *
+   * @param timeout the maximum time to wait for all pending requests to be
+   *                processed
+   * @return a future that completes when all requests are drained, or completes
+   *         exceptionally if {@code timeout} elapses first
+   * @throws NullPointerException if {@code timeout} is {@code null}
+   */
+  public CompletableFuture<Void> await(final Duration timeout);
+
 }
@@ -42,7 +42,7 @@ public interface AmazonSnsProducer<E> {
   /**
    * Shuts down the producer, preventing any further messages from being accepted.
    */
-  public void shutdown();
+  public void shutdown(final Runnable runnable);
 
 }
 // @formatter:on
@@ -18,6 +18,9 @@
 
 import static java.util.function.Function.identity;
 
+import java.time.Duration;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 
 import com.amazon.sns.messaging.lib.model.ResponseFailEntry;
@@ -64,5 +67,32 @@ default void addCallback(final Consumer<? super S> successCallback) {
    */
   void fail(final ResponseFailEntry entry);
 
+  /**
+   * Waits, if necessary, for this future to complete, then returns its success result.
+   * <p>
+   * Mirrors the contract of {@link java.util.concurrent.Future#get()}: if the future
+   * completed with a failure, this throws {@link ExecutionException} instead of returning.
+   * The exception's cause is the {@link Throwable} carried by the failure result, if any
+   * (see {@link ResponseFailEntry#getThrowable()}), or a new exception built from the
+   * failure result's message otherwise.
+   *
+   * @return the success result
+   * @throws InterruptedException if the current thread is interrupted while waiting
+   * @throws ExecutionException   if the future completed with a failure
+   */
+  S get() throws InterruptedException, ExecutionException;
+
+  /**
+   * Waits, if necessary, for at most the given timeout for this future to complete, then
+   * returns its success result. See {@link #get()} for failure semantics.
+   *
+   * @param timeout the maximum time to wait
+   * @return the success result
+   * @throws InterruptedException if the current thread is interrupted while waiting
+   * @throws ExecutionException   if the future completed with a failure
+   * @throws TimeoutException     if the timeout elapses before the future completes
+   */
+  S get(final Duration timeout) throws InterruptedException, ExecutionException, TimeoutException;
+
 }
 // @formatter:on