Merge thread and concurrency utils in one class

Author: acogoluegnes

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

@@ -357,7 +357,7 @@ public void initChannel(SocketChannel ch) {
               });
       this.dispatchingExecutorService =
           Executors.newSingleThreadExecutor(
-              threadFactory("dispatching-" + clientConnectionName + "-"));
+              threadFactory("rabbitmq-stream-dispatching-" + clientConnectionName + "-"));
     } else {
       this.closeDispatchingExecutorService =
           Utils.makeIdempotent(

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

@@ -35,7 +35,7 @@ final class DynamicBatch<T> implements AutoCloseable {
   private final int configuredBatchSize, minBatchSize, maxBatchSize;
   private final Thread thread;
 
-  DynamicBatch(BatchConsumer<T> consumer, int batchSize, int maxUnconfirmed) {
+  DynamicBatch(BatchConsumer<T> consumer, int batchSize, int maxUnconfirmed, String id) {
     this.consumer = consumer;
     if (batchSize < maxUnconfirmed) {
       this.minBatchSize = min(MIN_BATCH_SIZE, batchSize / 2);
@@ -44,7 +44,7 @@ final class DynamicBatch<T> implements AutoCloseable {
     }
     this.configuredBatchSize = batchSize;
     this.maxBatchSize = batchSize * 2;
-    this.thread = ConcurrencyUtils.defaultThreadFactory().newThread(this::loop);
+    this.thread = ThreadUtils.newInternalThread(id, this::loop);
     this.thread.start();
   }
 

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

@@ -48,7 +48,8 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
       CompressionCodec compressionCodec,
       ByteBufAllocator byteBufAllocator,
       ObservationCollector<?> observationCollector,
-      StreamProducer producer) {
+      StreamProducer producer,
+      long producerId) {
     this.helper =
         new ProducerUtils.MessageAccumulatorHelper(
             codec,
@@ -61,6 +62,7 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
     this.producer = producer;
     this.observationCollector = (ObservationCollector<Object>) observationCollector;
     boolean shouldObserve = !this.observationCollector.isNoop();
+    String batchId = "rabbitmq-stream-dynamic-batch-producer-" + producerId;
     if (subEntrySize <= 1) {
       this.dynamicBatch =
           new DynamicBatch<>(
@@ -77,7 +79,8 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
                 return result;
               },
               batchSize,
-              maxUnconfirmedMessages);
+              maxUnconfirmedMessages,
+              batchId);
     } else {
       byte compressionCode =
           compressionCodec == null ? Compression.NONE.code() : compressionCodec.code();
@@ -127,7 +130,8 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
                 return result;
               },
               batchSize * subEntrySize,
-              maxUnconfirmedMessages);
+              maxUnconfirmedMessages,
+              batchId);
     }
   }
 

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

