[Profiling] Add support for variable sampling frequency

x-pack/plugin/core/template-resources/src/main/resources/profiling/component-template/profiling-events.json

@@ -76,6 +76,10 @@
           "type": "short",
           "index": false
         },
+        "Stacktrace.sampling_frequency": {
+          "type": "long",
+          "index": false
+        },
         "agent.version": {
           "type": "keyword"
         },

x-pack/plugin/profiling/src/internalClusterTest/java/org/elasticsearch/xpack/profiling/action/GetStackTracesActionIT.java

@@ -40,7 +40,13 @@ public void testGetStackTracesUnfiltered() throws Exception {
 
         Map<TraceEventID, TraceEvent> traceEvents = response.getStackTraceEvents();
 
-        TraceEventID traceEventID = new TraceEventID("", "497295213074376", "8457605156473051743", "L7kj7UvlKbT-vN73el4faQ");
+        TraceEventID traceEventID = new TraceEventID(
+            "",
+            "497295213074376",
+            "8457605156473051743",
+            "L7kj7UvlKbT-vN73el4faQ",
+            TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY
+        );
         assertEquals(3L, response.getStackTraceEvents().get(traceEventID).count);
 
         assertNotNull(response.getStackTraces());
@@ -84,7 +90,13 @@ public void testGetStackTracesGroupedByServiceName() throws Exception {
 
         assertNotNull(response.getStackTraceEvents());
 
-        TraceEventID traceEventID = new TraceEventID("", "497295213074376", "8457605156473051743", "L7kj7UvlKbT-vN73el4faQ");
+        TraceEventID traceEventID = new TraceEventID(
+            "",
+            "497295213074376",
+            "8457605156473051743",
+            "L7kj7UvlKbT-vN73el4faQ",
+            TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY
+        );
         assertEquals(3L, response.getStackTraceEvents().get(traceEventID).count);
         assertEquals(Long.valueOf(2L), response.getStackTraceEvents().get(traceEventID).subGroups.getCount("basket"));
 
@@ -131,11 +143,17 @@ public void testGetStackTracesFromAPMWithMatchNoDownsampling() throws Exception
 
         assertNotNull(response.getStackTraceEvents());
 
-        TraceEventID traceEventID = new TraceEventID("", "", "", "Ce77w10WeIDow3kd1jowlA");
+        TraceEventID traceEventID = new TraceEventID(
+            "",
+            "",
+            "",
+            "Ce77w10WeIDow3kd1jowlA",
+            TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY
+        );
         assertEquals(3L, response.getStackTraceEvents().get(traceEventID).count);
         assertEquals(Long.valueOf(3L), response.getStackTraceEvents().get(traceEventID).subGroups.getCount("encodeSha1"));
 
-        traceEventID = new TraceEventID("", "", "", "JvISdnJ47BQ01489cwF9DA");
+        traceEventID = new TraceEventID("", "", "", "JvISdnJ47BQ01489cwF9DA", TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY);
         assertEquals(2L, response.getStackTraceEvents().get(traceEventID).count);
 
         assertNotNull(response.getStackTraces());
@@ -182,10 +200,16 @@ public void testGetStackTracesFromAPMWithMatchAndDownsampling() throws Exception
         assertNotNull(response.getStackTraceEvents());
 
         // as the sampling rate is 0.2, we see 5 times more samples (random sampler agg automatically adjusts sample count)
-        TraceEventID traceEventID = new TraceEventID("", "", "", "Ce77w10WeIDow3kd1jowlA");
+        TraceEventID traceEventID = new TraceEventID(
+            "",
+            "",
+            "",
+            "Ce77w10WeIDow3kd1jowlA",
+            TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY
+        );
         assertEquals(5 * 3L, response.getStackTraceEvents().get(traceEventID).count);
 
-        traceEventID = new TraceEventID("", "", "", "JvISdnJ47BQ01489cwF9DA");
+        traceEventID = new TraceEventID("", "", "", "JvISdnJ47BQ01489cwF9DA", TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY);
         assertEquals(5 * 2L, response.getStackTraceEvents().get(traceEventID).count);
 
         assertNotNull(response.getStackTraces());

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/CO2Calculator.java

@@ -12,7 +12,6 @@
 import java.util.Map;
 
 final class CO2Calculator {
-    private static final double DEFAULT_SAMPLING_FREQUENCY = 19.0d;
     private static final double DEFAULT_CO2_TONS_PER_KWH = 0.000379069d; // unit: metric tons / kWh
     private static final double DEFAULT_KILOWATTS_PER_CORE_X86 = 7.0d / 1000.0d; // unit: watt / core
     private static final double DEFAULT_KILOWATTS_PER_CORE_ARM64 = 2.8d / 1000.0d; // unit: watt / core
@@ -43,8 +42,8 @@ final class CO2Calculator {
             : customPerCoreWattARM64 / 1000.0d;
     }
 
-    public double getAnnualCO2Tons(String hostID, long samples) {
-        double annualCoreHours = CostCalculator.annualCoreHours(samplingDurationInSeconds, samples, DEFAULT_SAMPLING_FREQUENCY);
+    public double getAnnualCO2Tons(String hostID, long samples, double samplingFrequency) {
+        double annualCoreHours = CostCalculator.annualCoreHours(samplingDurationInSeconds, samples, samplingFrequency);
 
         HostMetadata host = hostMetadata.get(hostID);
         if (host == null) {

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/CostCalculator.java

@@ -10,7 +10,6 @@
 import java.util.Map;
 
 final class CostCalculator {
-    private static final double DEFAULT_SAMPLING_FREQUENCY = 19.0d;
     private static final double SECONDS_PER_HOUR = 60 * 60;
     private static final double SECONDS_PER_YEAR = SECONDS_PER_HOUR * 24 * 365.0d; // unit: seconds
     public static final double DEFAULT_COST_USD_PER_CORE_HOUR = 0.0425d; // unit: USD / (core * hour)
@@ -40,8 +39,8 @@ final class CostCalculator {
         );
     }
 
-    public double annualCostsUSD(String hostID, double samples) {
-        double annualCoreHours = annualCoreHours(samplingDurationInSeconds, samples, DEFAULT_SAMPLING_FREQUENCY);
+    public double annualCostsUSD(String hostID, double samples, double samplingFrequency) {
+        double annualCoreHours = annualCoreHours(samplingDurationInSeconds, samples, samplingFrequency);
 
         HostMetadata host = hostMetadata.get(hostID);
         if (host == null) {
@@ -59,7 +58,6 @@ public double annualCostsUSD(String hostID, double samples) {
     }
 
     public static double annualCoreHours(double duration, double samples, double samplingFrequency) {
-        // samplingFrequency will a variable value when we start supporting probabilistic profiling (soon).
         return (SECONDS_PER_YEAR / duration * samples / samplingFrequency) / SECONDS_PER_HOUR; // unit: core * hour
     }
 }

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/Resampler.java

@@ -14,7 +14,7 @@ class Resampler {
     private final boolean requiresResampling;
     private final RandomGenerator r;
     private final double adjustedSampleRate;
-    private final double p;
+    public final double p;
 
     Resampler(GetStackTracesRequest request, double sampleRate, long totalCount) {
         // Manually reduce sample count if totalCount exceeds sampleSize by 10%.
@@ -50,7 +50,7 @@ public int adjustSampleCount(int originalCount) {
         }
         // Adjust the sample counts from down-sampled to fully sampled.
         // Be aware that downsampling drops entries from stackTraceEvents, so that
-        // the sum of the upscaled count values is less that totalCount.
-        return (int) Math.floor(rawCount / (p * adjustedSampleRate));
+        // the sum of the upscaled count values is less than totalCount.
+        return (int) Math.round(rawCount / (p * adjustedSampleRate));
     }
 }

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/TraceEventID.java

@@ -7,4 +7,4 @@
 
 package org.elasticsearch.xpack.profiling.action;
 
-record TraceEventID(String executableName, String threadName, String hostID, String stacktraceID) {}
+record TraceEventID(String executableName, String threadName, String hostID, String stacktraceID, double samplingFrequency) {}

x-pack/plugin/profiling/src/main/java/org/elasticsearch/xpack/profiling/action/TransportGetStackTracesAction.java

@@ -115,6 +115,12 @@ public class TransportGetStackTracesAction extends TransportAction<GetStackTrace
      */
     private static final String CUSTOM_EVENT_SUB_AGGREGATION_NAME = "custom_event_group";
 
+    /**
+     * This is the default sampling rate for profiling events that we use if no sampling rate is
+     * stored in the backend (backwards compatibility).
+     */
+    public static final double DEFAULT_SAMPLING_FREQUENCY = 19.0d;
+
     private final NodeClient nodeClient;
     private final ProfilingLicenseChecker licenseChecker;
     private final ClusterService clusterService;
@@ -249,7 +255,6 @@ private void searchGenericEventGroupedByStackTrace(
         ActionListener<GetStackTracesResponse> submitListener,
         GetStackTracesResponseBuilder responseBuilder
     ) {
-
         CountedTermsAggregationBuilder groupByStackTraceId = new CountedTermsAggregationBuilder("group_by").size(
             MAX_TRACE_EVENTS_RESULT_SIZE
         ).field(request.getStackTraceIdsField());
@@ -286,7 +291,7 @@ private void searchGenericEventGroupedByStackTrace(
 
                     String stackTraceID = stacktraceBucket.getKeyAsString();
 
-                    TraceEventID eventID = new TraceEventID("", "", "", stackTraceID);
+                    TraceEventID eventID = new TraceEventID("", "", "", stackTraceID, DEFAULT_SAMPLING_FREQUENCY);
                     TraceEvent event = stackTraceEvents.computeIfAbsent(eventID, k -> new TraceEvent());
                     event.count += count;
                     subGroups.collectResults(stacktraceBucket, event);
@@ -337,6 +342,16 @@ private void searchEventGroupedByStackTrace(
             // Especially with high cardinality fields, this makes aggregations really slow.
             .executionHint("map")
             .subAggregation(groupByHostId);
+        TermsAggregationBuilder groupByExecutableName = new TermsAggregationBuilder("group_by")
+            // 'size' specifies the max number of host IDs we support per request.
+            .size(MAX_TRACE_EVENTS_RESULT_SIZE)
+            .field("process.executable.name")
+            // missing("") is used to include documents where the field is missing.
+            .missing("")
+            // 'execution_hint: map' skips the slow building of ordinals that we don't need.
+            // Especially with high cardinality fields, this makes aggregations really slow.
+            .executionHint("map")
+            .subAggregation(groupByThreadName);
         SubGroupCollector subGroups = SubGroupCollector.attach(groupByStackTraceId, request.getAggregationFields());
         client.prepareSearch(eventsIndex.getName())
             .setTrackTotalHits(false)
@@ -351,53 +366,89 @@ private void searchEventGroupedByStackTrace(
                 new TermsAggregationBuilder("group_by")
                     // 'size' specifies the max number of host ID we support per request.
                     .size(MAX_TRACE_EVENTS_RESULT_SIZE)
-                    .field("process.executable.name")
-                    // missing("") is used to include documents where the field is missing.
-                    .missing("")
+                    .field("Stacktrace.sampling_frequency")
+                    // missing(DEFAULT_SAMPLING_RATE) is used to include documents where the field is missing.
+                    .missing((long) DEFAULT_SAMPLING_FREQUENCY)
                     // 'execution_hint: map' skips the slow building of ordinals that we don't need.
                     // Especially with high cardinality fields, this makes aggregations really slow.
                     .executionHint("map")
-                    .subAggregation(groupByThreadName)
+                    .subAggregation(groupByExecutableName)
+                    .subAggregation(new SumAggregationBuilder("total_count").field("Stacktrace.count"))
             )
             .addAggregation(new SumAggregationBuilder("total_count").field("Stacktrace.count"))
             .execute(handleEventsGroupedByStackTrace(submitTask, client, responseBuilder, submitListener, searchResponse -> {
-                long totalCount = getAggValueAsLong(searchResponse, "total_count");
+                // The count values for events are scaled up to the highest sampling frequency.
+                // For example, if the highest sampling frequency is 100, an event with frequency=20 and count=1
+                // will be upscaled to count=5 (100/20 * count).
+                // For this, we need to find the highest frequency in the result set.
+                long maxSamplingFrequency = 0;
+                Terms samplingFrequencies = searchResponse.getAggregations().get("group_by");
+                for (Terms.Bucket samplingFrequencyBucket : samplingFrequencies.getBuckets()) {
+                    final double samplingFrequency = samplingFrequencyBucket.getKeyAsNumber().doubleValue();
+                    if (samplingFrequency > maxSamplingFrequency) {
+                        maxSamplingFrequency = (long) samplingFrequency;
+                    }
+                }
+
+                // Calculate a scaled-up total count (scaled up to the highest sampling frequency).
+                long totalCount = 0;
+                for (Terms.Bucket samplingFrequencyBucket : samplingFrequencies.getBuckets()) {
+                    InternalNumericMetricsAggregation.SingleValue count = samplingFrequencyBucket.getAggregations().get("total_count");
+                    final double samplingFrequency = samplingFrequencyBucket.getKeyAsNumber().doubleValue();
+                    final double samplingFactor = maxSamplingFrequency / samplingFrequency;
+                    totalCount += Math.round(count.value() * samplingFactor);
+                }
 
                 Resampler resampler = new Resampler(request, responseBuilder.getSamplingRate(), totalCount);
 
                 // Sort items lexicographically to access Lucene's term dictionary more efficiently when issuing an mget request.
-                // The term dictionary is lexicographically sorted and using the same order reduces the number of page faults
+                // The term dictionary is lexicographically sorted, and using the same order reduces the number of page faults
                 // needed to load it.
                 long totalFinalCount = 0;
                 Map<TraceEventID, TraceEvent> stackTraceEvents = new HashMap<>(MAX_TRACE_EVENTS_RESULT_SIZE);
 
-                Terms executableNames = searchResponse.getAggregations().get("group_by");
-                for (Terms.Bucket executableBucket : executableNames.getBuckets()) {
-                    String executableName = executableBucket.getKeyAsString();
-
-                    Terms threads = executableBucket.getAggregations().get("group_by");
-                    for (Terms.Bucket threadBucket : threads.getBuckets()) {
-                        String threadName = threadBucket.getKeyAsString();
-
-                        Terms hosts = threadBucket.getAggregations().get("group_by");
-                        for (Terms.Bucket hostBucket : hosts.getBuckets()) {
-                            String hostID = hostBucket.getKeyAsString();
-
-                            Terms stacktraces = hostBucket.getAggregations().get("group_by");
-                            for (Terms.Bucket stacktraceBucket : stacktraces.getBuckets()) {
-                                Sum count = stacktraceBucket.getAggregations().get("count");
-                                int finalCount = resampler.adjustSampleCount((int) count.value());
-                                if (finalCount <= 0) {
-                                    continue;
+                // Walk over all nested aggregations.
+                // The outermost aggregation is the sampling frequency.
+                // The next level is the executable name, followed by the thread name, host ID and stacktrace ID.
+                // the innermost aggregation contains the count of samples for each stacktrace ID.
+                for (Terms.Bucket samplingFrequencyBucket : samplingFrequencies.getBuckets()) {
+                    final double samplingFrequency = samplingFrequencyBucket.getKeyAsNumber().doubleValue();
+                    final double samplingFactor = maxSamplingFrequency / samplingFrequency;
+
+                    Terms executableNames = samplingFrequencyBucket.getAggregations().get("group_by");
+                    for (Terms.Bucket executableBucket : executableNames.getBuckets()) {
+                        String executableName = executableBucket.getKeyAsString();
+
+                        Terms threads = executableBucket.getAggregations().get("group_by");
+                        for (Terms.Bucket threadBucket : threads.getBuckets()) {
+                            String threadName = threadBucket.getKeyAsString();
+
+                            Terms hosts = threadBucket.getAggregations().get("group_by");
+                            for (Terms.Bucket hostBucket : hosts.getBuckets()) {
+                                String hostID = hostBucket.getKeyAsString();
+
+                                Terms stacktraces = hostBucket.getAggregations().get("group_by");
+                                for (Terms.Bucket stacktraceBucket : stacktraces.getBuckets()) {
+                                    Sum count = stacktraceBucket.getAggregations().get("count");
+                                    int finalCount = resampler.adjustSampleCount((int) Math.round(count.value() * samplingFactor));
+                                    if (finalCount <= 0) {
+                                        continue;
+                                    }
+
+                                    totalFinalCount += finalCount;
+
+                                    String stackTraceID = stacktraceBucket.getKeyAsString();
+                                    TraceEventID eventID = new TraceEventID(
+                                        executableName,
+                                        threadName,
+                                        hostID,
+                                        stackTraceID,
+                                        maxSamplingFrequency
+                                    );
+                                    TraceEvent event = stackTraceEvents.computeIfAbsent(eventID, k -> new TraceEvent());
+                                    event.count += finalCount;
+                                    subGroups.collectResults(stacktraceBucket, event);
                                 }
-                                totalFinalCount += finalCount;
-
-                                String stackTraceID = stacktraceBucket.getKeyAsString();
-
-                                TraceEventID eventID = new TraceEventID(executableName, threadName, hostID, stackTraceID);
-                                TraceEvent event = stackTraceEvents.computeIfAbsent(eventID, k -> new TraceEvent());
-                                event.count += finalCount;
-                                subGroups.collectResults(stacktraceBucket, event);
                             }
                         }
                     }
@@ -629,8 +680,8 @@ public void calculateCO2AndCosts() {
             );
 
             responseBuilder.getStackTraceEvents().forEach((eventId, event) -> {
-                event.annualCO2Tons += co2Calculator.getAnnualCO2Tons(eventId.hostID(), event.count);
-                event.annualCostsUSD += costCalculator.annualCostsUSD(eventId.hostID(), event.count);
+                event.annualCO2Tons += co2Calculator.getAnnualCO2Tons(eventId.hostID(), event.count, eventId.samplingFrequency());
+                event.annualCostsUSD += costCalculator.annualCostsUSD(eventId.hostID(), event.count, eventId.samplingFrequency());
             });
 
             log.debug(watch::report);