Add putIfAbsent to IAssignmentCache (#148)

This resolves a concurrency issue that allowed multiple AssignmentLogger callbacks on high contention first-experiment-access.

Author: hbordersTwitch

src/main/java/cloud/eppo/BaseEppoClient.java

@@ -267,17 +267,11 @@ protected EppoValue getTypedAssignment(
         boolean logAssignment = true;
         AssignmentCacheEntry cacheEntry = AssignmentCacheEntry.fromVariationAssignment(assignment);
         if (assignmentCache != null) {
-          if (assignmentCache.hasEntry(cacheEntry)) {
-            logAssignment = false;
-          }
+          logAssignment = assignmentCache.putIfAbsent(cacheEntry);
         }
 
         if (logAssignment) {
           assignmentLogger.logAssignment(assignment);
-
-          if (assignmentCache != null) {
-            assignmentCache.put(cacheEntry);
-          }
         }
 
       } catch (Exception e) {

src/main/java/cloud/eppo/api/AbstractAssignmentCache.java

@@ -17,6 +17,8 @@ public interface CacheDelegate {
     void put(String cacheKey, @NotNull String serializedEntry);
 
     @Nullable String get(String cacheKey);
+
+    boolean putIfAbsent(String cacheKey, @NotNull String serializedEntry);
   }
 
   protected final CacheDelegate delegate;
@@ -38,6 +40,19 @@ public void put(String cacheKey, @NotNull String serializedEntry) {
           public String get(String cacheKey) {
             return delegate.get(cacheKey);
           }
+
+          @Override
+          public boolean putIfAbsent(String cacheKey, @NotNull String serializedEntry) {
+            boolean hadNoPreviousEntry;
+            synchronized (delegate) {
+              String entry = delegate.get(cacheKey);
+              hadNoPreviousEntry = entry == null;
+              if (hadNoPreviousEntry) {
+                delegate.put(cacheKey, serializedEntry);
+              }
+            }
+            return hadNoPreviousEntry;
+          }
         });
   }
 
@@ -55,4 +70,9 @@ private String get(AssignmentCacheKey key) {
   public void put(AssignmentCacheEntry entry) {
     delegate.put(entry.getKeyString(), entry.getValueKeyString());
   }
+
+  @Override
+  public boolean putIfAbsent(AssignmentCacheEntry entry) {
+    return delegate.putIfAbsent(entry.getKeyString(), entry.getValueKeyString());
+  }
 }

src/main/java/cloud/eppo/api/IAssignmentCache.java

@@ -7,6 +7,9 @@
  * determine both presence and uniqueness of the cached value.
  */
 public interface IAssignmentCache {
+  /**
+   * Puts the entry into the cache, overwriting an existing entry.
+   */
   void put(AssignmentCacheEntry entry);
 
   /**
@@ -15,4 +18,10 @@ public interface IAssignmentCache {
    * comparing the `getValueKeyString()` method results.
    */
   boolean hasEntry(AssignmentCacheEntry entry);
+
+  /**
+   * Puts the entry into the cache if none exists.
+   * @return true if no previous cache entry exists, false otherwise
+   */
+  boolean putIfAbsent(AssignmentCacheEntry entry);
 }

src/test/java/cloud/eppo/BaseEppoClientTest.java

@@ -23,12 +23,18 @@
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
 import okhttp3.mockwebserver.RecordedRequest;
 import org.apache.commons.io.FileUtils;
+import org.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -557,6 +563,72 @@ public void testAssignmentEventCorrectlyDeduplicated() {
     verify(mockAssignmentLogger, times(3)).logAssignment(any(Assignment.class));
   }
 
+  @Test
+  public void testAssignmentEventCorrectlyDeduplicatedFromBackgroundThreads() {
+    initClientWithAssignmentCache(new LRUInMemoryAssignmentCache(1024));
+
+    Attributes subjectAttributes = new Attributes();
+    subjectAttributes.put("number", EppoValue.valueOf("123456789"));
+
+    int numThreads = 10;
+    final CountDownLatch threadStartCountDownLatch = new CountDownLatch(numThreads);
+    final CountDownLatch getAssignmentStartCountDownLatch = new CountDownLatch(1);
+    final List<Integer> assignments = Collections.synchronizedList(Arrays.asList(new Integer[numThreads]));
+    try (ExecutorService pool = Executors.newFixedThreadPool(numThreads, new ThreadFactory() {
+      private final AtomicInteger threadIndexAtomicInteger = new AtomicInteger(0);
+      @Override
+      public Thread newThread(@NotNull Runnable runnable) {
+        final int threadIndex = threadIndexAtomicInteger.getAndIncrement();
+        return new Thread(runnable, "testAssignmentEventCorrectlyDeduplicatedFromBackgroundThreads-" + threadIndex);
+      }
+    })) {
+      for (int i = 0; i < numThreads; i += 1) {
+        final int threadIndex = i;
+        pool.execute(
+          () -> {
+            threadStartCountDownLatch.countDown();
+            boolean shouldStart;
+            try {
+              shouldStart = getAssignmentStartCountDownLatch.await(1000, TimeUnit.SECONDS);
+            } catch (InterruptedException ignored) {
+              shouldStart = false;
+            }
+            final Integer assignment;
+            if (shouldStart) {
+              assignment = eppoClient.getIntegerAssignment("numeric-one-of", "alice", subjectAttributes, 0);
+            } else {
+              assignment = null;
+            }
+
+            assignments.set(threadIndex, assignment);
+          }
+        );
+      }
+
+      boolean shouldStart;
+      try {
+        shouldStart = threadStartCountDownLatch.await(2, TimeUnit.SECONDS);
+      } catch (InterruptedException ignored) {
+        shouldStart = false;
+      }
+
+      assertTrue(shouldStart, "All worker threads did not start");
+      getAssignmentStartCountDownLatch.countDown();
+    }
+
+    final List<Integer> expectedAssignments;
+    {
+      final Integer[] expectedAssignmentsArray = new Integer[numThreads];
+      // `2` matches the attribute `number` value of "123456789"
+      Arrays.fill(expectedAssignmentsArray, 2);
+      expectedAssignments = Arrays.asList(expectedAssignmentsArray);
+    }
+    assertEquals(expectedAssignments, assignments);
+
+    // `logAssignment` should be called only once.
+    verify(mockAssignmentLogger, times(1)).logAssignment(any(Assignment.class));
+  }
+
   @Test
   public void testAssignmentLogErrorNonFatal() {
     initClient();