Bound feature flag evaluation aggregation memory

products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/FlagEvaluationAggregator.java

@@ -4,13 +4,17 @@
 import static datadog.trace.util.HashingUtils.hash;
 
 import datadog.trace.api.featureflag.flagevaluation.FlagEvalEvent;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.TreeMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
+@SuppressFBWarnings(
+    value = {"AT_NONATOMIC_64BIT_PRIMITIVE", "AT_NONATOMIC_OPERATIONS_ON_SHARED_VARIABLE"},
+    justification = "The aggregator is confined to the single flag-evaluation serializer thread")
 final class FlagEvaluationAggregator {
 
   // Design assumptions — document the scale we sized for
@@ -32,6 +36,7 @@ final class FlagEvaluationAggregator {
   static final int GLOBAL_CAP = 131_072; // nearest power of two above FULL_BUCKET_SIZING_BASIS
   static final int PER_FLAG_CAP = PER_FLAG_BUCKET_SIZING_BASIS;
   static final int DEGRADED_CAP = 32_768; // nearest power of two above DEGRADED_BUCKET_SIZING_BASIS
+  static final long RETAINED_BYTE_BUDGET = 64L << 20;
 
   private static final byte CTX_TAG_STRING = 's';
   private static final byte CTX_TAG_BOOL = 'b';
@@ -45,7 +50,20 @@ final class FlagEvaluationAggregator {
   final Map<DegradedKey, EvalBucket> degradedTier = new HashMap<>();
   final Map<String, Integer> perFlagCount = new HashMap<>();
   final AtomicLong droppedDegradedOverflow = new AtomicLong(0);
+  final AtomicLong droppedByteBudget = new AtomicLong(0);
+  final AtomicLong degradedCardinalityCap = new AtomicLong(0);
+  final AtomicLong degradedByteBudget = new AtomicLong(0);
   final AtomicInteger globalFullCount = new AtomicInteger(0);
+  private final long retainedByteBudget;
+  private long retainedBytes;
+
+  FlagEvaluationAggregator() {
+    this(RETAINED_BYTE_BUDGET);
+  }
+
+  FlagEvaluationAggregator(final long retainedByteBudget) {
+    this.retainedByteBudget = Math.max(0, retainedByteBudget);
+  }
 
   void aggregate(final FlagEvalEvent event) {
     final boolean isDefault = event.variant == null;
@@ -66,50 +84,78 @@ void aggregate(final FlagEvalEvent event) {
 
     final int flagCount = perFlagCount.getOrDefault(event.flagKey, 0);
     if (globalFullCount.get() < GLOBAL_CAP && flagCount < PER_FLAG_CAP) {
-      fullTier.put(
-          fullKey,
-          new EvalBucket(
-              event.flagKey,
-              event.variant,
-              event.allocationKey,
-              event.targetingKey,
-              event.errorMessage,
-              event.evalTimeMs,
-              isDefault,
-              prunedAttrs,
-              observeFullEvaluationData));
-      globalFullCount.incrementAndGet();
-      perFlagCount.put(event.flagKey, flagCount + 1);
-      return;
+      final long bucketBytes =
+          FlagEvaluationMemoryEstimator.fullBucketBytes(event, prunedAttrs, ctxKey);
+      if (reserve(bucketBytes)) {
+        fullTier.put(
+            fullKey,
+            new EvalBucket(
+                event.flagKey,
+                event.variant,
+                event.allocationKey,
+                event.targetingKey,
+                event.errorMessage,
+                event.evalTimeMs,
+                isDefault,
+                prunedAttrs,
+                observeFullEvaluationData));
+        globalFullCount.incrementAndGet();
+        perFlagCount.put(event.flagKey, flagCount + 1);
+        return;
+      }
     }
 
+    final boolean degradedByByteBudget =
+        globalFullCount.get() < GLOBAL_CAP && flagCount < PER_FLAG_CAP;
     final DegradedKey degradedKey = buildDegradedKey(event);
     bucket = degradedTier.get(degradedKey);
     if (bucket != null) {
       bucket.merge(event.evalTimeMs, isDefault);
       bucket.observeFullEvaluationData &= observeFullEvaluationData;
+      countDegraded(degradedByByteBudget);
       return;
     }
 
     if (degradedTier.size() < DEGRADED_CAP) {
-      degradedTier.put(
-          degradedKey,
-          new EvalBucket(
-              event.flagKey,
-              event.variant,
-              event.allocationKey,
-              null,
-              event.errorMessage,
-              event.evalTimeMs,
-              isDefault,
-              null,
-              observeFullEvaluationData));
+      if (reserve(FlagEvaluationMemoryEstimator.degradedBucketBytes(event))) {
+        degradedTier.put(
+            degradedKey,
+            new EvalBucket(
+                event.flagKey,
+                event.variant,
+                event.allocationKey,
+                null,
+                event.errorMessage,
+                event.evalTimeMs,
+                isDefault,
+                null,
+                observeFullEvaluationData));
+        countDegraded(degradedByByteBudget);
+        return;
+      }
+      droppedByteBudget.incrementAndGet();
       return;
     }
 
     droppedDegradedOverflow.incrementAndGet();
   }
 
+  private void countDegraded(final boolean degradedByByteBudget) {
+    if (degradedByByteBudget) {
+      degradedByteBudget.incrementAndGet();
+    } else {
+      degradedCardinalityCap.incrementAndGet();
+    }
+  }
+
+  private boolean reserve(final long bytes) {
+    if (bytes > retainedByteBudget - retainedBytes) {
+      return false;
+    }
+    retainedBytes += bytes;
+    return true;
+  }
+
   boolean isEmpty() {
     return fullTier.isEmpty() && degradedTier.isEmpty();
   }
@@ -118,14 +164,6 @@ int fullTierSize() {
     return fullTier.size();
   }
 
-  long degradedEvaluationCount() {
-    long count = 0;
-    for (final EvalBucket bucket : degradedTier.values()) {
-      count += bucket.count;
-    }
-    return count;
-  }
-
   int bucketCount() {
     return fullTier.size() + degradedTier.size();
   }
@@ -143,11 +181,22 @@ void clear() {
     degradedTier.clear();
     perFlagCount.clear();
     globalFullCount.set(0);
+    retainedBytes = 0;
+  }
+
+  long retainedBytes() {
+    return retainedBytes;
   }
 
   AggregatedState snapshot() {
     return new AggregatedState(
-        new HashMap<>(fullTier), new HashMap<>(degradedTier), droppedDegradedOverflow.get());
+        new HashMap<>(fullTier),
+        new HashMap<>(degradedTier),
+        droppedDegradedOverflow.get(),
+        droppedByteBudget.get(),
+        degradedCardinalityCap.get(),
+        degradedByteBudget.get(),
+        retainedBytes);
   }
 
   void simulateFullTierAtCap() {
@@ -441,14 +490,26 @@ static class AggregatedState {
     final Map<FullKey, EvalBucket> fullTier;
     final Map<DegradedKey, EvalBucket> degradedTier;
     final long droppedDegradedOverflow;
+    final long droppedByteBudget;
+    final long degradedCardinalityCap;
+    final long degradedByteBudget;
+    final long retainedBytes;
 
     AggregatedState(
         final Map<FullKey, EvalBucket> fullTier,
         final Map<DegradedKey, EvalBucket> degradedTier,
-        final long droppedDegradedOverflow) {
+        final long droppedDegradedOverflow,
+        final long droppedByteBudget,
+        final long degradedCardinalityCap,
+        final long degradedByteBudget,
+        final long retainedBytes) {
       this.fullTier = fullTier;
       this.degradedTier = degradedTier;
       this.droppedDegradedOverflow = droppedDegradedOverflow;
+      this.droppedByteBudget = droppedByteBudget;
+      this.degradedCardinalityCap = degradedCardinalityCap;
+      this.degradedByteBudget = degradedByteBudget;
+      this.retainedBytes = retainedBytes;
     }
   }
 }

products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/FlagEvaluationMemoryEstimator.java

@@ -0,0 +1,74 @@
+package com.datadog.featureflag;
+
+import datadog.trace.api.featureflag.flagevaluation.FlagEvalEvent;
+import java.util.Map;
+
+/** Conservatively estimates memory retained by one aggregation bucket. */
+final class FlagEvaluationMemoryEstimator {
+
+  // These constants include aligned object headers, references, and amortized HashMap storage.
+  // They intentionally overestimate JOL measurements made with 100-bucket representative graphs.
+  private static final long BUCKET_AND_INDEX_BYTES = 256;
+  private static final long CONTEXT_MAP_BYTES = 64;
+  private static final long CONTEXT_ENTRY_BYTES = 48;
+  private static final long STRING_BYTES = 40;
+  private static final long OTHER_VALUE_BYTES = 64;
+  private static final long MAX_BYTES_PER_CHARACTER = 2;
+
+  private FlagEvaluationMemoryEstimator() {}
+
+  static long fullBucketBytes(
+      final FlagEvalEvent event,
+      final Map<String, Object> prunedAttrs,
+      final String canonicalContextKey) {
+    long bytes = BUCKET_AND_INDEX_BYTES;
+    bytes = add(bytes, stringBytes(event.flagKey));
+    bytes = add(bytes, stringBytes(event.variant));
+    bytes = add(bytes, stringBytes(event.allocationKey));
+    bytes = add(bytes, stringBytes(event.targetingKey));
+    bytes = add(bytes, stringBytes(event.errorMessage));
+    bytes = add(bytes, stringBytes(canonicalContextKey));
+    if (prunedAttrs == null || prunedAttrs.isEmpty()) {
+      return bytes;
+    }
+
+    bytes = add(bytes, CONTEXT_MAP_BYTES);
+    for (final Map.Entry<String, Object> entry : prunedAttrs.entrySet()) {
+      bytes = add(bytes, CONTEXT_ENTRY_BYTES);
+      bytes = add(bytes, stringBytes(entry.getKey()));
+      bytes = add(bytes, contextValueBytes(entry.getValue()));
+    }
+    return bytes;
+  }
+
+  static long degradedBucketBytes(final FlagEvalEvent event) {
+    long bytes = BUCKET_AND_INDEX_BYTES;
+    bytes = add(bytes, stringBytes(event.flagKey));
+    bytes = add(bytes, stringBytes(event.variant));
+    bytes = add(bytes, stringBytes(event.allocationKey));
+    return add(bytes, stringBytes(event.errorMessage));
+  }
+
+  private static long contextValueBytes(final Object value) {
+    if (value instanceof String) {
+      return stringBytes((String) value);
+    }
+    if (value == null) {
+      return 0;
+    }
+    return add(OTHER_VALUE_BYTES, characterBytes(value.toString().length()));
+  }
+
+  private static long stringBytes(final String value) {
+    return value == null ? 0 : add(STRING_BYTES, characterBytes(value.length()));
+  }
+
+  private static long characterBytes(final int length) {
+    return MAX_BYTES_PER_CHARACTER * length;
+  }
+
+  private static long add(final long left, final long right) {
+    // Both inputs describe live Java objects. Their sum cannot approach the long range in one JVM.
+    return left + right;
+  }
+}

products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/FlagEvaluationWriterImpl.java

@@ -40,15 +40,15 @@
  * string identity). Context pruning: deterministic (sort before cut), <=256 fields, string values
  * <=256 chars; the pruned attributes are what gets aggregated and serialized. Caps:
  * globalCap=131072, perFlagCap=10000, degradedCap=32768. Eval-time: min/max of
- * firstEvalMs/lastEvalMs across events in the same bucket. Runtime default: absent variant means
- * runtimeDefaultUsed=true. Flush interval: 10 seconds. Queue: bounded MessagePassingBlockingQueue
- * (capacity 2^16), non-blocking offer; on overflow the event is dropped and the
- * droppedQueueOverflow counter is incremented and surfaced on flush. Enqueue: lock-free. Producers
- * contend only on the MPSC queue, never on a monitor, so evaluation threads do not serialize
- * against each other. Shutdown: close() drains the queue and performs a final flush before the
- * worker thread exits. Because enqueue is lock-free, a producer can still offer during shutdown;
- * close() sweeps the queue once the worker has been joined, counting any remainder as a closed drop
- * so shutdown loss is observable rather than silent.
+ * firstEvalMs/lastEvalMs across events in the same bucket. Retained aggregation memory is limited
+ * to 64 MiB. Runtime default: absent variant means runtimeDefaultUsed=true. Flush interval: 10
+ * seconds. Queue: bounded MessagePassingBlockingQueue (capacity 2^12), non-blocking offer; on
+ * overflow the event is dropped and the droppedQueueOverflow counter is incremented and surfaced on
+ * flush. Enqueue: lock-free. Producers contend only on the MPSC queue, never on a monitor, so
+ * evaluation threads do not serialize against each other. Shutdown: close() drains the queue and
+ * performs a final flush before the worker thread exits. Because enqueue is lock-free, a producer
+ * can still offer during shutdown; close() sweeps the queue once the worker has been joined,
+ * counting any remainder as a closed drop so shutdown loss is observable rather than silent.
  */
 public class FlagEvaluationWriterImpl implements FlagEvaluationWriter {
 
@@ -65,8 +65,10 @@ public class FlagEvaluationWriterImpl implements FlagEvaluationWriter {
   static final String DROP_REASON_QUEUE_OVERFLOW = "queue_overflow";
   static final String DROP_REASON_CLOSED = "closed";
   static final String DROP_REASON_DEGRADED_CAP = "degraded_cap";
+  static final String DROP_REASON_BYTE_BUDGET = "byte_budget";
   static final String DROP_REASON_PAYLOAD_LIMIT = "payload_limit";
   static final String DEGRADED_REASON_CARDINALITY_CAP = "cardinality_cap";
+  static final String DEGRADED_REASON_BYTE_BUDGET = "byte_budget";
   static final String DEGRADED_REASON_PAYLOAD_LIMIT = "payload_limit";
   private static final String FLAG_EVALUATION_ROUTE = "flagevaluation";
   private static final CoreMetricCollector CORE_METRICS = CoreMetricCollector.getInstance();
@@ -449,6 +451,22 @@ void flush() {
                 + " (best-effort telemetry)",
             dgDrops);
       }
+      final long byteBudgetDrops = aggregator.droppedByteBudget.getAndSet(0);
+      countMetric(FLAG_EVALUATION_DROPPED_METRIC, byteBudgetDrops, DROP_REASON_BYTE_BUDGET);
+      if (byteBudgetDrops > 0) {
+        LOGGER.warn(
+            "flag evaluation aggregation byte budget full - dropped {} evaluation(s)"
+                + " (best-effort telemetry)",
+            byteBudgetDrops);
+      }
+      countMetric(
+          FLAG_EVALUATION_DEGRADED_METRIC,
+          aggregator.degradedCardinalityCap.getAndSet(0),
+          DEGRADED_REASON_CARDINALITY_CAP);
+      countMetric(
+          FLAG_EVALUATION_DEGRADED_METRIC,
+          aggregator.degradedByteBudget.getAndSet(0),
+          DEGRADED_REASON_BYTE_BUDGET);
 
       // Drain per-reason context-truncation counters and emit one metric per unique reason tag.
       for (final Map.Entry<String, AtomicLong> entry : contextTruncatedCounts.entrySet()) {
@@ -462,10 +480,6 @@ void flush() {
         return;
       }
       try {
-        countMetric(
-            FLAG_EVALUATION_DEGRADED_METRIC,
-            aggregator.degradedEvaluationCount(),
-            DEGRADED_REASON_CARDINALITY_CAP);
         final List<FlagEvaluationPayloads.FlagEvaluationEvent> events = buildEventList();
         if (events.isEmpty()) {
           return;