[APM] Add custom spans around async operations

package.json

@@ -121,6 +121,7 @@
     "@kbn/ace": "link:packages/kbn-ace",
     "@kbn/analytics": "link:packages/kbn-analytics",
     "@kbn/apm-config-loader": "link:packages/kbn-apm-config-loader",
+    "@kbn/apm-utils": "link:packages/kbn-apm-utils",
     "@kbn/config": "link:packages/kbn-config",
     "@kbn/config-schema": "link:packages/kbn-config-schema",
     "@kbn/i18n": "link:packages/kbn-i18n",
@@ -129,6 +130,7 @@
     "@kbn/logging": "link:packages/kbn-logging",
     "@kbn/monaco": "link:packages/kbn-monaco",
     "@kbn/std": "link:packages/kbn-std",
+    "@kbn/tinymath": "link:packages/kbn-tinymath",
     "@kbn/ui-framework": "link:packages/kbn-ui-framework",
     "@kbn/ui-shared-deps": "link:packages/kbn-ui-shared-deps",
     "@kbn/utils": "link:packages/kbn-utils",
@@ -312,7 +314,6 @@
     "tabbable": "1.1.3",
     "tar": "4.4.13",
     "tinygradient": "0.4.3",
-    "@kbn/tinymath": "link:packages/kbn-tinymath",
     "tree-kill": "^1.2.2",
     "ts-easing": "^0.2.0",
     "tslib": "^2.0.0",
@@ -390,10 +391,10 @@
     "@scant/router": "^0.1.0",
     "@storybook/addon-a11y": "^6.0.26",
     "@storybook/addon-actions": "^6.0.26",
+    "@storybook/addon-docs": "^6.0.26",
     "@storybook/addon-essentials": "^6.0.26",
     "@storybook/addon-knobs": "^6.0.26",
     "@storybook/addon-storyshots": "^6.0.26",
-    "@storybook/addon-docs": "^6.0.26",
     "@storybook/components": "^6.0.26",
     "@storybook/core": "^6.0.26",
     "@storybook/core-events": "^6.0.26",
@@ -851,4 +852,4 @@
     "yargs": "^15.4.1",
     "zlib": "^1.0.5"
   }
-}
+}

packages/kbn-apm-utils/package.json

@@ -0,0 +1,13 @@
+{
+  "name": "@kbn/apm-utils",
+  "main": "./target/index.js",
+  "types": "./target/index.d.ts",
+  "version": "1.0.0",
+  "license": "SSPL-1.0 OR Elastic License 2.0",
+  "private": true,
+  "scripts": {
+    "build": "../../node_modules/.bin/tsc",
+    "kbn:bootstrap": "yarn build",
+    "kbn:watch": "yarn build --watch"
+  }
+}

packages/kbn-apm-utils/src/index.ts

@@ -0,0 +1,87 @@
+/*
+ * 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 agent from 'elastic-apm-node';
+import asyncHooks from 'async_hooks';
+
+export interface SpanOptions {
+  name: string;
+  type?: string;
+  subtype?: string;
+  labels?: Record<string, string>;
+}
+
+export function parseSpanOptions(optionsOrName: SpanOptions | string) {
+  const options = typeof optionsOrName === 'string' ? { name: optionsOrName } : optionsOrName;
+
+  return options;
+}
+
+const runInNewContext = <T extends (...args: any[]) => any>(cb: T): ReturnType<T> => {
+  const resource = new asyncHooks.AsyncResource('fake_async');
+
+  return resource.runInAsyncScope(cb);
+};
+
+export async function withSpan<T>(
+  optionsOrName: SpanOptions | string,
+  cb: () => Promise<T>
+): Promise<T> {
+  const options = parseSpanOptions(optionsOrName);
+
+  const { name, type, subtype, labels } = options;
+
+  if (!agent.isStarted()) {
+    return cb();
+  }
+
+  // When a span starts, it's marked as the active span in its context.
+  // When it ends, it's not untracked, which means that if a span
+  // starts directly after this one ends, the newly started span is a
+  // child of this span, even though it should be a sibling.
+  // To mitigate this, we queue a microtask by awaiting a promise.
+  await Promise.resolve();
+
+  const span = agent.startSpan(name);
+
+  if (!span) {
+    return cb();
+  }
+
+  // If a span is created in the same context as the span that we just
+  // started, it will be a sibling, not a child. E.g., the Elasticsearch span
+  // that is created when calling search() happens in the same context. To
+  // mitigate this we create a new context.
+
+  return runInNewContext(() => {
+    // @ts-ignore
+    if (type) {
+      span.type = type;
+    }
+    if (subtype) {
+      span.subtype = subtype;
+    }
+
+    if (labels) {
+      span.addLabels(labels);
+    }
+
+    return cb()
+      .then((res) => {
+        span.outcome = 'success';
+        return res;
+      })
+      .catch((err) => {
+        span.outcome = 'failure';
+        throw err;
+      })
+      .finally(() => {
+        span.end();
+      });
+  });
+}

packages/kbn-apm-utils/tsconfig.json

@@ -0,0 +1,18 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "declaration": true,
+    "outDir": "./target",
+    "stripInternal": false,
+    "declarationMap": true,
+    "types": [
+      "node"
+    ]
+  },
+  "include": [
+    "./src/**/*.ts"
+  ],
+  "exclude": [
+    "target"
+  ]
+}

