Implement sampling

apm-agent-benchmarks/src/main/java/co/elastic/apm/impl/AbstractReporterBenchmark.java

@@ -10,6 +10,7 @@
 import co.elastic.apm.impl.payload.Service;
 import co.elastic.apm.impl.payload.SystemInfo;
 import co.elastic.apm.impl.payload.TransactionPayload;
+import co.elastic.apm.impl.sampling.ConstantSampler;
 import co.elastic.apm.impl.stacktrace.StacktraceFactory;
 import co.elastic.apm.impl.transaction.Span;
 import co.elastic.apm.impl.transaction.Transaction;
@@ -69,7 +70,7 @@ public void setUp() throws Exception {
         payload = new TransactionPayload(process, service, system);
         for (int i = 0; i < reporterConfiguration.getMaxQueueSize(); i++) {
             Transaction t = new Transaction();
-            t.start(tracer, 0, true);
+            t.start(tracer, 0, ConstantSampler.of(true));
             fillTransaction(t);
             payload.getTransactions().add(t);
         }

apm-agent-core/src/main/java/co/elastic/apm/configuration/CoreConfiguration.java

@@ -57,6 +57,16 @@ public class CoreConfiguration extends ConfigurationOptionProvider {
             "To reduce overhead and storage requirements, you can set the sample rate to a value between 0.0 and 1.0. " +
             "We still record overall time and the result for unsampled transactions, but no context information, tags, or spans.")
         .dynamic(true)
+        .addValidator(new ConfigurationOption.Validator<Double>() {
+            @Override
+            public void assertValid(Double value) {
+                if (value != null) {
+                    if (value < 0 || value > 1) {
+                        throw new IllegalArgumentException("The sample rate must be between 0 and 1");
+                    }
+                }
+            }
+        })
         .buildWithDefault(1.0);
 
     private final ConfigurationOption<Integer> transactionMaxSpans = ConfigurationOption.integerOption()
@@ -87,8 +97,8 @@ public String getEnvironment() {
         return environment.get();
     }
 
-    public double getSampleRate() {
-        return sampleRate.get();
+    public ConfigurationOption<Double> getSampleRate() {
+        return sampleRate;
     }
 
     public int getTransactionMaxSpans() {

apm-agent-core/src/main/java/co/elastic/apm/impl/ElasticApmTracer.java

@@ -6,6 +6,9 @@
 import co.elastic.apm.configuration.CoreConfiguration;
 import co.elastic.apm.configuration.PrefixingConfigurationSourceWrapper;
 import co.elastic.apm.impl.error.ErrorCapture;
+import co.elastic.apm.impl.sampling.ConstantSampler;
+import co.elastic.apm.impl.sampling.ProbabilitySampler;
+import co.elastic.apm.impl.sampling.Sampler;
 import co.elastic.apm.impl.stacktrace.Stacktrace;
 import co.elastic.apm.impl.stacktrace.StacktraceConfiguration;
 import co.elastic.apm.impl.stacktrace.StacktraceFactory;
@@ -21,13 +24,15 @@
 import com.blogspot.mydailyjava.weaklockfree.DetachedThreadLocal;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.stagemonitor.configuration.ConfigurationOption;
 import org.stagemonitor.configuration.ConfigurationOptionProvider;
 import org.stagemonitor.configuration.ConfigurationRegistry;
 import org.stagemonitor.configuration.source.EnvironmentVariableConfigurationSource;
 import org.stagemonitor.configuration.source.PropertyFileConfigurationSource;
 import org.stagemonitor.configuration.source.SimpleSource;
 import org.stagemonitor.configuration.source.SystemPropertyConfigurationSource;
 
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.util.ServiceLoader;
 import java.util.concurrent.TimeUnit;
@@ -56,6 +61,7 @@ public class ElasticApmTracer implements Tracer {
     private final CoreConfiguration coreConfiguration;
     private final Transaction noopTransaction;
     private final Span noopSpan;
+    private Sampler sampler;
 
     ElasticApmTracer(ConfigurationRegistry configurationRegistry, Reporter reporter, StacktraceFactory stacktraceFactory) {
         this.configurationRegistry = configurationRegistry;
@@ -91,9 +97,15 @@ public Stacktrace createInstance() {
             }
         });
         coreConfiguration = configurationRegistry.getConfig(CoreConfiguration.class);
-        noopTransaction = new Transaction().withName("noop").withType("noop").start(this, 0, false);
+        noopTransaction = new Transaction().withName("noop").withType("noop").start(this, 0, ConstantSampler.of(false));
         noopSpan = new Span().withName("noop").withType("noop").start(this, noopTransaction, null, 0, true);
-
+        sampler = ProbabilitySampler.of(coreConfiguration.getSampleRate().get());
+        coreConfiguration.getSampleRate().addChangeListener(new ConfigurationOption.ChangeListener<Double>() {
+            @Override
+            public void onChange(ConfigurationOption<?> configurationOption, Double oldValue, Double newValue) {
+                sampler = ProbabilitySampler.of(newValue);
+            }
+        });
     }
 
     public static Builder builder() {
@@ -121,13 +133,14 @@ ElasticApmTracer register() {
         return this;
     }
 
+    @Nonnull
     @Override
     public Transaction startTransaction() {
         Transaction transaction;
         if (!coreConfiguration.isActive()) {
             transaction = noopTransaction;
         } else {
-            transaction = transactionPool.createInstance().start(this, System.nanoTime(), true);
+            transaction = transactionPool.createInstance().start(this, System.nanoTime(), sampler);
         }
         currentTransaction.set(transaction);
         return transaction;
@@ -143,6 +156,7 @@ public Span currentSpan() {
         return currentSpan.get();
     }
 
+    @Nonnull
     @Override
     public Span startSpan() {
         Transaction transaction = currentTransaction();
@@ -167,7 +181,6 @@ private Span createRealSpan(Transaction transaction) {
             transaction.getSpanCount().getDropped().increment();
         } else {
             dropped = false;
-            transaction.addSpan(span);
         }
         span.start(this, transaction, currentSpan(), System.nanoTime(), dropped);
         return span;

apm-agent-core/src/main/java/co/elastic/apm/impl/sampling/ConstantSampler.java

@@ -0,0 +1,31 @@
+package co.elastic.apm.impl.sampling;
+
+import co.elastic.apm.impl.transaction.TransactionId;
+
+/**
+ * This is a implementation of {@link Sampler} which always returns the same sampling decision.
+ */
+public class ConstantSampler implements Sampler {
+
+    private static final Sampler TRUE = new ConstantSampler(true);
+    private static final Sampler FALSE = new ConstantSampler(false);
+
+    private final boolean decision;
+
+    private ConstantSampler(boolean decision) {
+        this.decision = decision;
+    }
+
+    public static Sampler of(boolean decision) {
+        if (decision) {
+            return TRUE;
+        } else {
+            return FALSE;
+        }
+    }
+
+    @Override
+    public boolean isSampled(TransactionId transactionId) {
+        return decision;
+    }
+}

apm-agent-core/src/main/java/co/elastic/apm/impl/sampling/ProbabilitySampler.java

@@ -0,0 +1,54 @@
+package co.elastic.apm.impl.sampling;
+
+import co.elastic.apm.impl.transaction.TransactionId;
+
+/**
+ * This implementation of {@link Sampler} samples based on a sampling probability (or sampling rate) between 0.0 and 1.0.
+ * <p>
+ * A sampling rate of 0.5 means that 50% of all transactions should be {@linkplain Sampler sampled}.
+ * </p>
+ * <p>
+ * Implementation notes:
+ * </p>
+ * <p>
+ * We are taking advantage of the fact, that the {@link TransactionId} is randomly generated.
+ * So instead of generating another random number,
+ * we just see if the long value returned by {@link TransactionId#getMostSignificantBits()}
+ * falls into the range between the <code>lowerBound</code> and the <code>higherBound</code>.
+ * This is a visual representation of the mechanism with a sampling rate of 0.5 (=50%):
+ * <pre>
+ * Long.MIN_VALUE        0                     Long.MAX_VALUE
+ * v                     v                     v
+ * [----------[----------|----------]----------]
+ *            ^                     ^
+ *            lowerBound            higherBound = Long.MAX_VALUE * samplingRate
+ *            = Long.MAX_VALUE * samplingRate * -1
+ * </pre>
+ * </p>
+ */
+public class ProbabilitySampler implements Sampler {
+
+    private final long lowerBound;
+    private final long higherBound;
+
+    public static Sampler of(double samplingRate) {
+        if (samplingRate == 1) {
+            return ConstantSampler.of(true);
+        }
+        if (samplingRate == 0) {
+            return ConstantSampler.of(false);
+        }
+        return new ProbabilitySampler(samplingRate);
+    }
+
+    private ProbabilitySampler(double samplingRate) {
+        higherBound = (long) (Long.MAX_VALUE * samplingRate);
+        lowerBound = -higherBound;
+    }
+
+    @Override
+    public boolean isSampled(TransactionId transactionId) {
+        final long mostSignificantBits = transactionId.getMostSignificantBits();
+        return mostSignificantBits > lowerBound && mostSignificantBits < higherBound;
+    }
+}

apm-agent-core/src/main/java/co/elastic/apm/impl/sampling/Sampler.java

@@ -0,0 +1,29 @@
+package co.elastic.apm.impl.sampling;
+
+import co.elastic.apm.api.Transaction;
+import co.elastic.apm.impl.transaction.TransactionId;
+
+/**
+ * A sampler is responsible for determining whether a {@link Transaction} should be sampled.
+ * <p>
+ * In contrast other tracing systems,
+ * in Elastic APM,
+ * non-sampled {@link Transaction}s do get reported to the APM server.
+ * However,
+ * to keep the size at a minimum,
+ * the reported {@link Transaction} only contains the transaction name,
+ * the duration and the id.
+ * Also,
+ * {@link co.elastic.apm.api.Span}s of non sampled {@link Transaction}s are not reported.
+ * </p>
+ */
+public interface Sampler {
+
+    /**
+     * Determines whether the given transaction should be sampled.
+     *
+     * @param transactionId The id of the transaction.
+     * @return The sampling decision.
+     */
+    boolean isSampled(TransactionId transactionId);
+}

apm-agent-core/src/main/java/co/elastic/apm/impl/transaction/Span.java

@@ -79,6 +79,7 @@ public Span start(ElasticApmTracer tracer, Transaction transaction, @Nullable Sp
         if (sampled) {
             start = (nanoTime - transaction.getDuration()) / MS_IN_NANOS;
             duration = nanoTime;
+            transaction.addSpan(this);
         }
         return this;
     }

apm-agent-core/src/main/java/co/elastic/apm/impl/transaction/Transaction.java

@@ -2,6 +2,7 @@
 
 import co.elastic.apm.impl.ElasticApmTracer;
 import co.elastic.apm.impl.context.Context;
+import co.elastic.apm.impl.sampling.Sampler;
 import co.elastic.apm.objectpool.Recyclable;
 import com.fasterxml.jackson.annotation.JsonFormat;
 import com.fasterxml.jackson.annotation.JsonInclude;
@@ -86,12 +87,12 @@ public class Transaction implements Recyclable, co.elastic.apm.api.Transaction {
     @JsonProperty("sampled")
     private boolean sampled;
 
-    public Transaction start(ElasticApmTracer tracer, long startTimestampNanos, boolean sampled) {
+    public Transaction start(ElasticApmTracer tracer, long startTimestampNanos, Sampler sampler) {
         this.tracer = tracer;
         this.duration = startTimestampNanos;
-        this.sampled = sampled;
         this.timestamp.setTime(System.currentTimeMillis());
         this.id.setToRandomValue();
+        this.sampled = sampler.isSampled(id);
         return this;
     }
 

apm-agent-core/src/main/java/co/elastic/apm/impl/transaction/TransactionId.java

@@ -6,6 +6,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Random;
 import java.util.UUID;
@@ -29,6 +30,10 @@ public void setToRandomValue(Random random) {
         random.nextBytes(data);
     }
 
+    public void setValue(long mostSignificantBits, long leastSignificantBits) {
+        ByteBuffer.wrap(data).putLong(mostSignificantBits).putLong(leastSignificantBits);
+    }
+
     @Override
     public void resetState() {
         for (int i = 0; i < data.length; i++) {

apm-agent-core/src/test/java/co/elastic/apm/configuration/SpyConfiguration.java

@@ -10,6 +10,8 @@
 
 public class SpyConfiguration {
 
+    public static final String CONFIG_SOURCE_NAME = "test config source";
+
     /**
      * Creates a configuration registry where all {@link ConfigurationOptionProvider}s are wrapped with
      * {@link org.mockito.Mockito#spy(Object)}
@@ -24,7 +26,7 @@ public static ConfigurationRegistry createSpyConfig() {
             builder.addOptionProvider(spy(options));
         }
         return builder
-            .addConfigSource(SimpleSource.forTest("service_name", "elastic-apm-test"))
+            .addConfigSource(new SimpleSource(CONFIG_SOURCE_NAME).add("service_name", "elastic-apm-test"))
             .build();
     }
 }

apm-agent-core/src/test/java/co/elastic/apm/impl/ElasticApmTracerTest.java

@@ -14,6 +14,8 @@
 import org.junit.jupiter.api.Test;
 import org.stagemonitor.configuration.ConfigurationRegistry;
 
+import java.io.IOException;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.when;
 
@@ -195,4 +197,17 @@ void testDisableMidTransaction() {
         assertThat(reporter.getFirstTransaction()).isSameAs(transaction);
 
     }
+
+    @Test
+    void testSamplingNone() throws IOException {
+        config.getConfig(CoreConfiguration.class).getSampleRate().update(0.0, SpyConfiguration.CONFIG_SOURCE_NAME);
+        try (Transaction transaction = tracerImpl.startTransaction()) {
+            transaction.setUser("1", "jon.doe@example.com", "jondoe");
+            try (Span span = tracerImpl.startSpan()) {
+            }
+        }
+        assertThat(reporter.getTransactions()).hasSize(1);
+        assertThat(reporter.getFirstTransaction().getSpans()).hasSize(0);
+        assertThat(reporter.getFirstTransaction().getContext().getUser().getEmail()).isNull();
+    }
 }

apm-agent-core/src/test/java/co/elastic/apm/impl/payload/TransactionPayloadJsonSchemaTest.java

@@ -1,6 +1,7 @@
 package co.elastic.apm.impl.payload;
 
 import co.elastic.apm.impl.ElasticApmTracer;
+import co.elastic.apm.impl.sampling.ConstantSampler;
 import co.elastic.apm.impl.transaction.Span;
 import co.elastic.apm.impl.transaction.Transaction;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -35,7 +36,7 @@ private TransactionPayload createPayloadWithRequiredValues() {
 
     private Transaction createTransactionWithRequiredValues() {
         Transaction t = new Transaction();
-        t.start(mock(ElasticApmTracer.class), 0, true);
+        t.start(mock(ElasticApmTracer.class), 0, ConstantSampler.of(true));
         t.setType("type");
         t.getContext().getRequest().withMethod("GET");
         Span s = new Span();