From 7cff78fb96827af3fd0dd29f0eb04821f6319294 Mon Sep 17 00:00:00 2001
From: david <mail@nonswag.dev>
Date: Wed, 21 Oct 2026 11:05:34 +0200
Subject: [PATCH] Fix featureflag fetching

In rare occasions fetch could complete before `whenComplete` was attached, causing the key to remove itself from `fetchesInProgress` before it was even added, throwing an IllegalStateException
---
 .../dev/faststats/SimpleFeatureFlagService.java     | 13 ++++++++++---
 .../test/java/dev/faststats/FeatureFlagTest.java    |  4 +---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/core/src/main/java/dev/faststats/SimpleFeatureFlagService.java b/core/src/main/java/dev/faststats/SimpleFeatureFlagService.java
index e5841057..03ff9fd8 100644
--- a/core/src/main/java/dev/faststats/SimpleFeatureFlagService.java
+++ b/core/src/main/java/dev/faststats/SimpleFeatureFlagService.java
@@ -45,7 +45,15 @@ final class SimpleFeatureFlagService extends SubmissionService implements Featur
 
     @SuppressWarnings("unchecked")
     <T> CompletableFuture<T> fetch(final SimpleFeatureFlag<T> flag) {
-        return (CompletableFuture<T>) fetchesInProgress.computeIfAbsent(flag.getId(), ignored -> createFetch(flag));
+        final var future = new CompletableFuture<T>();
+        final var existing = fetchesInProgress.putIfAbsent(flag.getId(), future);
+        if (existing != null) return (CompletableFuture<T>) existing;
+        createFetch(flag).whenComplete((value, throwable) -> {
+            fetchesInProgress.remove(flag.getId(), future);
+            if (throwable != null) future.completeExceptionally(throwable);
+            else future.complete(value);
+        });
+        return future;
     }
 
     <T> CompletableFuture<T> optIn(final SimpleFeatureFlag<T> flag) {
@@ -72,8 +80,7 @@ private <T> CompletableFuture<T> createFetch(final SimpleFeatureFlag<T> flag) {
         final var request = createRequest(CHECK_PATH, requestBody);
         logger.info("Fetching %s: %s", request.uri(), requestBody.toString());
         return HTTP_CLIENT.sendAsync(request, HttpResponse.BodyHandlers.ofString())
-                .thenApply(response -> handleFetchResponse(flag, response))
-                .whenComplete((ignored, throwable) -> fetchesInProgress.remove(flag.getId()));
+                .thenApply(response -> handleFetchResponse(flag, response));
     }
 
     private JsonObject createRequestBody(final SimpleFeatureFlag<?> flag) {
diff --git a/core/src/test/java/dev/faststats/FeatureFlagTest.java b/core/src/test/java/dev/faststats/FeatureFlagTest.java
index 9a6261df..4735b533 100644
--- a/core/src/test/java/dev/faststats/FeatureFlagTest.java
+++ b/core/src/test/java/dev/faststats/FeatureFlagTest.java
@@ -265,9 +265,7 @@ private void handle(final Socket client) {
                 }
 
                 final var body = new String(bodyChars, 0, read);
-                final var e = new Request(path, headers, body.isEmpty() ? new JsonObject() : JsonParser.parseString(body).getAsJsonObject());
-                System.out.println("parsed body: " + body + ", " + e.body);
-                requests.add(e);
+                requests.add(new Request(path, headers, body.isEmpty() ? new JsonObject() : JsonParser.parseString(body).getAsJsonObject()));
 
                 final var response = responses.poll(1, TimeUnit.SECONDS);
                 if (response == null) throw new AssertionError("No response enqueued");