fix flakey tests

core/src/main/java/net/staticstudios/data/impl/h2/H2DataAccessor.java

@@ -567,7 +567,15 @@ public void setRedisValue(String holderSchema, String holderTable, String identi
         Runnable runnable = () -> {
             if (value == null) {
                 taskQueue.submitTask((connection, jedis) -> {
-                    jedis.del(key);
+                    redisListener.expectLocalDeleteEvent(key);
+                    boolean deleted = false;
+                    try {
+                        deleted = jedis.del(key) > 0;
+                    } finally {
+                        if (!deleted) {
+                            redisListener.cancelLocalDeleteEvent(key);
+                        }
+                    }
                 });
             } else {
                 taskQueue.submitTask((connection, jedis) -> {

core/src/main/java/net/staticstudios/data/impl/redis/RedisListener.java

@@ -14,13 +14,18 @@
 import java.util.Arrays;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
 
 public class RedisListener extends JedisPubSub {
     private static final Logger logger = LoggerFactory.getLogger(RedisListener.class);
     private final Set<String> listenedPartialKeys = ConcurrentHashMap.newKeySet();
     private final Map<Pattern, RedisEventHandler> handlers = new ConcurrentHashMap<>();
+    private final Map<String, Integer> ignoredLocalDeleteEvents = new ConcurrentHashMap<>();
+    private final CompletableFuture<Void> subscriptionReady = new CompletableFuture<>();
     private final TaskQueue taskQueue;
 
     public RedisListener(DataSourceConfig ds, TaskQueue taskQueue) {
@@ -32,11 +37,18 @@ public RedisListener(DataSourceConfig ds, TaskQueue taskQueue) {
                 if (ThreadUtils.isShuttingDown()) {
                     return;
                 }
+                subscriptionReady.completeExceptionally(e);
                 logger.error("Redis connection lost in listener thread", e);
             }
         });
         listenerThread.start();
 
+        try {
+            subscriptionReady.get(10, TimeUnit.SECONDS);
+        } catch (Exception e) {
+            throw new IllegalStateException("Timed out waiting for the Redis event subscription", e);
+        }
+
         ThreadUtils.onShutdownRunSync(ShutdownStage.CLEANUP, () -> {
             this.punsubscribe();
             listenerThread.interrupt();
@@ -58,6 +70,9 @@ public void onPMessage(String pattern, String channel, String key) {
         if (!key.startsWith("static-data:")) {
             return;
         }
+        if (event == RedisEvent.DEL && consumeLocalDeleteEvent(key)) {
+            return;
+        }
 
         for (Map.Entry<Pattern, RedisEventHandler> entry : handlers.entrySet()) {
             if (entry.getKey().matcher(key).matches()) {
@@ -75,4 +90,26 @@ public void onPMessage(String pattern, String channel, String key) {
             }
         }
     }
+
+    @Override
+    public void onPSubscribe(String pattern, int subscribedChannels) {
+        subscriptionReady.complete(null);
+    }
+
+    public void expectLocalDeleteEvent(String key) {
+        ignoredLocalDeleteEvents.merge(key, 1, Integer::sum);
+    }
+
+    public void cancelLocalDeleteEvent(String key) {
+        consumeLocalDeleteEvent(key);
+    }
+
+    private boolean consumeLocalDeleteEvent(String key) {
+        AtomicBoolean consumed = new AtomicBoolean(false);
+        ignoredLocalDeleteEvents.computeIfPresent(key, (ignoredKey, count) -> {
+            consumed.set(true);
+            return count == 1 ? null : count - 1;
+        });
+        return consumed.get();
+    }
 }

core/src/test/java/net/staticstudios/data/CachedValueTest.java

@@ -11,6 +11,7 @@
 import org.junit.jupiter.api.Test;
 import redis.clients.jedis.Jedis;
 
+import java.util.Objects;
 import java.util.UUID;
 
 import static org.junit.jupiter.api.Assertions.*;
@@ -52,7 +53,7 @@ public void testFallback() {
         assertEquals(false, user.onCooldown.get());
         assertEquals(0, user.cooldownUpdates.get());
 
-        waitForDataPropagation();
+        flushDataManagers();
 
         Jedis jedis = getJedis();
 
@@ -114,9 +115,13 @@ public void testUpdateHandler() {
 
         Jedis jedis = getJedis();
         String onCooldownKey = RedisUtils.buildRedisKey("public", "users", "on_cooldown", user.getIdColumns());
+        dataManager.flushTaskQueue();
         jedis.del(onCooldownKey);
 
-        waitForDataPropagation();
+        awaitCondition(
+                () -> Objects.equals(false, user.onCooldown.get()) && Objects.equals(6, user.cooldownUpdates.get()),
+                "the external Redis deletion to reach the cached value and its update handler"
+        );
 
         assertEquals(false, user.onCooldown.get());
         assertEquals(6, user.cooldownUpdates.get());
@@ -139,19 +144,19 @@ public void testUpdateRedis() {
 
         user.onCooldown.set(true);
         user.cooldownUpdates.set(1);
-        waitForDataPropagation();
+        dataManager.flushTaskQueue();
         assertEquals("true", gson.fromJson(jedis.get(onCooldownKey), RedisEncodedValue.class).value());
         assertEquals("1", gson.fromJson(jedis.get(cooldownUpdatesKey), RedisEncodedValue.class).value());
 
         user.onCooldown.set(null);
         user.cooldownUpdates.set(null);
-        waitForDataPropagation();
+        dataManager.flushTaskQueue();
         assertNull(jedis.get(onCooldownKey));
         assertNull(jedis.get(cooldownUpdatesKey));
 
         user.onCooldown.set(false); //fallback
         user.cooldownUpdates.set(0); //fallback
-        waitForDataPropagation();
+        dataManager.flushTaskQueue();
         assertNull(jedis.get(onCooldownKey));
         assertNull(jedis.get(cooldownUpdatesKey));
     }
@@ -175,7 +180,7 @@ public void testLoadCachedValues() {
 
         jedis.set(cooldownUpdatesKey, gson.toJson(new RedisEncodedValue(null, "5")));
 
-        waitForDataPropagation();
+        awaitCondition(() -> Objects.equals(5, user1.cooldownUpdates.get()), "the external Redis value to reach H2");
 
         assertEquals(5, user1.cooldownUpdates.get());
 
@@ -204,15 +209,15 @@ public void testRefreshCachedValues() throws InterruptedException {
         assertEquals(2, user.counter.refresh());
         assertEquals(2, user.counter.get());
 
-        Thread.sleep(10_000); //wait for the cached value to expire
-
         String counterKey = RedisUtils.buildRedisKey("public", "users", "counter", user.getIdColumns());
 
         Jedis jedis = getJedis();
+        awaitCondition(() -> !jedis.exists(counterKey), "the cached counter to expire");
         assertFalse(jedis.exists(counterKey));
 
-        assertEquals(0, user.counter.get()); //trigger a refresh
-        waitForDataPropagation();
+        awaitCondition(() -> Objects.equals(0, user.counter.get()), "the expiration event to clear and refresh the H2 cached value");
+        assertEquals(0, user.counter.get());
+        dataManager.flushTaskQueue();
         assertEquals("0", gson.fromJson(jedis.get(counterKey), RedisEncodedValue.class).value());
     }
 
@@ -234,17 +239,20 @@ public void testUpdateInterval() throws Exception {
         }
 
         assertEquals(4, user.throttledCounter.get());
-        waitForDataPropagation();
+        dataManager.flushTaskQueue();
 
         Jedis jedis = getJedis();
         String throttledCounterKey = RedisUtils.buildRedisKey("public", "users", "throttled_counter", user.getIdColumns());
 
         assertNull(jedis.get(throttledCounterKey));
 
-        Thread.sleep(6000);
+        awaitCondition(() -> {
+            RedisEncodedValue value = gson.fromJson(jedis.get(throttledCounterKey), RedisEncodedValue.class);
+            return value != null && Objects.equals("4", value.value());
+        }, "the throttled cached value to be written");
 
         RedisEncodedValue encoded = gson.fromJson(jedis.get(throttledCounterKey), RedisEncodedValue.class);
         assertNotNull(encoded);
         assertEquals("4", encoded.value());
     }
-}
+}

core/src/test/java/net/staticstudios/data/PersistentManyToManyCollectionTest.java

@@ -319,7 +319,7 @@ public void testAddHandlerUpdate() throws SQLException {
                 .insert(InsertMode.SYNC);
         List<MockUser> friends = createFriends(5);
         user.friends.addAll(friends);
-        waitForDataPropagation();
+        flushDataManagers();
         assertEquals(5, user.friendAdditions.get());
 
         List<MockUser> otherFriends = createFriends(5);
@@ -334,8 +334,12 @@ public void testAddHandlerUpdate() throws SQLException {
                 preparedStatement.setObject(3, friend.id.get());
                 preparedStatement.executeUpdate();
             }
-            waitForDataPropagation();
-            assertEquals(5 + (++i), user.friendAdditions.get());
+            int expectedAdditions = 5 + (++i);
+            awaitCondition(
+                    () -> user.friendAdditions.get() == expectedAdditions,
+                    "the many-to-many update addition handler"
+            );
+            assertEquals(expectedAdditions, user.friendAdditions.get());
         }
     }
 
@@ -347,7 +351,7 @@ public void testRemoveHandlerUpdate() throws SQLException {
                 .insert(InsertMode.SYNC);
         List<MockUser> friends = createFriends(5);
         user.friends.addAll(friends);
-        waitForDataPropagation();
+        flushDataManagers();
         assertEquals(5, user.friendAdditions.get());
 
         Connection pgConnection = getConnection();
@@ -359,8 +363,12 @@ public void testRemoveHandlerUpdate() throws SQLException {
                 preparedStatement.setObject(2, friend.id.get());
                 preparedStatement.executeUpdate();
             }
-            waitForDataPropagation();
-            assertEquals(++i, user.friendRemovals.get());
+            int expectedRemovals = ++i;
+            awaitCondition(
+                    () -> user.friendRemovals.get() == expectedRemovals,
+                    "the many-to-many delete handler"
+            );
+            assertEquals(expectedRemovals, user.friendRemovals.get());
         }
     }
 
@@ -402,4 +410,4 @@ public void testRemoveHandlerDelete() {
             assertEquals(++i, user.friendRemovals.get());
         }
     }
-}
+}

core/src/test/java/net/staticstudios/data/PersistentOneToManyValueCollectionTest.java

@@ -337,9 +337,12 @@ public void testAddHandlerUpdate() {
             } catch (Exception e) {
                 throw new RuntimeException(e);
             }
-            waitForDataPropagation();
-
-            assertEquals(++i, user.favoriteNumberAdditions.get());
+            int expectedAdditions = ++i;
+            awaitCondition(
+                    () -> user.favoriteNumberAdditions.get() == expectedAdditions,
+                    "the one-to-many value update addition handler"
+            );
+            assertEquals(expectedAdditions, user.favoriteNumberAdditions.get());
         }
     }
 
@@ -381,9 +384,12 @@ public void testRemoveHandlerUpdate() {
             } catch (Exception e) {
                 throw new RuntimeException(e);
             }
-            waitForDataPropagation();
-
-            assertEquals(++i, user.favoriteNumberRemovals.get());
+            int expectedRemovals = ++i;
+            awaitCondition(
+                    () -> user.favoriteNumberRemovals.get() == expectedRemovals,
+                    "the one-to-many value update removal handler"
+            );
+            assertEquals(expectedRemovals, user.favoriteNumberRemovals.get());
         }
     }
 
@@ -425,4 +431,4 @@ public void testRemoveHandlerDelete() {
             assertEquals(++i, user.favoriteNumberRemovals.get());
         }
     }
-}
+}