Skip to content

Commit

Permalink
Remove ReferenceQueue from TaintedMap implementation and purge inline (…
Browse files Browse the repository at this point in the history 
…#6241)

Move to an inline purge map and drop the reference queue
  • Loading branch information
@manuel-alvarez-alvarez
manuel-alvarez-alvarez authored Dec 13, 2023
1 parent 6e9eadb commit b0cf14d
Show file tree
Hide file tree
Showing 18 changed files with 467 additions and 668 deletions.
14 changes: 12 additions & 2 deletions dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/taint/TaintedMapEmptyBenchmark.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.BenchmarkParams;
import org.openjdk.jmh.infra.Blackhole;

@Warmup(iterations = 2, time = 1000, timeUnit = MILLISECONDS)
Expand All @@ -30,8 +31,17 @@ public class TaintedMapEmptyBenchmark {
private final Object anyObject = new Object();

@Setup(Level.Iteration)
public void setup() {
map = new TaintedMap();
public void setup(BenchmarkParams params) {
final boolean baseline = params.getBenchmark().endsWith("baseline");
map = baseline ? TaintedMap.NoOp.INSTANCE : new TaintedMap.TaintedMapImpl();
}

@Benchmark
@OperationsPerInvocation(OP_COUNT)
public void baseline(final Blackhole bh) {
for (int i = 0; i < OP_COUNT; i++) {
bh.consume(map.get(anyObject));
}
}

@Benchmark
Expand Down
16 changes: 9 additions & 7 deletions dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/taint/TaintedMapGetsBenchmark.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.BenchmarkParams;
import org.openjdk.jmh.infra.Blackhole;

@Warmup(iterations = 1, time = 1000, timeUnit = MILLISECONDS)
Expand All @@ -28,35 +29,36 @@
@State(Scope.Benchmark)
public class TaintedMapGetsBenchmark {

private static final int INITIAL_OP_COUNT = TaintedMap.DEFAULT_FLAT_MODE_THRESHOLD;
private static final int INITIAL_OP_COUNT = 1 << 12;
private static final int OP_COUNT = 1024;

private TaintedMap map;
private List<Object> objectList;
private List<Object> initialObjectList;

@Setup(Level.Iteration)
public void setup() {
map = new TaintedMap();
public void setup(BenchmarkParams params) {
final boolean baseline = params.getBenchmark().endsWith("baseline");
map = baseline ? TaintedMap.NoOp.INSTANCE : new TaintedMap.TaintedMapImpl();
initialObjectList = new ArrayList<>(INITIAL_OP_COUNT);
objectList = new ArrayList<>(OP_COUNT);
for (int i = 0; i < INITIAL_OP_COUNT; i++) {
final Object k = new Object();
initialObjectList.add(k);
map.put(new TaintedObject(k, new Range[0], map.getReferenceQueue()));
map.put(new TaintedObject(k, new Range[0]));
}
for (int i = 0; i < OP_COUNT; i++) {
final Object k = new Object();
objectList.add(k);
map.put(new TaintedObject(k, new Range[0], map.getReferenceQueue()));
map.put(new TaintedObject(k, new Range[0]));
}
}

@Benchmark
@OperationsPerInvocation(OP_COUNT)
public void getsBaseline(final Blackhole bh) {
public void baseline(final Blackhole bh) {
for (int i = 0; i < OP_COUNT; i++) {
bh.consume(objectList.get(i));
bh.consume(map.get(objectList.get(i)));
}
}

Expand Down
61 changes: 47 additions & 14 deletions dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/taint/TaintedMapPutsBenchmark.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@
import com.datadog.iast.model.Range;
import datadog.trace.test.util.CircularBuffer;
import java.util.ArrayList;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
Expand All @@ -20,45 +22,47 @@
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Timeout;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.Blackhole;
import org.openjdk.jmh.infra.BenchmarkParams;

@Warmup(iterations = 1, time = 1000, timeUnit = MILLISECONDS)
@Measurement(iterations = 3, time = 1000, timeUnit = MILLISECONDS)
@Warmup(iterations = 3, time = 1000, timeUnit = MILLISECONDS)
@Measurement(iterations = 5, time = 1000, timeUnit = MILLISECONDS)
@Timeout(time = 10000, timeUnit = MILLISECONDS)
@Fork(3)
@OutputTimeUnit(NANOSECONDS)
@BenchmarkMode(Mode.AverageTime)
@State(Scope.Benchmark)
public class TaintedMapPutsBenchmark {

private static final int INITIAL_OP_COUNT = TaintedMap.DEFAULT_FLAT_MODE_THRESHOLD;
private static final int INITIAL_OP_COUNT = 1 << 12;
private static final int OP_COUNT = 1024;

private static final Range[] EMPTY_RANGES = new Range[0];

private TaintedMap map;
private List<Object> initialObjectList;
private CircularBuffer<Object> objectBuffer;
private GarbageCollectorHandler gcHandler;

@Setup(Level.Iteration)
public void setup() {
map = new TaintedMap();
objectBuffer = new CircularBuffer<>(OP_COUNT);
public void setup(BenchmarkParams params) {
final boolean baseline = params.getBenchmark().endsWith("baseline");
map = baseline ? TaintedMap.NoOp.INSTANCE : new TaintedMap.TaintedMapImpl();
gcHandler = new GarbageCollectorHandler(OP_COUNT);
initialObjectList = new ArrayList<>(INITIAL_OP_COUNT);
for (int i = 0; i < INITIAL_OP_COUNT; i++) {
final Object k = new Object();
initialObjectList.add(k);
map.put(new TaintedObject(k, EMPTY_RANGES, map.getReferenceQueue()));
map.put(new TaintedObject(k, EMPTY_RANGES));
}
}

@Benchmark
@OperationsPerInvocation(OP_COUNT)
public void putsBaseline(final Blackhole bh) {
public void baseline() {
for (int i = 0; i < OP_COUNT; i++) {
final Object k = new Object();
objectBuffer.add(k);
bh.consume(new TaintedObject(k, EMPTY_RANGES, map.getReferenceQueue()));
final TaintedObject to = new TaintedObject(k, EMPTY_RANGES);
gcHandler.add(to);
map.put(to);
}
}

Expand All @@ -67,8 +71,37 @@ public void putsBaseline(final Blackhole bh) {
public void puts() {
for (int i = 0; i < OP_COUNT; i++) {
final Object k = new Object();
objectBuffer.add(k);
map.put(new TaintedObject(k, EMPTY_RANGES, map.getReferenceQueue()));
final TaintedObject to = new TaintedObject(k, EMPTY_RANGES);
gcHandler.add(to);
map.put(to);
}
}

/**
* Reference queue that holds a circular buffer of alive objects and enqueues to be purged when
* they are removed
*/
private static class GarbageCollectorHandler {

private final Map<Object, TaintedObject> map;
private final CircularBuffer<Object> alive;

public GarbageCollectorHandler(final int aliveCount) {
map = new IdentityHashMap<>(aliveCount);
alive = new CircularBuffer<>(aliveCount);
}

public void add(TaintedObject reference) {
if (reference == null || reference.get() == null) {
return;
}
final Object referent = reference.get();
final Object toRemove = alive.add(referent);
if (toRemove != null) {
final TaintedObject taintedObject = map.remove(toRemove);
taintedObject.enqueue();
}
map.put(reference.get(), reference);
}
}
}
Loading

0 comments on commit b0cf14d

Please sign in to comment.