[Profiling] Manually downsample to 20k events if required (#131196)

* [Profiling] Manually downsample to 20k events if required

Since the random sampler aggregation is limited to p values of 0..0.5,
add an explicit downsampling for p values of 0.5..1.0.

With this, the latency of fetching stacktraces and stackframes can be
further reduced.

* Fix number of reported samples

Author: rockdaboot

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

@@ -32,6 +32,7 @@ public class GetStackTracesResponse extends ActionResponse implements ChunkedToX
     private final Map<TraceEventID, TraceEvent> stackTraceEvents;
     private final int totalFrames;
     private final double samplingRate;
+    private final double samplingFrequency;
     private final long totalSamples;
 
     public GetStackTracesResponse(
@@ -42,6 +43,28 @@ public GetStackTracesResponse(
         int totalFrames,
         double samplingRate,
         long totalSamples
+    ) {
+        this(
+            stackTraces,
+            stackFrames,
+            executables,
+            stackTraceEvents,
+            totalFrames,
+            samplingRate,
+            totalSamples,
+            TransportGetStackTracesAction.DEFAULT_SAMPLING_FREQUENCY
+        );
+    }
+
+    public GetStackTracesResponse(
+        Map<String, StackTrace> stackTraces,
+        Map<String, StackFrame> stackFrames,
+        Map<String, String> executables,
+        Map<TraceEventID, TraceEvent> stackTraceEvents,
+        int totalFrames,
+        double samplingRate,
+        long totalSamples,
+        double samplingFrequency
     ) {
         this.stackTraces = stackTraces;
         this.stackFrames = stackFrames;
@@ -50,6 +73,7 @@ public GetStackTracesResponse(
         this.totalFrames = totalFrames;
         this.samplingRate = samplingRate;
         this.totalSamples = totalSamples;
+        this.samplingFrequency = samplingFrequency;
     }
 
     @Override
@@ -101,7 +125,7 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
                     Iterators.map(v.entrySet().iterator(), e -> (b, p) -> b.field(e.getKey().stacktraceID(), e.getValue().count))
                 )
             ),
-            Iterators.single((b, p) -> b.field("sampling_rate", samplingRate).endObject())
+            Iterators.single((b, p) -> b.field("sampling_rate", samplingRate).field("sampling_frequency", samplingFrequency).endObject())
             // the following fields are intentionally not written to the XContent representation (only needed on the transport layer):
             //
             // * start
@@ -129,6 +153,7 @@ public boolean equals(Object o) {
         GetStackTracesResponse response = (GetStackTracesResponse) o;
         return totalFrames == response.totalFrames
             && samplingRate == response.samplingRate
+            && samplingFrequency == response.samplingFrequency
             && Objects.equals(stackTraces, response.stackTraces)
             && Objects.equals(stackFrames, response.stackFrames)
             && Objects.equals(executables, response.executables)
@@ -137,6 +162,6 @@ public boolean equals(Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(stackTraces, stackFrames, executables, stackTraceEvents, totalFrames, samplingRate);
+        return Objects.hash(stackTraces, stackFrames, executables, stackTraceEvents, totalFrames, samplingRate, samplingFrequency);
     }
 }

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

@@ -82,6 +82,14 @@ public double getSamplingRate() {
         return samplingRate;
     }
 
+    public void setSamplingFrequency(long samplingFrequency) {
+        this.samplingFrequency = samplingFrequency;
+    }
+
+    public long getSamplingFrequency() {
+        return samplingFrequency;
+    }
+
     public void setRequestedDuration(Double requestedDuration) {
         this.requestedDuration = requestedDuration;
     }
@@ -154,14 +162,15 @@ public GetStackTracesResponse build() {
                 }
             }
         }
-        return new GetStackTracesResponse(stackTraces, stackFrames, executables, stackTraceEvents, totalFrames, samplingRate, totalSamples);
-    }
-
-    public void setSamplingFrequency(long samplingFrequency) {
-        this.samplingFrequency = samplingFrequency;
-    }
-
-    public long getSamplingFrequency() {
-        return samplingFrequency;
+        return new GetStackTracesResponse(
+            stackTraces,
+            stackFrames,
+            executables,
+            stackTraceEvents,
+            totalFrames,
+            samplingRate,
+            totalSamples,
+            samplingFrequency
+        );
     }
 }

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

