Add a proper LRU cache for the exposures

Author: manuel-alvarez-alvarez

dd-java-agent/agent-featureflag/build.gradle

@@ -12,6 +12,12 @@ java {
   targetCompatibility = JavaVersion.VERSION_1_8
 }
 
+excludedClassesCoverage += [
+  // POJOs
+  'com.datadog.featureflag.ExposureCache.Key',
+  'com.datadog.featureflag.ExposureCache.Value'
+]
+
 dependencies {
   api libs.slf4j
   implementation libs.moshi

dd-java-agent/agent-featureflag/src/main/java/com/datadog/featureflag/ExposureCache.java

@@ -0,0 +1,61 @@
+package com.datadog.featureflag;
+
+import datadog.trace.api.featureflag.exposure.ExposureEvent;
+import java.util.Objects;
+
+public interface ExposureCache {
+
+  boolean add(ExposureEvent event);
+
+  Value get(Key key);
+
+  int size();
+
+  final class Key {
+    public final String flag;
+    public final String subject;
+
+    public Key(final ExposureEvent event) {
+      this.flag = event.flag == null ? null : event.flag.key;
+      this.subject = event.subject == null ? null : event.subject.id;
+    }
+
+    @Override
+    public boolean equals(final Object o) {
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      final Key key = (Key) o;
+      return Objects.equals(flag, key.flag) && Objects.equals(subject, key.subject);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(flag, subject);
+    }
+  }
+
+  final class Value {
+    public final String variant;
+    public final String allocation;
+
+    public Value(final ExposureEvent event) {
+      this.variant = event.variant == null ? null : event.variant.key;
+      this.allocation = event.allocation == null ? null : event.allocation.key;
+    }
+
+    @Override
+    public boolean equals(final Object o) {
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      final Value value = (Value) o;
+      return Objects.equals(variant, value.variant) && Objects.equals(allocation, value.allocation);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(variant, allocation);
+    }
+  }
+}

dd-java-agent/agent-featureflag/src/main/java/com/datadog/featureflag/ExposureWriterImpl.java

@@ -13,13 +13,10 @@
 import datadog.trace.api.featureflag.FeatureFlagGateway;
 import datadog.trace.api.featureflag.exposure.ExposureEvent;
 import datadog.trace.api.featureflag.exposure.ExposuresRequest;
-import datadog.trace.core.util.LRUCache;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import okhttp3.Headers;
 import okhttp3.HttpUrl;
@@ -35,7 +32,7 @@ public class ExposureWriterImpl implements ExposureWriter {
   private static final Logger LOGGER = LoggerFactory.getLogger(ExposureWriterImpl.class);
   private static final int DEFAULT_CAPACITY = 1 << 16; // 65536 elements
   private static final int DEFAULT_FLUSH_INTERVAL_IN_SECONDS = 1;
-  private static final int FLUSH_THRESHOLD = 50;
+  private static final int FLUSH_THRESHOLD = 100;
   private static final String EXPOSURES_API_PATH = "api/v2/exposures";
   private static final String EVP_SUBDOMAIN_HEADER_NAME = "X-Datadog-EVP-Subdomain";
   private static final String EVP_SUBDOMAIN_HEADER_VALUE = "event-platform-intake";
@@ -62,7 +59,7 @@ public ExposureWriterImpl(final HttpUrl agentUrl, final Config config) {
             agentUrl.toString()
                 + DDAgentFeaturesDiscovery.V2_EVP_PROXY_ENDPOINT
                 + EXPOSURES_API_PATH);
-    final Map<String, String> context = new HashMap<>();
+    final Map<String, String> context = new HashMap<>(4);
     context.put("service", config.getServiceName() == null ? "unknown" : config.getServiceName());
     if (config.getEnv() != null) {
       context.put("env", config.getEnv());
@@ -103,7 +100,7 @@ private static class ExposureSerializingHandler implements Runnable {
     private final Headers headers;
 
     private final Map<String, String> context;
-    private final Set<ExposureEvent> cache;
+    private final ExposureCache cache;
 
     private final List<ExposureEvent> buffer = new ArrayList<>();
 
@@ -115,7 +112,7 @@ public ExposureSerializingHandler(
         final Headers headers,
         final Map<String, String> context) {
       this.queue = queue;
-      this.cache = Collections.newSetFromMap(new LRUCache<>(queue.capacity()));
+      this.cache = new LRUExposureCache(queue.capacity());
       this.jsonAdapter = new Moshi.Builder().build().adapter(ExposuresRequest.class);
       this.httpClient = new OkHttpClient();
       this.submissionUrl = submissionUrl;
@@ -159,7 +156,7 @@ private void consumeBatch() {
     }
 
     /** Adds an element to the buffer taking care of duplicated exposures thanks to the LRU cache */
-    private boolean addToBuffer(ExposureEvent event) {
+    private boolean addToBuffer(final ExposureEvent event) {
       if (cache.add(event)) {
         buffer.add(event);
         return true;

dd-java-agent/agent-featureflag/src/main/java/com/datadog/featureflag/LRUExposureCache.java

@@ -0,0 +1,55 @@
+package com.datadog.featureflag;
+
+import datadog.trace.api.featureflag.exposure.ExposureEvent;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+public class LRUExposureCache implements ExposureCache {
+
+  private final Map<Key, Value> cache;
+
+  public LRUExposureCache(final int capacity) {
+    cache = new FIFOCache<>(capacity);
+  }
+
+  @Override
+  public boolean add(final ExposureEvent event) {
+    final Key key = new Key(event);
+    final Value oldValue = cache.get(key);
+    if (oldValue == null) {
+      cache.put(key, new Value(event));
+      return true;
+    }
+    final Value newValue = new Value(event);
+    if (!newValue.equals(oldValue)) {
+      cache.remove(key); // ensure LRU semantics
+      cache.put(key, newValue);
+      return true;
+    }
+    return false;
+  }
+
+  @Override
+  public Value get(final Key key) {
+    return cache.get(key);
+  }
+
+  @Override
+  public int size() {
+    return cache.size();
+  }
+
+  private static class FIFOCache<K, V> extends LinkedHashMap<K, V> {
+
+    private final int capacity;
+
+    private FIFOCache(final int capacity) {
+      this.capacity = capacity;
+    }
+
+    @Override
+    protected boolean removeEldestEntry(final Map.Entry<K, V> eldest) {
+      return size() > capacity;
+    }
+  }
+}

dd-java-agent/agent-featureflag/src/test/groovy/com/datadog/featureflag/ExposureWriterTests.groovy

@@ -69,7 +69,7 @@ class ExposureWriterTests extends DDSpecification {
       getEnv() >> env
       getVersion() >> version
     }
-    def exposures = (1..5).collect { buildExposure(System.currentTimeMillis(), it) }
+    def exposures = (1..5).collect { buildExposure() }
     def writer = new ExposureWriterImpl(1 << 4, 100, MILLISECONDS, serverUrl, config)
     writer.init()
 
@@ -89,7 +89,7 @@ class ExposureWriterTests extends DDSpecification {
         }
       }
       final received = requests*.exposures.flatten() as List<ExposureEvent>
-      assert received.containsAll(exposures)
+      assertExposures(received, exposures)
     }
 
     cleanup:
@@ -105,30 +105,29 @@ class ExposureWriterTests extends DDSpecification {
 
   void 'test lru cache'() {
     setup:
-    def ts = System.currentTimeMillis()
-    def exposures = (0..5).collect { buildExposure(ts, it) }
+    def exposures = (0..5).collect { buildExposure() }
     def writer = new ExposureWriterImpl(1 << 4, 100, MILLISECONDS, serverUrl, Config.get())
     writer.init()
 
     when: 'populating the cache'
     exposures.each { writer.accept(it) }
 
-    then: 'all exposures are written'
-    poll.eventually {
+    then: 'all events are written'
+    new PollingConditions(timeout: 1).eventually {
       requests*.exposures.flatten().size() == exposures.size()
     }
 
-    when: 'publishing duplicate exposures'
+    when: 'publishing duplicate events'
     exposures.each { writer.accept(it) }
 
-    then: 'no exposures are written'
+    then: 'no events are written'
     MILLISECONDS.sleep(300) // wait until a flush happens
     requests*.exposures.flatten().size() == exposures.size()
 
-    when: 'a new exposure is generated'
-    writer.accept(buildExposure(System.currentTimeMillis(), 'new'))
+    when: 'a new event is generated'
+    writer.accept(buildExposure())
 
-    then: 'oldest exposure is evicted and the new one is submitted'
+    then: 'oldest event is evicted and the new one is submitted'
     poll.eventually {
       requests*.exposures.flatten().size() == exposures.size() + 1
     }
@@ -144,7 +143,7 @@ class ExposureWriterTests extends DDSpecification {
     def threads = Runtime.runtime.availableProcessors()
     def executor = Executors.newFixedThreadPool(threads)
     def exposures = (1..(threads * exposuresPerThread)).collect {
-      buildExposure(System.currentTimeMillis(), it)
+      buildExposure()
     }
     def latch = new CountDownLatch(1)
     def writer = new ExposureWriterImpl(serverUrl, Config.get())
@@ -166,11 +165,8 @@ class ExposureWriterTests extends DDSpecification {
     then:
     futures.each { it.get() } // wait for all threads to finish
     poll.eventually {
-      final received = requests*.exposures.flatten()
-      final receivedSize = received.size()
-      final expectedSize = exposures.size()
-      assert receivedSize == expectedSize
-      assert received.containsAll(exposures)
+      final received = requests*.exposures.flatten() as List<ExposureEvent>
+      assertExposures(received, exposures)
     }
 
     cleanup:
@@ -187,7 +183,7 @@ class ExposureWriterTests extends DDSpecification {
     writer.init()
 
     when:
-    writer.accept(buildExposure(System.currentTimeMillis(), serviceName))
+    writer.accept(buildExposure())
 
     then:
     MILLISECONDS.sleep(500) // wait for a flush to happen
@@ -207,13 +203,66 @@ class ExposureWriterTests extends DDSpecification {
     'fail-forever' | true
   }
 
-  private static ExposureEvent buildExposure(final long ts, final Object suffix) {
+  private static void assertExposures(final List<ExposureEvent> receivedExposures, final List<ExposureEvent> expectedExposures) {
+    assert receivedExposures.size() == expectedExposures.size()
+    final received = new TreeSet<ExposureEvent>(ExposureWriterTests::compare)
+    received.addAll(expectedExposures)
+    assert received.containsAll(expectedExposures)
+  }
+
+  private static int compare(final ExposureEvent a, final ExposureEvent b) {
+    if (a.is(b)) {
+      return 0
+    }
+    if (a == null) {
+      return -1
+    }
+    if (b == null) {
+      return 1
+    }
+
+    def result = a.timestamp <=> b.timestamp
+    if (result) {
+      return result
+    }
+
+    result = (a.flag?.key ?: '') <=> (b.flag?.key ?: '')
+    if (result) {
+      return result
+    }
+
+    result = (a.variant?.key ?: '') <=> (b.variant?.key ?: '')
+    if (result) {
+      return result
+    }
+
+    result = (a.allocation?.key ?: '') <=> (b.allocation?.key ?: '')
+    if (result) {
+      return result
+    }
+
+    result = (a.subject?.id ?: '') <=> (b.subject?.id ?: '')
+    if (result) {
+      return result
+    }
+
+    final aEntry = a.subject?.attributes?.entrySet()?.iterator()?.next()
+    final bEntry = b.subject?.attributes?.entrySet()?.iterator()?.next()
+    result = (aEntry?.key ?: '') <=> (bEntry?.key ?: '')
+    if (result) {
+      return result
+    }
+    return (aEntry?.value?.toString() ?: '') <=> (bEntry?.value?.toString() ?: '')
+  }
+
+  private static ExposureEvent buildExposure() {
+    final idx = UUID.randomUUID().toString()
     return new ExposureEvent(
-      ts,
-      new Allocation("Allocation_$suffix"),
-      new Flag("Flag_$suffix"),
-      new Variant("Variant_$suffix"),
-      new Subject("Subject_$suffix", [("key_$suffix".toString()): "value_$suffix".toString()])
+      System.currentTimeMillis(),
+      new Allocation("Allocation_$idx"),
+      new Flag("Flag_$idx"),
+      new Variant("Variant_$idx"),
+      new Subject("Subject_$idx", [("key_$idx".toString()): "value_$idx".toString()])
       )
   }
 }