Test the at-least-once path of the dedup state machine

easyssf-receiver-jdbc/AGENTS.md

@@ -17,8 +17,10 @@ The JDBC stores of the receiver in plain SQL, over the small `SsfJdbcOperations`
   schema changes. All share `AbstractJdbcSsfStoresTests`; add store tests there.
   `AbstractJdbcSharedReceiversTests` (same three databases) runs two receivers against one
   `TestTransmitter` on shared stores: a SET handled by one instance and skipped by the others,
-  acknowledgements handed over between instances, and the race where one instance acknowledges a SET another is still handling (covered
-  by the in-progress state of the dedup store, see `scratch/plans/dedup-in-progress-state.md`).
+  acknowledgements handed over between instances, the race where one instance acknowledges a SET
+  another is still handling, the at-least-once path (an instance dies after its handler and before
+  `processed`, the lease hands the SET to the next one), an abandoned claim, and concurrent claims
+  yielding exactly one `NEW`. Keep it the place for multi-instance scenarios.
 - The module is an explicit Java module; test compilation adds `java.naming` because H2's data
   source implements a `javax.naming` interface (see `pom.xml`). Keep that if you touch the
   surefire or compiler configuration.

easyssf-receiver-jdbc/src/test/java/org/easyssf/receiver/jdbc/AbstractJdbcSharedReceiversTests.java

@@ -2,21 +2,31 @@
 
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
+import java.time.Clock;
 import java.time.Duration;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
 import java.util.ArrayList;
 import java.util.Base64;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import org.easyssf.core.event.SsfEventToken;
 import org.easyssf.core.event.SsfSubjectIdentifiers;
 import org.easyssf.receiver.event.SsfEventHandler;
 import org.easyssf.receiver.http.JdkSsfHttpClient;
 import org.easyssf.receiver.http.SsfHttpClient;
 import org.easyssf.receiver.poll.SsfPoller;
 import org.easyssf.receiver.set.NimbusSsfSetVerifier;
+import org.easyssf.receiver.set.SsfJtiDedupStore;
 import org.easyssf.receiver.set.SsfSetProcessor;
 import org.easyssf.test.TestTransmitter;
 import org.junit.jupiter.api.AfterEach;
