Extend MemoryGovernor to heap usage and 5 more antagonists; bump chaos heap to 3072m

Cgroup-only signal missed heap exhaustion when Xmx is a small fraction of the
cgroup limit (job 1860026763). Governor now also watches heap usage; pace()
wired into weakref-wave, hidden-class-churn, context-hop, bounded-pool,
consumer-group alongside the existing alloc-storm/direct-memory/dump-storm.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: jbachorik

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/BoundedThreadPoolAntagonist.java

@@ -96,6 +96,7 @@ private void schedulePoolTasks(final int poolIdx) {
                     @Override
                     public void run() {
                         if (!running) return;
+                        MemoryGovernor.pace();
                         long seed = System.nanoTime();
                         sink.addAndGet(burn(seed));
                         ScheduledExecutorService sibling = pools[nextIdx];

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/ConsumerGroupAntagonist.java

@@ -104,6 +104,7 @@ private void rebalanceLoop() {
                 return;
             }
             if (!running) return;
+            MemoryGovernor.pace();
 
             // Interrupt all victims simultaneously (burst)
             Thread[] victims = new Thread[BURST_SIZE];

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/ContextHopAntagonist.java

@@ -76,6 +76,7 @@ public void stopGracefully(Duration timeout) {
 
     private void startChain(final int chainId) {
         if (!running) return;
+        MemoryGovernor.pace();
         final long seed = ThreadLocalRandom.current().nextLong() ^ ((long) chainId << 32);
         CompletableFuture
                 .runAsync(new Runnable() {

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/HiddenClassChurnAntagonist.java

@@ -110,6 +110,7 @@ private void loop() {
                 Thread.currentThread().interrupt();
                 return;
             }
+            MemoryGovernor.pace();
         }
     }
 

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/MemoryGovernor.java

@@ -10,30 +10,43 @@
 package com.datadoghq.profiler.chaos;
 
 import java.io.IOException;
+import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryUsage;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 
 /**
- * Shared memory-pressure gate for antagonists that burst allocations fast
- * enough to threaten the container's cgroup ceiling before GC/Cleaner
- * catches up (alloc-storm, direct-memory, dump-storm).
+ * Shared memory-pressure gate for antagonists whose allocation/thread-churn
+ * rate can outrun GC/cleanup and threaten either the container's cgroup
+ * ceiling or the JVM heap itself (alloc-storm, direct-memory, dump-storm,
+ * weakref-wave, hidden-class-churn, context-hop, bounded-pool,
+ * consumer-group).
  *
- * <p>A single background thread samples the cgroup memory usage/limit (v2
- * {@code memory.current}/{@code memory.max}, falling back to v1 {@code
- * memory.usage_in_bytes}/{@code memory.limit_in_bytes}) and flips a shared
- * {@code throttled} flag with hysteresis: entering throttle above {@link
- * #HIGH_WATERMARK} of the limit, clearing it below {@link #LOW_WATERMARK}.
- * Antagonists poll {@link #pace()} once per loop iteration — a volatile read
- * plus, only while throttled, a short sleep. If no cgroup limit is readable
- * (e.g. running outside a container), the governor stays permanently
- * disabled rather than guessing at a ceiling.
+ * <p>A single background thread samples two independent signals every tick:
+ * cgroup memory usage/limit (v2 {@code memory.current}/{@code memory.max},
+ * falling back to v1 {@code memory.usage_in_bytes}/{@code
+ * memory.limit_in_bytes}) and JVM heap usage ({@link
+ * java.lang.management.MemoryMXBean#getHeapMemoryUsage()}). Either signal
+ * crossing its high watermark sets the shared {@code throttled} flag;
+ * clearing it requires both signals to be back below their low watermark
+ * (hysteresis, so a single antagonist bouncing near one ceiling doesn't
+ * flap the gate for antagonists driving the other). A signal that has no
+ * readable ceiling (no cgroup limit, or a JVM without {@code -Xmx}) is
+ * treated as permanently low rather than guessed at.
+ *
+ * <p>Antagonists poll {@link #pace()} once per loop iteration — a volatile
+ * read plus, only while throttled, a short sleep.
  */
 final class MemoryGovernor {
 
-    private static final double HIGH_WATERMARK = 0.85;
-    private static final double LOW_WATERMARK = 0.70;
+    private static final double CGROUP_HIGH_WATERMARK = 0.85;
+    private static final double CGROUP_LOW_WATERMARK = 0.70;
+    // Heap-space OOME is less recoverable than a cgroup kill (no headroom to
+    // shed off-heap first), so back off earlier and require more slack.
+    private static final double HEAP_HIGH_WATERMARK = 0.80;
+    private static final double HEAP_LOW_WATERMARK = 0.65;
     private static final long SAMPLE_INTERVAL_MS = 500;
     private static final long THROTTLE_SLEEP_MS = 5;
 
@@ -51,14 +64,17 @@ final class MemoryGovernor {
     private MemoryGovernor() {
     }
 
-    /** Starts the background sampler if a cgroup memory limit is readable. Idempotent-ish: call once from Main. */
+    /** Starts the background sampler. Cgroup and/or heap signals no-op individually if unreadable. Call once from Main. */
     static void start() {
         Path usagePath = firstReadable(USAGE_PATHS);
         long limitBytes = firstLimit(LIMIT_PATHS);
         if (usagePath == null || limitBytes <= 0) {
-            return; // no container ceiling visible; never throttle
+            usagePath = null;
+            limitBytes = -1;
         }
-        Thread sampler = new Thread(() -> sampleLoop(usagePath, limitBytes), "chaos-memory-governor");
+        Path finalUsagePath = usagePath;
+        long finalLimitBytes = limitBytes;
+        Thread sampler = new Thread(() -> sampleLoop(finalUsagePath, finalLimitBytes), "chaos-memory-governor");
         sampler.setDaemon(true);
         sampler.start();
     }
@@ -78,20 +94,37 @@ private static void sampleLoop(Path usagePath, long limitBytes) {
         while (true) {
             try {
                 Thread.sleep(SAMPLE_INTERVAL_MS);
-                long usage = readLong(usagePath);
-                double fraction = (double) usage / (double) limitBytes;
-                if (fraction >= HIGH_WATERMARK) {
-                    throttled = true;
-                } else if (fraction <= LOW_WATERMARK) {
-                    throttled = false;
-                }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
                 return;
-            } catch (IOException | NumberFormatException e) {
-                // Transient read failure; keep last known state and retry next tick.
             }
+            try {
+                update(usagePath, limitBytes);
+            } catch (IOException | NumberFormatException | OutOfMemoryError e) {
+                // Transient read/allocation failure; keep last known state and retry next tick.
+            }
+        }
+    }
+
+    private static void update(Path usagePath, long limitBytes) throws IOException {
+        double cgroupFraction = 0.0;
+        if (usagePath != null) {
+            cgroupFraction = (double) readLong(usagePath) / (double) limitBytes;
+        }
+        double heapFraction = 0.0;
+        MemoryUsage heap = ManagementFactory.getMemoryMXBean().getHeapMemoryUsage();
+        if (heap.getMax() > 0) {
+            heapFraction = (double) heap.getUsed() / (double) heap.getMax();
+        }
+
+        boolean high = cgroupFraction >= CGROUP_HIGH_WATERMARK || heapFraction >= HEAP_HIGH_WATERMARK;
+        boolean low = cgroupFraction <= CGROUP_LOW_WATERMARK && heapFraction <= HEAP_LOW_WATERMARK;
+        if (high) {
+            throttled = true;
+        } else if (low) {
+            throttled = false;
         }
+        // else: in the dead zone between watermarks — keep current state.
     }
 
     private static Path firstReadable(Path[] candidates) {

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/WeakRefWaveAntagonist.java

@@ -84,6 +84,7 @@ private void waveLoop() {
                 obj[0] = (byte) i;
                 strongRefs.add(obj);
                 weakRefs.add(new WeakReference<byte[]>(obj));
+                MemoryGovernor.pace();
             }
             // Publish filled list to reader
             currentWave = weakRefs;

utils/run-chaos-harness.sh

@@ -194,9 +194,11 @@ echo "cgroup memory limit: ${CGROUP_MEM_LIMIT}"
 # dump-storm) deliberately burst off-heap/native memory hard and fast, sharing
 # the same cgroup limit as the heap. Reproduced under a 6GiB-capped cgroup and
 # confirmed MemoryGovernor (see chaos/MemoryGovernor.java, wired into the
-# antagonists' hot loops) keeps that same config under the limit, so the heap
-# no longer needs to be trimmed to make room for off-heap bursts.
-HEAP_MB=2560
+# antagonists' hot loops) keeps that same config under the limit. 2560m still
+# ran within a few percent of java.lang.OutOfMemoryError under the profiler
+# config (job 1860026763); 3072m keeps comparable headroom to the off-heap
+# side while giving the heap itself real margin.
+HEAP_MB=3072
 
 HS_ERR="${CHAOS_ERROR_FILE:-${WORK_DIR}/hs_err.log}"
 rm -f "${HS_ERR}"