Fix subscription worker starvation under a small thread pool

Signed-off-by: Marvin Froeder <velo.br@gmail.com>

Author: velo

graphql/src/main/java/feign/graphql/GraphqlCapability.java

@@ -28,9 +28,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.concurrent.Executor;
-import java.util.concurrent.SynchronousQueue;
-import java.util.concurrent.ThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicLong;
 
 @Experimental
@@ -90,18 +88,14 @@ public GraphqlCapability(
   }
 
   /**
-   * Bounded and daemon, so a runaway subscription count fails fast rather than exhausting threads.
-   * A synchronous handoff queue is deliberate: subscription workers are long-lived, so queueing
-   * them would hide exhaustion until the heap gave out.
+   * Each open {@code Flow.Publisher} or {@code CompletableFuture} subscription holds one worker for
+   * its lifetime, so the default pool grows on demand and reaps idle threads rather than capping
+   * concurrent subscriptions at a guess. Supply a bounded executor to cap them deliberately: the
+   * excess is refused with {@code RejectedExecutionException} rather than left hanging.
    */
   private static Executor defaultExecutor() {
     var threads = new AtomicLong();
-    return new ThreadPoolExecutor(
-        0,
-        Math.max(8, Runtime.getRuntime().availableProcessors() * 4),
-        60L,
-        TimeUnit.SECONDS,
-        new SynchronousQueue<>(),
+    return Executors.newCachedThreadPool(
         runnable -> {
           var thread =
               new Thread(runnable, "feign-graphql-subscription-" + threads.incrementAndGet());

graphql/src/main/java/feign/graphql/GraphqlDecoder.java

@@ -34,6 +34,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Flow;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.SubmissionPublisher;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Stream;
@@ -210,14 +211,19 @@ private CompletableFuture<Object> futureOf(Subscription subscription, Type eleme
             subscription.unsubscribe();
           }
         });
-    executor.execute(
-        () -> {
-          try {
-            future.complete(first(subscription, elementType, 0).orElse(null));
-          } catch (Throwable e) {
-            future.completeExceptionally(e);
-          }
-        });
+    try {
+      executor.execute(
+          () -> {
+            try {
+              future.complete(first(subscription, elementType, 0).orElse(null));
+            } catch (Throwable e) {
+              future.completeExceptionally(e);
+            }
+          });
+    } catch (RejectedExecutionException e) {
+      subscription.unsubscribe();
+      future.completeExceptionally(e);
+    }
     return future;
   }
 
@@ -235,7 +241,9 @@ private Stream<Object> elements(Subscription subscription, Type elementType, lon
   }
 
   private Flow.Publisher<Object> publish(Subscription subscription, Type elementType) {
-    var publisher = new SubmissionPublisher<>(executor, Flow.defaultBufferSize());
+    // Runnable::run delivers on the pump thread: one worker per subscription in total, delivery can
+    // never be rejected by a busy pool, and onNext is inherently ordered.
+    var publisher = new SubmissionPublisher<>(Runnable::run, Flow.defaultBufferSize());
     var started = new AtomicBoolean();
     // Pumping starts on the first subscribe, so hasSubscribers() is meaningful from the first
     // element onwards and there is no pre-subscribe window to latch around.
@@ -244,18 +252,24 @@ private Flow.Publisher<Object> publish(Subscription subscription, Type elementTy
       if (!started.compareAndSet(false, true)) {
         return;
       }
-      executor.execute(
-          () -> {
-            try (var elements = elements(subscription, elementType, 0)) {
-              var iterator = elements.iterator();
-              while (iterator.hasNext() && publisher.hasSubscribers()) {
-                publisher.submit(iterator.next());
+      try {
+        executor.execute(
+            () -> {
+              try (var elements = elements(subscription, elementType, 0)) {
+                var iterator = elements.iterator();
+                while (iterator.hasNext() && publisher.hasSubscribers()) {
+                  publisher.submit(iterator.next());
+                }
+                publisher.close();
+              } catch (Throwable e) {
+                publisher.closeExceptionally(e);
               }
-              publisher.close();
-            } catch (Throwable e) {
-              publisher.closeExceptionally(e);
-            }
-          });
+            });
+      } catch (RejectedExecutionException e) {
+        // A subscriber must always get a terminal signal; stranding it is worse than failing it.
+        subscription.unsubscribe();
+        publisher.closeExceptionally(e);
+      }
     };
   }
 

graphql/src/test/java/feign/graphql/GraphqlSubscriptionConcurrencyTest.java

@@ -24,10 +24,14 @@
 import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Flow;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.IntStream;
@@ -62,6 +66,11 @@ class GraphqlSubscriptionConcurrencyTest {
 
   private final AtomicInteger openSockets = new AtomicInteger();
 
+  /**
+   * When false the server acknowledges the subscribe and then stays quiet, as a real feed would.
+   */
+  private volatile boolean emitEvents = true;
+
   public static class Price {
     public String symbol;
     public double price;
@@ -78,6 +87,11 @@ interface StockApi {
         "subscription onPrice($symbol: String!) {"
             + " priceChanged(symbol: $symbol) { symbol price } }")
     Flow.Publisher<Price> publishPrice(String symbol);
+
+    @GraphqlQuery(
+        "subscription onPrice($symbol: String!) {"
+            + " priceChanged(symbol: $symbol) { symbol price } }")
+    CompletableFuture<Price> futurePrice(String symbol);
   }
 
   @BeforeEach
@@ -120,6 +134,9 @@ public void onMessage(WebSocket webSocket, String text) {
         if (!"subscribe".equals(type)) {
           return;
         }
+        if (!emitEvents) {
+          return;
+        }
         var id = message.get("id").asText();
         var symbol = message.get("payload").get("variables").get("symbol").asText();
         for (var i = 0; i < EVENTS_EACH; i++) {
@@ -156,6 +173,42 @@ private StockApi buildClient() {
         .target(StockApi.class, server.url("/graphql").toString());
   }
 
+  private StockApi buildClient(Executor executor) {
+    return Feign.builder()
+        .addCapability(
+            new GraphqlCapability(new JacksonCodec(mapper), Duration.ofSeconds(30), executor))
+        .target(StockApi.class, server.url("/graphql").toString());
+  }
+
+  private int drain(Flow.Publisher<Price> publisher) throws Exception {
+    var delivered = new AtomicInteger();
+    var done = new CountDownLatch(1);
+    publisher.subscribe(
+        new Flow.Subscriber<Price>() {
+          @Override
+          public void onSubscribe(Flow.Subscription subscription) {
+            subscription.request(Long.MAX_VALUE);
+          }
+
+          @Override
+          public void onNext(Price item) {
+            delivered.incrementAndGet();
+          }
+
+          @Override
+          public void onError(Throwable throwable) {
+            done.countDown();
+          }
+
+          @Override
+          public void onComplete() {
+            done.countDown();
+          }
+        });
+    assertThat(done.await(60, TimeUnit.SECONDS)).isTrue();
+    return delivered.get();
+  }
+
   private <T> List<T> runAllAtOnce(List<Callable<T>> tasks) throws Exception {
     var pool = Executors.newFixedThreadPool(tasks.size());
     try {
@@ -257,6 +310,62 @@ public void onComplete() {
     assertThat(closed.await(30, TimeUnit.SECONDS)).isTrue();
   }
 
+  @Test
+  void aPoolSizedForTheSubscriptionsIsEnough() throws Exception {
+    // One worker per open subscription is the documented cost. Needing a second thread per
+    // subscription for delivery would starve this pool and hang instead.
+    var pool = Executors.newFixedThreadPool(SUBSCRIPTIONS);
+    try {
+      var api = buildClient(pool);
+      List<Callable<Integer>> tasks =
+          IntStream.range(0, SUBSCRIPTIONS)
+              .<Callable<Integer>>mapToObj(index -> () -> drain(api.publishPrice("SYM" + index)))
+              .toList();
+
+      assertThat(runAllAtOnce(tasks)).containsOnly(EVENTS_EACH);
+    } finally {
+      pool.shutdownNow();
+    }
+  }
+
+  @Test
+  void aPoolTooSmallRefusesRatherThanStranding() throws Exception {
+    var pool = new ThreadPoolExecutor(0, 4, 60L, TimeUnit.SECONDS, new SynchronousQueue<>());
+    try {
+      var api = buildClient(pool);
+      List<Callable<Integer>> tasks =
+          IntStream.range(0, SUBSCRIPTIONS)
+              .<Callable<Integer>>mapToObj(index -> () -> drain(api.publishPrice("SYM" + index)))
+              .toList();
+
+      // Far more subscriptions than workers. Some are refused, but every subscriber must reach a
+      // terminal signal — drain() asserts that. Leaving one waiting forever is the failure mode.
+      assertThat(runAllAtOnce(tasks)).hasSize(SUBSCRIPTIONS);
+    } finally {
+      pool.shutdownNow();
+    }
+  }
+
+  @Test
+  void longLivedSubscriptionsAreNotCappedByCoreCount() throws Exception {
+    emitEvents = false;
+    var api = buildClient();
+
+    // Each of these holds its worker parked on the queue for as long as it is open, which is what a
+    // real feed does. A pool sized from the core count would refuse the excess synchronously.
+    var futures =
+        IntStream.range(0, SUBSCRIPTIONS)
+            .mapToObj(index -> api.futurePrice("SYM" + index))
+            .toList();
+
+    assertThat(futures)
+        .as("no subscription should have been refused a worker")
+        .allSatisfy(future -> assertThat(future).isNotCompleted());
+
+    futures.forEach(future -> future.cancel(true));
+    assertThat(closed.await(30, TimeUnit.SECONDS)).isTrue();
+  }
+
   @Test
   void closingMidStreamFromAnotherThreadTerminatesPromptly() throws Exception {
     var api = buildClient();