Add stackframe metadata to TopN stack trace response (#64)

* Refactor grouping metadata by stacktraces

* Collate stacktrace events same as flamegraph

* Transform raw aggregation result

We now use the same TopN histogram bucket format used by prodfiler.

* Return no hits when querying stacktrace events

* Simplify topN query for stacktrace events

* Reduce query time by avoiding unneeded searches

* Refactor out query for stacktraces

* Refactor out queries for frames and executables

* Convert Elastic response to stackframe metadata

* Search for frames and executables in parallel

* Add default value for contexts

* Add types for function parameters

* Fix compiler warnings and failing tests

src/core/server/elasticsearch/client/configure_client.ts

@@ -8,6 +8,9 @@
 
 import { Client, HttpConnection } from '@elastic/elasticsearch';
 import type { Logger } from '@kbn/logging';
+import type { KibanaClient } from '@elastic/elasticsearch/lib/api/kibana';
+import agent, { Span } from 'elastic-apm-node';
+import LRUCache from 'lru-cache';
 import { parseClientOptions, ElasticsearchClientConfig } from './client_config';
 import { instrumentEsQueryAndDeprecationLogger } from './log_query_and_deprecation';
 import { createTransport } from './create_transport';
@@ -31,12 +34,45 @@ export const configureClient = (
   const clientOptions = parseClientOptions(config, scoped);
   const KibanaTransport = createTransport({ getExecutionContext });
 
+  const cache = new LRUCache<any, Span | undefined | null>({
+    max: 100,
+  });
+
   const client = new Client({
     ...clientOptions,
     Transport: KibanaTransport,
     Connection: HttpConnection,
   });
 
+  function startSpan(name: string) {
+    const span = agent.startSpan(name, 'db', 'elasticsearch', { exitSpan: true });
+    return span;
+  }
+
+  client.diagnostic.on('serialization', (err, result) => {
+    if (!err) {
+      cache.set(result?.meta.request.id, startSpan('serialization'));
+    }
+  });
+
+  client.diagnostic.on('request', (err, result) => {
+    cache.get(result?.meta.request.id)?.end();
+    if (!err) {
+      cache.set(result?.meta.request.id, startSpan('request'));
+    }
+  });
+
+  client.diagnostic.on('deserialization', (err, result) => {
+    cache.get(result?.requestId)?.end();
+    if (!err) {
+      cache.set(result?.requestId, startSpan('deserialization'));
+    }
+  });
+
+  client.diagnostic.on('response', (err, result) => {
+    cache.get(result?.meta.request.id)?.end();
+  });
+
   instrumentEsQueryAndDeprecationLogger({ logger, client, type });
 
   return client;

src/plugins/profiling/common/flamegraph.ts

@@ -19,7 +19,7 @@ import {
   StackFrame,
   Executable,
   createStackFrameMetadata,
-  StackFrameMetadata,
+  groupStackFrameMetadataByStackTrace,
 } from './profiling';
 
 interface PixiFlameGraph extends CallerCalleeNode {
@@ -75,36 +75,6 @@ export class FlameGraph {
     this.logger = logger;
   }
 
-  // getFrameMetadataForTraces collects all of the per-stack-frame metadata for a
-  // given set of trace IDs and their respective stack frames.
-  //
-  // This is similar to GetTraceMetaData in pf-storage-backend/storagebackend/storagebackendv1/reads_webservice.go
-  private getFrameMetadataForTraces(): Map<StackTraceID, StackFrameMetadata[]> {
-    const frameMetadataForTraces = new Map<StackTraceID, StackFrameMetadata[]>();
-    for (const [stackTraceID, trace] of this.stacktraces) {
-      const frameMetadata = new Array<StackFrameMetadata>();
-      for (let i = 0; i < trace.FrameID.length; i++) {
-        const frame = this.stackframes.get(trace.FrameID[i])!;
-        const executable = this.executables.get(trace.FileID[i])!;
-
-        const metadata = createStackFrameMetadata({
-          FileID: Buffer.from(trace.FileID[i], 'base64url').toString('hex'),
-          FrameType: trace.Type[i],
-          AddressOrLine: frame.LineNumber,
-          FunctionName: frame.FunctionName,
-          FunctionOffset: frame.FunctionOffset,
-          SourceLine: frame.LineNumber,
-          ExeFileName: executable.FileName,
-          Index: i,
-        });
-
-        frameMetadata.push(metadata);
-      }
-      frameMetadataForTraces.set(stackTraceID, frameMetadata);
-    }
-    return frameMetadataForTraces;
-  }
-
   private getExeFileName(exe: any, type: number) {
     if (exe?.FileName === undefined) {
       this.logger.warn('missing executable FileName');
@@ -188,7 +158,11 @@ export class FlameGraph {
 
   toPixi(): PixiFlameGraph {
     const rootFrame = createStackFrameMetadata();
-    const frameMetadataForTraces = this.getFrameMetadataForTraces();
+    const frameMetadataForTraces = groupStackFrameMetadataByStackTrace(
+      this.stacktraces,
+      this.stackframes,
+      this.executables
+    );
     const diagram = createCallerCalleeIntermediateRoot(
       rootFrame,
       this.events,

src/plugins/profiling/common/index.ts

@@ -28,20 +28,10 @@ function toMilliseconds(seconds: string): number {
   return parseInt(seconds, 10) * 1000;
 }
 
-export function getTopN(obj) {
+export function getTopN(obj: any) {
   const data = [];
 
-  if (obj.topN?.histogram?.buckets!) {
-    // needed for data served from Elasticsearch
-    for (let i = 0; i < obj.topN.histogram.buckets.length; i++) {
-      const bucket = obj.topN.histogram.buckets[i];
-      for (let j = 0; j < bucket.group_by.buckets.length; j++) {
-        const v = bucket.group_by.buckets[j];
-        data.push({ x: bucket.key, y: v.Count.value, g: v.key });
-      }
-    }
-  } else if (obj.TopN!) {
-    // needed for data served from fixtures
+  if (obj.TopN!) {
     for (const x in obj.TopN) {
       if (obj.TopN.hasOwnProperty(x)) {
         const values = obj.TopN[x];
@@ -56,7 +46,7 @@ export function getTopN(obj) {
   return data;
 }
 
-export function groupSamplesByCategory(samples) {
+export function groupSamplesByCategory(samples: any) {
   const series = new Map();
   for (let i = 0; i < samples.length; i++) {
     const v = samples[i];
@@ -74,3 +64,15 @@ export function timeRangeFromRequest(request: any): [number, number] {
   const timeTo = parseInt(request.query.timeTo!, 10);
   return [timeFrom, timeTo];
 }
+
+// Converts from a Map object to a Record object since Map objects are not
+// serializable to JSON by default
+export function fromMapToRecord<K extends string, V>(m: Map<K, V>): Record<string, V> {
+  let output: Record<string, V> = {};
+
+  for (const [key, value] of m) {
+    output[key] = value;
+  }
+
+  return output;
+}

src/plugins/profiling/common/profiling.ts

@@ -96,6 +96,40 @@ export function createStackFrameMetadata(
   return metadata;
 }
 
+// groupStackFrameMetadataByStackTrace collects all of the per-stack-frame
+// metadata for a given set of trace IDs and their respective stack frames.
+//
+// This is similar to GetTraceMetaData in pf-storage-backend/storagebackend/storagebackendv1/reads_webservice.go
+export function groupStackFrameMetadataByStackTrace(
+  stackTraces: Map<StackTraceID, StackTrace>,
+  stackFrames: Map<StackFrameID, StackFrame>,
+  executables: Map<FileID, Executable>
+): Map<StackTraceID, StackFrameMetadata[]> {
+  const frameMetadataForTraces = new Map<StackTraceID, StackFrameMetadata[]>();
+  for (const [stackTraceID, trace] of stackTraces) {
+    const frameMetadata = new Array<StackFrameMetadata>();
+    for (let i = 0; i < trace.FrameID.length; i++) {
+      const frame = stackFrames.get(trace.FrameID[i])!;
+      const executable = executables.get(trace.FileID[i])!;
+
+      const metadata = createStackFrameMetadata({
+        FileID: Buffer.from(trace.FileID[i], 'base64url').toString('hex'),
+        FrameType: trace.Type[i],
+        AddressOrLine: frame.LineNumber,
+        FunctionName: frame.FunctionName,
+        FunctionOffset: frame.FunctionOffset,
+        SourceLine: frame.LineNumber,
+        ExeFileName: executable.FileName,
+        Index: i,
+      });
+
+      frameMetadata.push(metadata);
+    }
+    frameMetadataForTraces.set(stackTraceID, frameMetadata);
+  }
+  return frameMetadataForTraces;
+}
+
 export type FrameGroup = Pick<
   StackFrameMetadata,
   'FileID' | 'ExeFileName' | 'FunctionName' | 'AddressOrLine' | 'SourceFilename'

src/plugins/profiling/common/topn.ts

@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import {
+  AggregationsHistogramAggregate,
+  AggregationsHistogramBucket,
+} from '@elastic/elasticsearch/lib/api/types';
+
+import { StackFrameMetadata } from './profiling';
+
+type TopNBucket = {
+  Value: string;
+  Count: number;
+};
+
+type TopNBucketsByDate = {
+  TopN: Record<number, TopNBucket[]>;
+};
+
+type TopNContainers = TopNBucketsByDate;
+type TopNDeployments = TopNBucketsByDate;
+type TopNHosts = TopNBucketsByDate;
+type TopNThreads = TopNBucketsByDate;
+
+type TopNTraces = TopNBucketsByDate & {
+  Metadata: Record<string, StackFrameMetadata[]>;
+};
+
+type TopN = TopNContainers | TopNDeployments | TopNHosts | TopNThreads | TopNTraces;
+
+export function createTopNBucketsByDate(
+  histogram: AggregationsHistogramAggregate
+): TopNBucketsByDate {
+  const topNBucketsByDate: Record<number, TopNBucket[]> = {};
+
+  const histogramBuckets = (histogram?.buckets as AggregationsHistogramBucket[]) ?? [];
+  for (let i = 0; i < histogramBuckets.length; i++) {
+    const key = histogramBuckets[i].key / 1000;
+    topNBucketsByDate[key] = [];
+    histogramBuckets[i].group_by.buckets.forEach((item: any) => {
+      topNBucketsByDate[key].push({
+        Value: item.key,
+        Count: item.count.value,
+      });
+    });
+  }
+
+  return { TopN: topNBucketsByDate };
+}

src/plugins/profiling/public/components/PixiFlamechart/FlameGraph/helpers/BitmapTextEllipse.ts

@@ -95,4 +95,4 @@ export class BitmapTextEllipse extends Pixi.BitmapText {
      */
     this.dirty = true
   }
-}
+}

src/plugins/profiling/public/components/PixiFlamechart/FlameGraph/helpers/binarySearchLowerLimit.ts

@@ -30,4 +30,4 @@ export const binarySearchLowerLimit = (
   }
 
   return min;
-}
+}

src/plugins/profiling/public/components/PixiFlamechart/FlameGraph/helpers/debounce.ts

@@ -8,4 +8,4 @@ const debounce = (callback: Function, wait: number) => {
   };
 }
 
-export default debounce
+export default debounce

src/plugins/profiling/public/components/PixiFlamechart/FlameGraph/helpers/lifecycle.ts

@@ -151,4 +151,4 @@ export const useResizeListenerEffect = (
       callback(canvas, renderer, viewport)
     }, 500)
   }, [gameCanvasRef, sidebar, renderer, viewport, callback])
-}
+}

src/plugins/profiling/public/components/PixiFlamechart/FlameGraph/helpers/regexCreator.ts

@@ -6,4 +6,4 @@ export const safeRegexCreator = (pattern: string, flags?: string | undefined): R
     // in the future this could be handled by the UI and we could show a message to the user
     return null
   }
-}
+}

src/plugins/profiling/public/components/PixiFlamechart/FlameGraph/services/featureFlags.ts

@@ -15,4 +15,4 @@ export const toggleSandboxMode = () => {
   } else {
     setSandboxModeTo(ENABLED_KEY)
   }
-}
+}

src/plugins/profiling/public/components/chart-grid.tsx

@@ -26,8 +26,8 @@ export interface ChartGridProps {
 
 export const ChartGrid: React.FC<ChartGridProps> = ({ maximum }) => {
   const ctx = useContext(TopNContext);
-  const printSubCharts = (series) => {
-    let keys = Array.from(series.keys());
+  const printSubCharts = (series: any) => {
+    let keys: string[] = Array.from(series.keys());
     const ncharts = Math.min(maximum, series.size);
     keys = keys.slice(0, ncharts);
 

src/plugins/profiling/public/components/contexts/flamegraph.tsx

@@ -8,4 +8,4 @@
 
 import { createContext } from 'react';
 
-export const FlameGraphContext = createContext();
+export const FlameGraphContext = createContext({});

src/plugins/profiling/public/components/contexts/topn.tsx

@@ -8,4 +8,4 @@
 
 import { createContext } from 'react';
 
-export const TopNContext = createContext();
+export const TopNContext = createContext({});

src/plugins/profiling/public/components/flamegraph.tsx

@@ -32,7 +32,7 @@ export const FlameGraph: React.FC<FlameGraphProps> = ({ id, height }) => {
     }
 
     const { leaves } = ctx;
-    const maxDepth = Math.max(...leaves.map((node) => node.depth));
+    const maxDepth = Math.max(...leaves.map((node: any) => node.depth));
 
     const result = [...new Array(maxDepth)].map((_, depth) => {
       return {

src/plugins/profiling/public/components/stacktrace-nav.tsx

@@ -45,7 +45,7 @@ export const StackTraceNavigation = ({ fetchTopN, setTopN }) => {
 
   const [toggleTopNSelected, setToggleTopNSelected] = useState(`${topnButtonGroupPrefix}__0`);
 
-  const onTopNChange = (optionId) => {
+  const onTopNChange = (optionId: string) => {
     if (optionId === toggleTopNSelected) {
       return;
     }
@@ -84,7 +84,7 @@ export const StackTraceNavigation = ({ fetchTopN, setTopN }) => {
 
   const [toggleDateSelected, setToggleDateSelected] = useState(`${dateButtonGroupPrefix}__0`);
 
-  const onDateChange = (optionId) => {
+  const onDateChange = (optionId: string) => {
     if (optionId === toggleDateSelected) {
       return;
     }

src/plugins/profiling/server/routes/flamechart.test.ts

@@ -7,7 +7,7 @@
  */
 
 import { DownsampledEventsIndex, getSampledTraceEventsIndex } from './downsampling';
-import { extractFileIDFromFrameID, parallelMget } from './flamechart';
+import { parallelMget } from './flamechart';
 import { ElasticsearchClient } from 'kibana/server';
 
 describe('Using down-sampled indexes', () => {
@@ -63,28 +63,6 @@ describe('Using down-sampled indexes', () => {
   });
 });
 
-describe('Extract FileID from FrameID', () => {
-  test('extractFileIDFromFrameID', () => {
-    const tests: Array<{
-      frameID: string;
-      expected: string;
-    }> = [
-      {
-        frameID: 'aQpJmTLWydNvOapSFZOwKgAAAAAAB924',
-        expected: 'aQpJmTLWydNvOapSFZOwKg==',
-      },
-      {
-        frameID: 'hz_u-HGyrN6qeIk6UIJeCAAAAAAAAAZZ',
-        expected: 'hz_u-HGyrN6qeIk6UIJeCA==',
-      },
-    ];
-
-    for (const t of tests) {
-      expect(extractFileIDFromFrameID(t.frameID)).toEqual(t.expected);
-    }
-  });
-});
-
 describe('Calling mget from events to stacktraces', () => {
   test('parallel queries to ES are resolved as promises', async () => {
     const numberOfFrames = 4;