@@ -41,6 +51,10 @@ abstract class AbstractJdbcSharedReceiversTests {
 
     private static final String PREFIX = JdbcSsfSchema.DEFAULT_TABLE_PREFIX;
 
+    private SsfJdbcOperations jdbc;
+
+    private final MutableClock clock = new MutableClock();
+
     private JdbcSsfJtiDedupStore dedupStore;
 
     private JdbcSsfPollAckStore ackStore;
@@ -53,13 +67,13 @@ abstract class AbstractJdbcSharedReceiversTests {
     @BeforeEach
     void setUp() {
         this.transmitter = new TestTransmitter();
-        SsfJdbcOperations jdbc = emptyDatabase();
-        JdbcSsfSchema.prepareTable(jdbc, JdbcSsfSchema.processedSetTable(PREFIX),
+        this.jdbc = emptyDatabase();
+        JdbcSsfSchema.prepareTable(this.jdbc, JdbcSsfSchema.processedSetTable(PREFIX),
                 JdbcSsfSchema.createProcessedSetTable(PREFIX), true, null);
-        JdbcSsfSchema.prepareTable(jdbc, JdbcSsfSchema.pollAckTable(PREFIX), JdbcSsfSchema.createPollAckTable(PREFIX),
-                true, null);
-        this.dedupStore = new JdbcSsfJtiDedupStore(jdbc, PREFIX);
-        this.ackStore = new JdbcSsfPollAckStore(jdbc, PREFIX);
+        JdbcSsfSchema.prepareTable(this.jdbc, JdbcSsfSchema.pollAckTable(PREFIX),
+                JdbcSsfSchema.createPollAckTable(PREFIX), true, null);
+        this.dedupStore = dedupStore();
+        this.ackStore = new JdbcSsfPollAckStore(this.jdbc, PREFIX);
     }
 
     @AfterEach
@@ -177,15 +191,138 @@ void setBeingHandledByOneInstanceIsNeitherHandledNorAcknowledgedByAnother() thro
         assertThat(transmitter.acknowledgedSets()).containsExactly(jti);
     }
 
+    @Test
+    void setHandledByAnInstanceThatDiedBeforeRecordingTheSuccessIsHandledAgainAfterTheLease() {
+        // a's instance dies between its handler's side effects and the record of the
+        // success
+        Receiver a = receiver(new SsfJtiDedupStore() {
+            @Override
+            public Claim claim(SsfEventToken eventToken) {
+                return AbstractJdbcSharedReceiversTests.this.dedupStore.claim(eventToken);
+            }
+
+            @Override
+            public void processed(SsfEventToken eventToken) {
+                throw new IllegalStateException("the instance died");
+            }
+
+            @Override
+            public void forget(SsfEventToken eventToken) {
+                AbstractJdbcSharedReceiversTests.this.dedupStore.forget(eventToken);
+            }
+        });
+        Receiver b = receiver(dedupStore());
+        String jti = queue("session-1");
+
+        assertThatIllegalStateException().isThrownBy(a.poller::pollNow).withMessageContaining("died");
+        assertThat(a.handled).containsExactly(jti);
+        assertThat(this.transmitter.acknowledgedSets()).isEmpty();
+        assertThat(state(jti)).isEqualTo("IN_PROGRESS");
+
+        // within the lease: b leaves the SET alone
+        assertThat(b.poller.pollNow()).isEqualTo(1);
+        assertThat(b.handled).isEmpty();
+        assertThat(this.transmitter.acknowledgedSets()).isEmpty();
+
+        // after the lease: b takes the claim over and handles the SET a second time, at
+        // least once
+        this.clock.advance(Duration.ofSeconds(61));
+        assertThat(b.poller.pollNow()).isEqualTo(1);
+        assertThat(b.handled).containsExactly(jti);
+        assertThat(this.transmitter.acknowledgedSets()).containsExactly(jti);
+        assertThat(state(jti)).isEqualTo("PROCESSED");
+        // handled twice in total, with the same idempotency key both times
+        assertThat(a.keys).containsExactlyElementsOf(b.keys);
+    }
+
+    @Test
+    void setClaimedByAnInstanceThatDiedBeforeItsHandlerRanIsHandledAfterTheLease() {
+        Receiver b = receiver(dedupStore());
+        String jti = queue("session-1");
+        // the dead instance's claim, left behind in the shared table
+        assertThat(this.dedupStore.claim(token(jti))).isEqualTo(SsfJtiDedupStore.Claim.NEW);
+
+        assertThat(b.poller.pollNow()).isEqualTo(1);
+        assertThat(b.handled).isEmpty();
+        assertThat(this.transmitter.acknowledgedSets()).isEmpty();
+
+        this.clock.advance(Duration.ofSeconds(61));
+        assertThat(b.poller.pollNow()).isEqualTo(1);
+        assertThat(b.handled).containsExactly(jti);
+        assertThat(this.transmitter.acknowledgedSets()).containsExactly(jti);
+    }
+
+    @Test
+    void concurrentClaimsOfOneSetYieldExactlyOneNew() throws Exception {
+        String jti = "jti-concurrent";
+        assertThat(claimConcurrently(jti, 8)).containsExactly(SsfJtiDedupStore.Claim.NEW);
+        assertThat(state(jti)).isEqualTo("IN_PROGRESS");
+
+        // the lease expired: exactly one of the concurrent take-overs wins
+        this.clock.advance(Duration.ofSeconds(61));
+        assertThat(claimConcurrently(jti, 8)).containsExactly(SsfJtiDedupStore.Claim.NEW);
+
+        this.dedupStore.processed(token(jti));
+        assertThat(claimConcurrently(jti, 8)).isEmpty();
+        assertThat(state(jti)).isEqualTo("PROCESSED");
+    }
+
+    /**
+     * Claims the SET from several threads at the same moment, each with a store of its
+     * own.
+     * @return the claims that came back NEW
+     */
+    private List<SsfJtiDedupStore.Claim> claimConcurrently(String jti, int threads) throws Exception {
+        CountDownLatch start = new CountDownLatch(1);
+        List<Future<SsfJtiDedupStore.Claim>> claims = new ArrayList<>();
+        try (ExecutorService executor = Executors.newFixedThreadPool(threads)) {
+            for (int i = 0; i < threads; i++) {
+                JdbcSsfJtiDedupStore store = dedupStore();
+                claims.add(executor.submit(() -> {
+                    start.await();
+                    return store.claim(token(jti));
+                }));
+            }
+            start.countDown();
+            List<SsfJtiDedupStore.Claim> results = new ArrayList<>();
+            for (Future<SsfJtiDedupStore.Claim> claim : claims) {
+                results.add(claim.get(30, TimeUnit.SECONDS));
+            }
+            assertThat(results).doesNotContain((SsfJtiDedupStore.Claim) null);
+            return results.stream().filter((claim) -> claim == SsfJtiDedupStore.Claim.NEW).toList();
+        }
+    }
+
+    private JdbcSsfJtiDedupStore dedupStore() {
+        JdbcSsfJtiDedupStore store = new JdbcSsfJtiDedupStore(this.jdbc, PREFIX);
+        store.setClock(this.clock);
+        return store;
+    }
+
+    private String state(String jti) {
+        return this.jdbc.query("SELECT STATE FROM EASYSSF_PROCESSED_SET WHERE JTI = ?", (row) -> row.getString(1), jti)
+            .get(0);
+    }
+
+    private SsfEventToken token(String jti) {
+        return new SsfEventToken(jti, this.transmitter.issuer(), Instant.EPOCH, List.of(), Map.of(), null, null,
+                Map.of());
+    }
+
     private Receiver receiver() {
+        return receiver(this.dedupStore);
+    }
+
+    private Receiver receiver(SsfJtiDedupStore dedupStore) {
         Receiver receiver = new Receiver();
         SsfHttpClient http = new JdkSsfHttpClient();
         NimbusSsfSetVerifier verifier = new NimbusSsfSetVerifier(this.transmitter.issuer(), transmitter::jwksUri, http);
         SsfEventHandler handler = (eventContext) -> {
             receiver.handled.add(eventContext.eventToken().jti());
+            receiver.keys.add(eventContext.idempotencyKey());
             receiver.afterHandling.run();
         };
-        SsfSetProcessor processor = new SsfSetProcessor(verifier, this.dedupStore, List.of(handler));
+        SsfSetProcessor processor = new SsfSetProcessor(verifier, dedupStore, List.of(handler));
         receiver.poller = new SsfPoller(http, () -> TestTransmitter.ACCESS_TOKEN,
                 () -> URI.create(this.transmitter.pollUri()), processor);
         receiver.poller.setAckStore(this.ackStore);
@@ -220,10 +357,41 @@ private static final class Receiver {
 
         private final List<String> handled = new ArrayList<>();
 
+        private final List<String> keys = new ArrayList<>();
+
         /** runs after a SET was handled, before it is acknowledged */
         private Runnable afterHandling = () -> {
         };
 
     }
 
+    /**
+     * A clock the tests move forward past the lease; it starts at the real time so that
+     * the SETs of the transmitter are not in the future.
+     */
+    private static final class MutableClock extends Clock {
+
+        private Instant now = Instant.now();
+
+        void advance(Duration duration) {
+            this.now = this.now.plus(duration);
+        }
+
+        @Override
+        public ZoneId getZone() {
+            return ZoneOffset.UTC;
+        }
+
+        @Override
+        public Clock withZone(ZoneId zone) {
+            return this;
+        }
+
+        @Override
+        public Instant instant() {
+            return this.now;
+        }
+
+    }
+
 }