@@ -40,7 +40,8 @@ static MessageAccumulator createMessageAccumulator(
       Clock clock,
       String stream,
       ObservationCollector<?> observationCollector,
-      StreamProducer producer) {
+      StreamProducer producer,
+      long producerId) {
     if (dynamicBatch) {
       return new DynamicBatchMessageAccumulator(
           subEntrySize,
@@ -55,7 +56,8 @@ static MessageAccumulator createMessageAccumulator(
           compressionCodec,
           byteBufAllocator,
           observationCollector,
-          producer);
+          producer,
+          producerId);
     } else {
       if (subEntrySize <= 1) {
         return new SimpleMessageAccumulator(

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

@@ -190,7 +190,8 @@ public int fragmentLength(Object entity) {
             environment.clock(),
             stream,
             environment.observationCollector(),
-            this);
+            this,
+            this.id);
 
     boolean backgroundBatchPublishingTaskRequired =
         !dynamicBatch && batchPublishingDelay.toMillis() > 0;

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

@@ -17,11 +17,9 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Arrays;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.function.Function;
 import java.util.function.Predicate;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -31,48 +29,30 @@ final class ThreadUtils {
   private static final Logger LOGGER = LoggerFactory.getLogger(ThreadUtils.class);
 
   private static final ThreadFactory THREAD_FACTORY;
-  private static final Function<String, ExecutorService> EXECUTOR_SERVICE_FACTORY;
   private static final Predicate<Thread> IS_VIRTUAL;
 
   static {
+    ThreadFactory tf;
     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 {
+        Class<?> builderClass =
+            Arrays.stream(Thread.class.getDeclaredClasses())
+                .filter(c -> "Builder".equals(c.getSimpleName()))
+                .findFirst()
+                .get();
+        // Reflection code is the same as:
+        // Thread.ofVirtual().factory();
         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);
+        tf = (ThreadFactory) builderClass.getDeclaredMethod("factory").invoke(builder);
+      } catch (IllegalAccessException
+          | InvocationTargetException
+          | NoSuchMethodException
+          | RuntimeException e) {
+        LOGGER.debug("Error when creating virtual thread factory on Java 21+: {}", e.getMessage());
+        LOGGER.debug("Falling back to default thread factory");
+        tf = Executors.defaultThreadFactory();
       }
-      EXECUTOR_SERVICE_FACTORY =
-          prefix -> {
-            try {
-              // Reflection code is the same as the 2 following lines:
-              // ThreadFactory factory = Thread.ofVirtual().name(prefix, 0).factory();
-              // Executors.newThreadPerTaskExecutor(factory);
-              Object builder = Thread.class.getDeclaredMethod("ofVirtual").invoke(null);
-              if (prefix != null) {
-                builder =
-                    builderClass
-                        .getDeclaredMethod("name", String.class, Long.TYPE)
-                        .invoke(builder, prefix, 0L);
-              }
-              ThreadFactory factory =
-                  (ThreadFactory) builderClass.getDeclaredMethod("factory").invoke(builder);
-              return (ExecutorService)
-                  Executors.class
-                      .getDeclaredMethod("newThreadPerTaskExecutor", ThreadFactory.class)
-                      .invoke(null, factory);
-            } catch (IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
-              throw new RuntimeException(e);
-            }
-          };
       IS_VIRTUAL =
           thread -> {
             Method method = null;
@@ -85,14 +65,21 @@ final class ThreadUtils {
             }
           };
     } else {
-      THREAD_FACTORY = Executors.defaultThreadFactory();
-      EXECUTOR_SERVICE_FACTORY = prefix -> Executors.newCachedThreadPool(threadFactory(prefix));
+      tf = Executors.defaultThreadFactory();
       IS_VIRTUAL = ignored -> false;
     }
+    THREAD_FACTORY = tf;
   }
 
   private ThreadUtils() {}
 
+  /**
+   * Create a thread factory that prefixes thread names. Based on {@link
+   * java.util.concurrent.Executors#defaultThreadFactory()}, so always creates platform threads.
+   *
+   * @param prefix used to prefix thread names
+   * @return thread factory
+   */
   static ThreadFactory threadFactory(String prefix) {
     if (prefix == null) {
       return Executors.defaultThreadFactory();
@@ -101,8 +88,31 @@ static ThreadFactory threadFactory(String prefix) {
     }
   }
 
+  /**
+   * Returns a thread factory that creates virtual threads if available.
+   *
+   * @param prefix
+   * @return
+   */
   static ThreadFactory internalThreadFactory(String prefix) {
-    return new NamedThreadFactory(THREAD_FACTORY, prefix);
+    if (prefix == null) {
+      return THREAD_FACTORY;
+    } else {
+      return new NamedThreadFactory(THREAD_FACTORY, prefix);
+    }
+  }
+
+  /**
+   * Creates a virtual thread if available.
+   *
+   * @param name
+   * @param task
+   * @return
+   */
+  static Thread newInternalThread(String name, Runnable task) {
+    Thread t = THREAD_FACTORY.newThread(task);
+    t.setName(name);
+    return t;
   }
 
   static boolean isVirtual(Thread thread) {

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

@@ -17,6 +17,7 @@
 import static com.rabbitmq.stream.BackOffDelayPolicy.fixedWithInitialDelay;
 import static com.rabbitmq.stream.impl.Assertions.assertThat;
 import static com.rabbitmq.stream.impl.TestUtils.sync;
+import static com.rabbitmq.stream.impl.ThreadUtils.internalThreadFactory;
 import static java.time.Duration.*;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.*;
@@ -204,7 +205,6 @@ void completeExceptionally(ScheduledExecutorService scheduler) throws Exception
   static List<ScheduledExecutorService> schedulers() {
     return List.of(
         Executors.newSingleThreadScheduledExecutor(),
-        Executors.newScheduledThreadPool(
-            0, ThreadUtils.internalThreadFactory("async-retry-test-")));
+        Executors.newScheduledThreadPool(0, internalThreadFactory("async-retry-test-")));
   }
 }

src/test/java/com/rabbitmq/stream/impl/AsyncRetryTest.java

@@ -30,6 +30,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.IntStream;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
 
 public class DynamicBatchTest {
 
@@ -58,7 +59,7 @@ private static void printHistogram(Histogram histogram) {
   }
 
   @Test
-  void itemAreProcessed() {
+  void itemAreProcessed(TestInfo info) {
     MetricRegistry metrics = new MetricRegistry();
     Histogram batchSizeMetrics = metrics.histogram("batch-size");
     int itemCount = 3000;
@@ -71,7 +72,7 @@ void itemAreProcessed() {
           sync.down(items.size());
           return true;
         };
-    try (DynamicBatch<String> batch = new DynamicBatch<>(action, 100, 10_000)) {
+    try (DynamicBatch<String> batch = new DynamicBatch<>(action, 100, 10_000, batchId(info))) {
       RateLimiter rateLimiter = RateLimiter.create(10000);
       IntStream.range(0, itemCount)
           .forEach(
@@ -85,7 +86,7 @@ void itemAreProcessed() {
   }
 
   @Test
-  void failedProcessingIsReplayed() throws Exception {
+  void failedProcessingIsReplayed(TestInfo info) throws Exception {
     int itemCount = 10000;
     AtomicInteger collected = new AtomicInteger(0);
     AtomicInteger processed = new AtomicInteger(0);
@@ -102,7 +103,7 @@ void failedProcessingIsReplayed() throws Exception {
           }
           return result;
         };
-    try (DynamicBatch<String> batch = new DynamicBatch<>(action, 100, 10_000)) {
+    try (DynamicBatch<String> batch = new DynamicBatch<>(action, 100, 10_000, batchId(info))) {
       int firstRoundCount = itemCount / 5;
       IntStream.range(0, firstRoundCount)
           .forEach(
@@ -123,7 +124,7 @@ void failedProcessingIsReplayed() throws Exception {
   }
 
   @Test
-  void lowThrottlingValueShouldStillHighPublishingRate() throws Exception {
+  void lowThrottlingValueShouldStillHighPublishingRate(TestInfo info) throws Exception {
     int batchSize = 10;
     Semaphore semaphore = new Semaphore(batchSize);
     DynamicBatch.BatchConsumer<Long> action =
@@ -132,7 +133,7 @@ void lowThrottlingValueShouldStillHighPublishingRate() throws Exception {
           return true;
         };
 
-    try (DynamicBatch<Long> batch = new DynamicBatch<>(action, batchSize, 10_000)) {
+    try (DynamicBatch<Long> batch = new DynamicBatch<>(action, batchSize, 10_000, batchId(info))) {
       MetricRegistry metrics = new MetricRegistry();
       Meter rate = metrics.meter("publishing-rate");
       AtomicBoolean keepGoing = new AtomicBoolean(true);
@@ -154,4 +155,8 @@ void lowThrottlingValueShouldStillHighPublishingRate() throws Exception {
               System.nanoTime() - start > TimeUnit.SECONDS.toNanos(1) && rate.getMeanRate() > 1000);
     }
   }
+
+  private static String batchId(TestInfo info) {
+    return info.getTestMethod().get().getName();
+  }
 }