@@ -400,10 +400,10 @@ private void searchEventGroupedByStackTrace(
                 StopWatch watch = new StopWatch("createStackTraceEvents");
                 SingleBucketAggregation sample = searchResponse.getAggregations().get("sample");
                 InternalComposite stacktraces = sample.getAggregations().get("group_by");
+                RandomGenerator rng = new Random(rngSeed);
                 long indexSamplingFactor = Math.round(1 / eventsIndex.getSampleRate()); // for example, 5^2 = 25 for profiling-events-5pow02
                 int bucketCount = stacktraces.getBuckets().size();
                 long eventCount = sample.getDocCount();
-                AtomicLong upscaledEventCount = new AtomicLong(eventCount * indexSamplingFactor);
                 long maxSamplingFrequency = getAggValueAsLong(searchResponse, "max_freq") <= 0
                     ? (long) DEFAULT_SAMPLING_FREQUENCY
                     : getAggValueAsLong(searchResponse, "max_freq");
@@ -413,13 +413,35 @@ private void searchEventGroupedByStackTrace(
                     eventCount,
                     bucketCount,
                     indexSamplingFactor,
-                    upscaledEventCount.get()
+                    eventCount * indexSamplingFactor
                 );
 
+                // Since the random sampler aggregation does not support sampling rates between 0.5 and 1.0,
+                // we can have up to 2x more events in the response as requested by the user.
+                // In order to reduce latency for stacktrace and stackframe lookups, we add another sampling factor
+                // to reduce the number of events to match the user request (which reduces the number of unique stacktrace ids).
+                boolean needAdditionalDownsampling = eventCount > request.getSampleSize();
+                double downSamplingRate = needAdditionalDownsampling
+                    ? (double) request.getSampleSize() / eventCount
+                    : responseBuilder.getSamplingRate();
+
+                eventCount = 0;
                 boolean mixedFrequency = false;
                 Map<TraceEventID, TraceEvent> stackTraceEvents = new HashMap<>(bucketCount);
                 for (InternalComposite.InternalBucket stacktraceBucket : stacktraces.getBuckets()) {
-                    long count = stacktraceBucket.getDocCount() * indexSamplingFactor;
+                    long sampledCount;
+                    if (needAdditionalDownsampling) {
+                        sampledCount = downsampleEvents(rng, downSamplingRate, stacktraceBucket.getDocCount());
+                        if (sampledCount <= 0) {
+                            bucketCount--;
+                            continue; // skip bucket
+                        }
+                    } else {
+                        sampledCount = stacktraceBucket.getDocCount();
+                    }
+
+                    long count = roundWithRandom((sampledCount * indexSamplingFactor) / downSamplingRate, rng);
+                    eventCount += sampledCount;
 
                     TraceEventID eventID = getTraceEventID(stacktraceBucket);
                     stackTraceEvents.compute(eventID, (k, event) -> {
@@ -450,46 +472,52 @@ private void searchEventGroupedByStackTrace(
                     }
                 }
 
+                AtomicLong upscaledEventCount = new AtomicLong(eventCount * indexSamplingFactor);
                 if (mixedFrequency) {
-                    RandomGenerator r = new Random(rngSeed);
                     upscaledEventCount.set(0);
 
                     // Events have different frequencies.
-                    // Now upscale the count values to the max sampling frequency,
-                    // also taking into account the stratified downsampling factor (5, 25, 125, etc.).
+                    // Scale the count up to the maximum sampling frequency.
                     stackTraceEvents.forEach((eventID, event) -> {
-                        if (eventID.samplingFrequency() == maxSamplingFrequency) {
-                            upscaledEventCount.addAndGet(event.count);
-                            return; // no need to upscale
+                        if (eventID.samplingFrequency() != maxSamplingFrequency) {
+                            double samplingFactor = maxSamplingFrequency / eventID.samplingFrequency();
+                            event.count = roundWithRandom(event.count * samplingFactor, rng);
                         }
-
-                        // Use randomization, to avoid a systematic rounding issue that would happen
-                        // if we naively do `event.count = Math.round(event.count * samplingFactor)`.
-                        // For example, think of event.count = 1 and samplingFactor = 1.4: the naive approach would not change anything.
-                        double samplingFactor = maxSamplingFrequency / eventID.samplingFrequency();
-                        double newCount = event.count * samplingFactor;
-                        long integerPart = (long) newCount;
-                        double fractionalPart = newCount - integerPart;
-                        event.count = integerPart + (r.nextDouble() < fractionalPart ? 1 : 0);
                         upscaledEventCount.addAndGet(event.count);
                     });
                 }
 
                 log.debug(watch::report);
 
+                responseBuilder.setSamplingRate(downSamplingRate);
                 responseBuilder.setSamplingFrequency(maxSamplingFrequency);
                 responseBuilder.setTotalSamples(upscaledEventCount.get());
-                log.debug(
-                    "Found [{}] events in [{}] buckets, upscaled to [{}] events).",
-                    eventCount,
-                    bucketCount,
-                    upscaledEventCount.get()
-                );
+                log.debug("Use [{}] events in [{}] buckets, upscaled to [{}] events).", eventCount, bucketCount, upscaledEventCount.get());
 
                 return stackTraceEvents;
             }));
     }
 
+    private static long roundWithRandom(double value, RandomGenerator r) {
+        // Use randomization, to avoid a systematic rounding issue that would happen
+        // if we naively do `Math.round(value)`.
+        // For example, think of rounding value = 1.4: the naive approach would always drop 0.4 and return 1.
+        long integerPart = (long) value;
+        double fractionalPart = value - integerPart;
+        return integerPart + (r.nextDouble() < fractionalPart ? 1 : 0);
+    }
+
+    private static long downsampleEvents(RandomGenerator r, double samplingRate, long count) {
+        // Downsampling needs to be applied to each event individually.
+        long sampledCount = 0;
+        for (long i = 0; i < count; i++) {
+            if (r.nextDouble() < samplingRate) {
+                sampledCount++;
+            }
+        }
+        return sampledCount;
+    }
+
     private static TraceEventID getTraceEventID(InternalComposite.InternalBucket stacktraceBucket) {
         Map<String, Object> key = stacktraceBucket.getKey();
         Object samplingFrequency = key.get("freq");

x-pack/plugin/profiling/src/test/java/org/elasticsearch/xpack/profiling/action/GetStackTracesResponseTests.java

@@ -58,7 +58,7 @@ private GetStackTracesResponse createTestInstance() {
 
     public void testChunking() {
         AbstractChunkedSerializingTestCase.assertChunkCount(createTestInstance(), instance -> {
-            // start and {sampling_rate; end}; see GetStackTracesResponse.toXContentChunked()
+            // start and {sampling_rate; sampling_freq; end}; see GetStackTracesResponse.toXContentChunked()
             int chunks = 2;
             chunks += size(instance.getExecutables());
             chunks += size(instance.getStackFrames());