x-pack/plugins/apm/server/lib/alerts/chart_preview/get_transaction_duration.ts

@@ -15,80 +15,83 @@ import {
 import { ProcessorEvent } from '../../../../common/processor_event';
 import { rangeFilter } from '../../../../common/utils/range_filter';
 import { AlertParams } from '../../../routes/alerts/chart_preview';
+import { withApmSpan } from '../../../utils/with_apm_span';
 import { getEnvironmentUiFilterES } from '../../helpers/convert_ui_filters/get_environment_ui_filter_es';
 import { getBucketSize } from '../../helpers/get_bucket_size';
 import { Setup, SetupTimeRange } from '../../helpers/setup_request';
 
-export async function getTransactionDurationChartPreview({
+export function getTransactionDurationChartPreview({
   alertParams,
   setup,
 }: {
   alertParams: AlertParams;
   setup: Setup & SetupTimeRange;
 }) {
-  const { apmEventClient, start, end } = setup;
-  const {
-    aggregationType,
-    environment,
-    serviceName,
-    transactionType,
-  } = alertParams;
+  return withApmSpan('get_transaction_duration_chart_preview', async () => {
+    const { apmEventClient, start, end } = setup;
+    const {
+      aggregationType,
+      environment,
+      serviceName,
+      transactionType,
+    } = alertParams;
 
-  const query = {
-    bool: {
-      filter: [
-        { range: rangeFilter(start, end) },
-        { term: { [PROCESSOR_EVENT]: ProcessorEvent.transaction } },
-        ...(serviceName ? [{ term: { [SERVICE_NAME]: serviceName } }] : []),
-        ...(transactionType
-          ? [{ term: { [TRANSACTION_TYPE]: transactionType } }]
-          : []),
-        ...getEnvironmentUiFilterES(environment),
-      ],
-    },
-  };
+    const query = {
+      bool: {
+        filter: [
+          { range: rangeFilter(start, end) },
+          { term: { [PROCESSOR_EVENT]: ProcessorEvent.transaction } },
+          ...(serviceName ? [{ term: { [SERVICE_NAME]: serviceName } }] : []),
+          ...(transactionType
+            ? [{ term: { [TRANSACTION_TYPE]: transactionType } }]
+            : []),
+          ...getEnvironmentUiFilterES(environment),
+        ],
+      },
+    };
 
-  const { intervalString } = getBucketSize({ start, end, numBuckets: 20 });
+    const { intervalString } = getBucketSize({ start, end, numBuckets: 20 });
 
-  const aggs = {
-    timeseries: {
-      date_histogram: {
-        field: '@timestamp',
-        fixed_interval: intervalString,
-      },
-      aggs: {
-        agg:
-          aggregationType === 'avg'
-            ? { avg: { field: TRANSACTION_DURATION } }
-            : {
-                percentiles: {
-                  field: TRANSACTION_DURATION,
-                  percents: [aggregationType === '95th' ? 95 : 99],
+    const aggs = {
+      timeseries: {
+        date_histogram: {
+          field: '@timestamp',
+          fixed_interval: intervalString,
+        },
+        aggs: {
+          agg:
+            aggregationType === 'avg'
+              ? { avg: { field: TRANSACTION_DURATION } }
+              : {
+                  percentiles: {
+                    field: TRANSACTION_DURATION,
+                    percents: [aggregationType === '95th' ? 95 : 99],
+                  },
                 },
-              },
+        },
       },
-    },
-  };
-  const params = {
-    apm: { events: [ProcessorEvent.transaction] },
-    body: { size: 0, query, aggs },
-  };
-  const resp = await apmEventClient.search(params);
+    };
+    const params = {
+      apm: { events: [ProcessorEvent.transaction] },
+      body: { size: 0, query, aggs },
+    };
+    const resp = await apmEventClient.search(params);
 
-  if (!resp.aggregations) {
-    return [];
-  }
+    if (!resp.aggregations) {
+      return [];
+    }
 
-  return resp.aggregations.timeseries.buckets.map((bucket) => {
-    const percentilesKey = aggregationType === '95th' ? '95.0' : '99.0';
-    const x = bucket.key;
-    const y =
-      aggregationType === 'avg'
-        ? (bucket.agg as MetricsAggregationResponsePart).value
-        : (bucket.agg as { values: Record<string, number | null> }).values[
-            percentilesKey
-          ];
+    return resp.aggregations.timeseries.buckets.map((bucket) => {
+      const percentilesKey = aggregationType === '95th' ? '95.0' : '99.0';
+      const x = bucket.key;
+      const y =
+        aggregationType === 'avg'
+          ? (bucket.agg as MetricsAggregationResponsePart).value
+          : (bucket.agg as { values: Record<string, number | null> }).values[
+              percentilesKey
+            ];
 
-    return { x, y };
+      return { x, y };
+    });
   });
 }

x-pack/plugins/apm/server/lib/alerts/chart_preview/get_transaction_error_count.ts

@@ -9,56 +9,59 @@ import { SERVICE_NAME } from '../../../../common/elasticsearch_fieldnames';
 import { ProcessorEvent } from '../../../../common/processor_event';
 import { rangeFilter } from '../../../../common/utils/range_filter';
 import { AlertParams } from '../../../routes/alerts/chart_preview';
+import { withApmSpan } from '../../../utils/with_apm_span';
 import { getEnvironmentUiFilterES } from '../../helpers/convert_ui_filters/get_environment_ui_filter_es';
 import { getBucketSize } from '../../helpers/get_bucket_size';
 import { Setup, SetupTimeRange } from '../../helpers/setup_request';
 
-export async function getTransactionErrorCountChartPreview({
+export function getTransactionErrorCountChartPreview({
   setup,
   alertParams,
 }: {
   setup: Setup & SetupTimeRange;
   alertParams: AlertParams;
 }) {
-  const { apmEventClient, start, end } = setup;
-  const { serviceName, environment } = alertParams;
+  return withApmSpan('get_transaction_error_count_chart_preview', async () => {
+    const { apmEventClient, start, end } = setup;
+    const { serviceName, environment } = alertParams;
 
-  const query = {
-    bool: {
-      filter: [
-        { range: rangeFilter(start, end) },
-        ...(serviceName ? [{ term: { [SERVICE_NAME]: serviceName } }] : []),
-        ...getEnvironmentUiFilterES(environment),
-      ],
-    },
-  };
+    const query = {
+      bool: {
+        filter: [
+          { range: rangeFilter(start, end) },
+          ...(serviceName ? [{ term: { [SERVICE_NAME]: serviceName } }] : []),
+          ...getEnvironmentUiFilterES(environment),
+        ],
+      },
+    };
 
-  const { intervalString } = getBucketSize({ start, end, numBuckets: 20 });
+    const { intervalString } = getBucketSize({ start, end, numBuckets: 20 });
 
-  const aggs = {
-    timeseries: {
-      date_histogram: {
-        field: '@timestamp',
-        fixed_interval: intervalString,
+    const aggs = {
+      timeseries: {
+        date_histogram: {
+          field: '@timestamp',
+          fixed_interval: intervalString,
+        },
       },
-    },
-  };
+    };
 
-  const params = {
-    apm: { events: [ProcessorEvent.error] },
-    body: { size: 0, query, aggs },
-  };
+    const params = {
+      apm: { events: [ProcessorEvent.error] },
+      body: { size: 0, query, aggs },
+    };
 
-  const resp = await apmEventClient.search(params);
+    const resp = await apmEventClient.search(params);
 
-  if (!resp.aggregations) {
-    return [];
-  }
+    if (!resp.aggregations) {
+      return [];
+    }
 
-  return resp.aggregations.timeseries.buckets.map((bucket) => {
-    return {
-      x: bucket.key,
-      y: bucket.doc_count,
-    };
+    return resp.aggregations.timeseries.buckets.map((bucket) => {
+      return {
+        x: bucket.key,
+        y: bucket.doc_count,
+      };
+    });
   });
 }