Detect pattern sources, use OptionalInt result, inline awaitAll

- Calculator returns OptionalInt; pre-pass on Source.topicPattern() returns
  empty when any sub-topology subscribes via regex (was silently undercount).
- Resolver also returns OptionalInt; non-positive replicas and AdminClient
  failures map to empty rather than each falling back independently.
- Framework collapses every "cannot resolve" path through one .orElse so
  FALLBACK_NUM_STREAM_THREADS is the only place the literal substitution
  happens.
- describePartitions inlines the awaitAll helper into a single deadline-
  driven loop. Behaviour identical (allOf + per-future get and now+timeout
  + per-future get(remaining) both share one 5s budget; AdminClient fires
  RPCs concurrently before returning futures), readability improved.
- New test covers regex-source topology -> empty all the way through.

Author: suresh-prakash

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/KafkaStreamsApp.java

@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -292,10 +293,11 @@ protected Optional<StreamThreadsCountResolver> getStreamThreadsCountResolver() {
     return Optional.empty();
   }
 
-  // Caller must check StreamThreadsCountResolver.isDynamic(...) before invoking. Always returns
-  // a concrete value: either the resolver's computed count or FALLBACK_NUM_STREAM_THREADS.
-  // The string sentinel "DYNAMIC" must never leak back to Kafka Streams config — it can't be
-  // parsed as an integer.
+  // Caller must check StreamThreadsCountResolver.isDynamic(...) before invoking. Always returns a
+  // concrete int — every "cannot resolve" path (resolver-supplier throws, no resolver wired,
+  // resolver returns empty for unsupported topologies / non-positive replicas / AdminClient
+  // failure) collapses to FALLBACK_NUM_STREAM_THREADS so the literal "DYNAMIC" sentinel never
+  // reaches Kafka Streams.
   private int resolveDynamicStreamThreads(Map<String, Object> streamsProperties) {
     Optional<StreamThreadsCountResolver> resolver;
     try {
@@ -316,7 +318,14 @@ private int resolveDynamicStreamThreads(Map<String, Object> streamsProperties) {
               StreamThreadsCountResolver.FALLBACK_NUM_STREAM_THREADS);
       return StreamThreadsCountResolver.FALLBACK_NUM_STREAM_THREADS;
     }
-    return resolver.get().resolve(this.topology, streamsProperties);
+    final OptionalInt resolved = resolver.get().resolve(this.topology, streamsProperties);
+    if (resolved.isEmpty()) {
+      getLogger()
+          .warn(
+              "Resolver could not compute dynamic num.stream.threads; falling back to {}",
+              StreamThreadsCountResolver.FALLBACK_NUM_STREAM_THREADS);
+    }
+    return resolved.orElse(StreamThreadsCountResolver.FALLBACK_NUM_STREAM_THREADS);
   }
 
   public Map<String, Object> getStreamsConfig(Config jobConfig) {

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/threading/DynamicStreamThreadsCountCalculator.java

@@ -6,6 +6,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.OptionalInt;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -30,8 +31,10 @@
  * stream tasks. Summing across sub-topologies and dividing by the replica count yields the threads
  * each instance should run to keep all tasks active without idle threads.
  *
- * <p>Regex/pattern subscriptions ({@link Source#topicPattern()}) are not supported — those
- * sub-topologies contribute zero tasks. No current app subscribes via regex.
+ * <p>Returns {@link OptionalInt#empty()} when the topology contains a regex/pattern subscription
+ * ({@link Source#topicPattern()}) — those sub-topologies cannot be enumerated against the broker
+ * up-front, so dynamic sizing would silently under-count tasks. The caller falls back to its
+ * configured default in that case.
  */
 public class DynamicStreamThreadsCountCalculator {
 
@@ -47,12 +50,33 @@ private static Set<String> sourceTopicsOf(final Subtopology subtopology) {
         .collect(toUnmodifiableSet());
   }
 
-  public int compute(final Topology topology, final AdminClient adminClient, final int replicas) {
+  private static boolean hasPatternSource(final Subtopology subtopology) {
+    return subtopology.nodes().stream()
+        .filter(node -> node instanceof Source)
+        .map(node -> (Source) node)
+        .anyMatch(source -> source.topicPattern() != null);
+  }
+
+  public OptionalInt compute(
+      final Topology topology, final AdminClient adminClient, final int replicas) {
     if (replicas <= 0) {
       throw new IllegalArgumentException("replicas must be positive, got " + replicas);
     }
 
     final TopologyDescription description = topology.describe();
+
+    // Bail out if any sub-topology subscribes via regex — topicSet() is empty for those, so
+    // dynamic sizing would silently under-count tasks. The caller substitutes its fallback.
+    final boolean anyPatternSource =
+        description.subtopologies().stream()
+            .anyMatch(DynamicStreamThreadsCountCalculator::hasPatternSource);
+    if (anyPatternSource) {
+      logger.warn(
+          "Topology contains a regex/pattern source; dynamic num.stream.threads is not supported. "
+              + "Caller will fall back to its configured default.");
+      return OptionalInt.empty();
+    }
+
     final Set<String> sourceTopics =
         description.subtopologies().stream()
             .flatMap(subtopology -> sourceTopicsOf(subtopology).stream())
@@ -87,28 +111,39 @@ public int compute(final Topology topology, final AdminClient adminClient, final
         subtopologyCount,
         replicas,
         threads);
-    return threads;
+    return OptionalInt.of(threads);
   }
 
+  // Single-loop implementation: AdminClient.describeTopics() already fires all RPCs concurrently
+  // before returning futures, so iteration here only consumes a shared deadline (now+timeout) —
+  // total wall-clock is capped at DESCRIBE_TOPICS_TIMEOUT_MILLIS regardless of topic count.
   private Map<String, Integer> describePartitions(
       final AdminClient adminClient, final Set<String> topics) {
     if (topics.isEmpty()) {
       return Map.of();
     }
     final DescribeTopicsResult result = adminClient.describeTopics(topics);
     final Map<String, KafkaFuture<TopicDescription>> futures = result.topicNameValues();
-
-    // Single shared deadline across all topics — worst-case wait is one timeout, not N.
-    // AdminClient retries the underlying RPC internally per its own retry config.
-    awaitAll(futures.values(), DESCRIBE_TOPICS_TIMEOUT_MILLIS);
-
+    final long deadlineMillis = System.currentTimeMillis() + DESCRIBE_TOPICS_TIMEOUT_MILLIS;
     final Map<String, Integer> partitions = new HashMap<>();
+
     for (final Entry<String, KafkaFuture<TopicDescription>> entry : futures.entrySet()) {
+      final long remainingMillis = deadlineMillis - System.currentTimeMillis();
+      if (remainingMillis <= 0) {
+        throw new RuntimeException(
+            "Timed out describing topics after " + DESCRIBE_TOPICS_TIMEOUT_MILLIS + "ms");
+      }
       try {
-        // Non-blocking after awaitAll: each future is already complete (or completed
-        // exceptionally). Per-topic handling stays so UnknownTopicOrPartitionException
-        // doesn't fail the whole batch.
-        partitions.put(entry.getKey(), entry.getValue().get().partitions().size());
+        partitions.put(
+            entry.getKey(),
+            entry.getValue().get(remainingMillis, TimeUnit.MILLISECONDS).partitions().size());
+      } catch (final TimeoutException timeoutException) {
+        throw new RuntimeException(
+            "Timed out describing topic " + entry.getKey(), timeoutException);
+      } catch (final InterruptedException interruptedException) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException(
+            "Interrupted while describing topic " + entry.getKey(), interruptedException);
       } catch (final ExecutionException executionException) {
         if (executionException.getCause() instanceof UnknownTopicOrPartitionException) {
           logger.warn(
@@ -119,29 +154,8 @@ private Map<String, Integer> describePartitions(
           throw new RuntimeException(
               "Failed to describe topic " + entry.getKey(), executionException);
         }
-      } catch (final InterruptedException interruptedException) {
-        Thread.currentThread().interrupt();
-        throw new RuntimeException(
-            "Interrupted while describing topic " + entry.getKey(), interruptedException);
       }
     }
     return Map.copyOf(partitions);
   }
-
-  private static void awaitAll(
-      final java.util.Collection<KafkaFuture<TopicDescription>> futures, final long timeoutMillis) {
-    try {
-      KafkaFuture.allOf(futures.toArray(KafkaFuture[]::new))
-          .get(timeoutMillis, TimeUnit.MILLISECONDS);
-    } catch (final TimeoutException timeoutException) {
-      throw new RuntimeException(
-          "Timed out describing topics after " + timeoutMillis + "ms", timeoutException);
-    } catch (final InterruptedException interruptedException) {
-      Thread.currentThread().interrupt();
-      throw new RuntimeException("Interrupted describing topics", interruptedException);
-    } catch (final ExecutionException executionException) {
-      // At least one topic future completed exceptionally; per-future handling below classifies
-      // (UnknownTopicOrPartitionException → 0 partitions; everything else → fail).
-    }
-  }
 }

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/threading/StreamThreadsCountResolver.java

@@ -2,6 +2,7 @@
 
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.Properties;
 import java.util.function.Function;
 import java.util.function.IntSupplier;
@@ -58,28 +59,25 @@ public static boolean isDynamic(final Map<String, Object> streamsProperties) {
   }
 
   /**
-   * Resolve a concrete thread count for the given topology. Returns {@link
-   * #FALLBACK_NUM_STREAM_THREADS} when prerequisites are missing (non-positive replica count) or
-   * the AdminClient/calculator call fails, so the application can still start.
+   * Resolve a thread count for the given topology. Returns {@link OptionalInt#empty()} when
+   * prerequisites are missing (non-positive replica count), the topology is unsupported (regex
+   * sources), or the AdminClient/calculator call fails — the caller substitutes its fallback so the
+   * application can still start.
    */
-  public int resolve(final Topology topology, final Map<String, Object> streamsProperties) {
+  public OptionalInt resolve(final Topology topology, final Map<String, Object> streamsProperties) {
     final int replicas = replicaCountSupplier.getAsInt();
     if (replicas <= 0) {
-      logger.warn(
-          "replica.count is non-positive ({}); falling back to {} stream threads",
-          replicas,
-          FALLBACK_NUM_STREAM_THREADS);
-      return FALLBACK_NUM_STREAM_THREADS;
+      logger.warn("replica.count is non-positive ({}); skipping dynamic resolution", replicas);
+      return OptionalInt.empty();
     }
     try (final AdminClient adminClient =
         adminClientFactory.apply(toProperties(streamsProperties))) {
       return calculator.compute(topology, adminClient, replicas);
     } catch (final RuntimeException runtimeException) {
       logger.error(
-          "Failed to compute dynamic num.stream.threads; falling back to {}",
-          FALLBACK_NUM_STREAM_THREADS,
+          "Failed to compute dynamic num.stream.threads; skipping dynamic resolution",
           runtimeException);
-      return FALLBACK_NUM_STREAM_THREADS;
+      return OptionalInt.empty();
     }
   }
 

kafka-streams-framework/src/test/java/org/hypertrace/core/kafkastreams/framework/SampleAppTest.java

@@ -16,11 +16,15 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.regex.Pattern;
 import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.streams.StreamsBuilder;
 import org.apache.kafka.streams.TestInputTopic;
 import org.apache.kafka.streams.TestOutputTopic;
 import org.apache.kafka.streams.TopologyTestDriver;
 import org.apache.kafka.streams.errors.LogAndContinueExceptionHandler;
+import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.kstream.KStream;
 import org.hypertrace.core.kafkastreams.framework.rocksdb.BoundedMemoryConfigSetter;
 import org.hypertrace.core.kafkastreams.framework.threading.StreamThreadsCountResolver;
 import org.hypertrace.core.serviceframework.config.ConfigClientFactory;
@@ -111,6 +115,44 @@ public Map<String, Object> getStreamsConfig(Config jobConfig) {
         is(StreamThreadsCountResolver.FALLBACK_NUM_STREAM_THREADS));
   }
 
+  // Pattern-source topology: calculator returns OptionalInt.empty() because regex subscriptions
+  // can't be enumerated up-front against the broker. The framework must substitute the integer
+  // fallback so Kafka Streams gets a parseable value (not the literal "DYNAMIC").
+  @Test
+  public void dynamicWithPatternSourceFallsBackToEight() {
+    SampleApp dynamicApp =
+        new SampleApp(ConfigClientFactory.getClient()) {
+          @Override
+          public Map<String, Object> getStreamsConfig(Config jobConfig) {
+            Map<String, Object> properties = super.getStreamsConfig(jobConfig);
+            properties.put(NUM_STREAM_THREADS_CONFIG, StreamThreadsCountResolver.DYNAMIC_SENTINEL);
+            return properties;
+          }
+
+          @Override
+          public StreamsBuilder buildTopology(
+              Map<String, Object> streamsConfig,
+              StreamsBuilder streamsBuilder,
+              Map<String, KStream<?, ?>> sourceStreams) {
+            streamsBuilder.stream(
+                    Pattern.compile("input-.*"), Consumed.with(Serdes.String(), Serdes.String()))
+                .foreach((key, value) -> {});
+            return streamsBuilder;
+          }
+
+          @Override
+          protected Optional<StreamThreadsCountResolver> getStreamThreadsCountResolver() {
+            return Optional.of(new StreamThreadsCountResolver(() -> 8));
+          }
+        };
+
+    dynamicApp.doInit();
+
+    assertThat(
+        dynamicApp.streamsConfig.get(NUM_STREAM_THREADS_CONFIG),
+        is(StreamThreadsCountResolver.FALLBACK_NUM_STREAM_THREADS));
+  }
+
   // Verifies the fix for "exception while obtaining the resolver should not leak DYNAMIC" — when
   // getStreamThreadsCountResolver() throws, the framework must catch and fall back to the integer
   // default rather than letting the literal sentinel reach Kafka Streams.

kafka-streams-framework/src/test/java/org/hypertrace/core/kafkastreams/framework/threading/DynamicStreamThreadsCountCalculatorTest.java

@@ -10,7 +10,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
 import java.util.Set;
+import java.util.regex.Pattern;
 import java.util.stream.IntStream;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.DescribeTopicsResult;
@@ -44,23 +46,23 @@ void singleSubtopologyWithThirtyPartitionsAndEightReplicasGivesFourThreads() {
     final Topology topology = topologyForSubtopologies(Set.of("topic-a"));
     stubPartitions(Map.of("topic-a", 30));
 
-    assertEquals(4, calculator.compute(topology, adminClient, 8));
+    assertEquals(OptionalInt.of(4), calculator.compute(topology, adminClient, 8));
   }
 
   @Test
   void twoSubtopologiesAreSummedThenDividedByReplicas() {
     final Topology topology = topologyForSubtopologies(Set.of("topic-a"), Set.of("topic-b"));
     stubPartitions(Map.of("topic-a", 10, "topic-b", 20));
 
-    assertEquals(4, calculator.compute(topology, adminClient, 8));
+    assertEquals(OptionalInt.of(4), calculator.compute(topology, adminClient, 8));
   }
 
   @Test
   void absentTopicCountsAsZeroPartitions() {
     final Topology topology = topologyForSubtopologies(Set.of("topic-a"), Set.of("topic-missing"));
     stubPartitions(Map.of("topic-a", 16, "topic-missing", ABSENT));
 
-    assertEquals(2, calculator.compute(topology, adminClient, 8));
+    assertEquals(OptionalInt.of(2), calculator.compute(topology, adminClient, 8));
   }
 
   @Test
@@ -78,7 +80,20 @@ void totalTasksZeroReturnsOne() {
     final Topology topology = topologyForSubtopologies(Set.of("topic-a"));
     stubPartitions(Map.of("topic-a", ABSENT));
 
-    assertEquals(1, calculator.compute(topology, adminClient, 8));
+    assertEquals(OptionalInt.of(1), calculator.compute(topology, adminClient, 8));
+  }
+
+  // Regex/pattern subscriptions cannot be enumerated up-front against the broker, so dynamic
+  // sizing would silently under-count tasks. Calculator must signal "not applicable" via empty
+  // so the caller falls back to its configured default.
+  @Test
+  void patternSubscriptionReturnsEmpty() {
+    final StreamsBuilder builder = new StreamsBuilder();
+    builder.stream(Pattern.compile("topic-.*"), Consumed.with(Serdes.String(), Serdes.String()))
+        .foreach((key, value) -> {});
+    final Topology topology = builder.build();
+
+    assertEquals(OptionalInt.empty(), calculator.compute(topology, adminClient, 8));
   }
 
   @SafeVarargs