Instrument task manager

dgieselaar · May 3, 2021 · 9b6bb25 · 9b6bb25
1 parent cbcc800
commit 9b6bb25
Show file tree

Hide file tree

Showing 7 changed files with 202 additions and 148 deletions.
diff --git a/packages/kbn-apm-utils/src/index.ts b/packages/kbn-apm-utils/src/index.ts
@@ -16,6 +16,8 @@ export interface SpanOptions {
   labels?: Record<string, string>;
 }
 
+type Span = Exclude<typeof agent.currentSpan, undefined | null>;
+
 export function parseSpanOptions(optionsOrName: SpanOptions | string) {
   const options = typeof optionsOrName === 'string' ? { name: optionsOrName } : optionsOrName;
 
@@ -30,7 +32,7 @@ const runInNewContext = <T extends (...args: any[]) => any>(cb: T): ReturnType<T
 
 export async function withSpan<T>(
   optionsOrName: SpanOptions | string,
-  cb: () => Promise<T>
+  cb: (span?: Span) => Promise<T>
 ): Promise<T> {
   const options = parseSpanOptions(optionsOrName);
 
@@ -71,13 +73,17 @@ export async function withSpan<T>(
       span.addLabels(labels);
     }
 
-    return cb()
+    return cb(span)
       .then((res) => {
-        span.outcome = 'success';
+        if (!span.outcome || span.outcome === 'unknown') {
+          span.outcome = 'success';
+        }
         return res;
       })
       .catch((err) => {
-        span.outcome = 'failure';
+        if (!span.outcome || span.outcome === 'unknown') {
+          span.outcome = 'failure';
+        }
         throw err;
       })
       .finally(() => {

diff --git a/x-pack/plugins/actions/server/lib/action_executor.ts b/x-pack/plugins/actions/server/lib/action_executor.ts
@@ -7,6 +7,7 @@
 
 import type { PublicMethodsOf } from '@kbn/utility-types';
 import { Logger, KibanaRequest } from 'src/core/server';
+import { withSpan } from '@kbn/apm-utils';
 import { validateParams, validateConfig, validateSecrets } from './validate_with_schema';
 import {
   ActionTypeExecutorResult,
@@ -78,113 +79,135 @@ export class ActionExecutor {
       );
     }
 
-    const {
-      logger,
-      spaces,
-      getServices,
-      encryptedSavedObjectsClient,
-      actionTypeRegistry,
-      eventLogger,
-      preconfiguredActions,
-      getActionsClientWithRequest,
-    } = this.actionExecutorContext!;
-
-    const services = getServices(request);
-    const spaceId = spaces && spaces.getSpaceId(request);
-    const namespace = spaceId && spaceId !== 'default' ? { namespace: spaceId } : {};
-
-    const { actionTypeId, name, config, secrets } = await getActionInfo(
-      await getActionsClientWithRequest(request, source),
-      encryptedSavedObjectsClient,
-      preconfiguredActions,
-      actionId,
-      namespace.namespace
-    );
+    return withSpan(
+      {
+        name: `execute_action`,
+        type: 'actions',
+        labels: {
+          actionId,
+        },
+      },
+      async (span) => {
+        const {
+          logger,
+          spaces,
+          getServices,
+          encryptedSavedObjectsClient,
+          actionTypeRegistry,
+          eventLogger,
+          preconfiguredActions,
+          getActionsClientWithRequest,
+        } = this.actionExecutorContext!;
 
-    if (!actionTypeRegistry.isActionExecutable(actionId, actionTypeId, { notifyUsage: true })) {
-      actionTypeRegistry.ensureActionTypeEnabled(actionTypeId);
-    }
-    const actionType = actionTypeRegistry.get(actionTypeId);
-
-    let validatedParams: Record<string, unknown>;
-    let validatedConfig: Record<string, unknown>;
-    let validatedSecrets: Record<string, unknown>;
-
-    try {
-      validatedParams = validateParams(actionType, params);
-      validatedConfig = validateConfig(actionType, config);
-      validatedSecrets = validateSecrets(actionType, secrets);
-    } catch (err) {
-      return { status: 'error', actionId, message: err.message, retry: false };
-    }
+        const services = getServices(request);
+        const spaceId = spaces && spaces.getSpaceId(request);
+        const namespace = spaceId && spaceId !== 'default' ? { namespace: spaceId } : {};
+
+        const { actionTypeId, name, config, secrets } = await getActionInfo(
+          await getActionsClientWithRequest(request, source),
+          encryptedSavedObjectsClient,
+          preconfiguredActions,
+          actionId,
+          namespace.namespace
+        );
+
+        if (span) {
+          span.name = `execute_action ${actionTypeId}`;
+          span.addLabels({
+            actionTypeId,
+          });
+        }
+
+        if (!actionTypeRegistry.isActionExecutable(actionId, actionTypeId, { notifyUsage: true })) {
+          actionTypeRegistry.ensureActionTypeEnabled(actionTypeId);
+        }
+        const actionType = actionTypeRegistry.get(actionTypeId);
+
+        let validatedParams: Record<string, unknown>;
+        let validatedConfig: Record<string, unknown>;
+        let validatedSecrets: Record<string, unknown>;
+
+        try {
+          validatedParams = validateParams(actionType, params);
+          validatedConfig = validateConfig(actionType, config);
+          validatedSecrets = validateSecrets(actionType, secrets);
+        } catch (err) {
+          span?.setOutcome('failure');
+          return { status: 'error', actionId, message: err.message, retry: false };
+        }
 
-    const actionLabel = `${actionTypeId}:${actionId}: ${name}`;
-    logger.debug(`executing action ${actionLabel}`);
-
-    const event: IEvent = {
-      event: { action: EVENT_LOG_ACTIONS.execute },
-      kibana: {
-        saved_objects: [
-          {
-            rel: SAVED_OBJECT_REL_PRIMARY,
-            type: 'action',
-            id: actionId,
-            ...namespace,
+        const actionLabel = `${actionTypeId}:${actionId}: ${name}`;
+        logger.debug(`executing action ${actionLabel}`);
+
+        const event: IEvent = {
+          event: { action: EVENT_LOG_ACTIONS.execute },
+          kibana: {
+            saved_objects: [
+              {
+                rel: SAVED_OBJECT_REL_PRIMARY,
+                type: 'action',
+                id: actionId,
+                ...namespace,
+              },
+            ],
           },
-        ],
-      },
-    };
+        };
 
-    eventLogger.startTiming(event);
-    let rawResult: ActionTypeExecutorResult<unknown>;
-    try {
-      rawResult = await actionType.executor({
-        actionId,
-        services,
-        params: validatedParams,
-        config: validatedConfig,
-        secrets: validatedSecrets,
-      });
-    } catch (err) {
-      rawResult = {
-        actionId,
-        status: 'error',
-        message: 'an error occurred while running the action executor',
-        serviceMessage: err.message,
-        retry: false,
-      };
-    }
-    eventLogger.stopTiming(event);
+        eventLogger.startTiming(event);
+        let rawResult: ActionTypeExecutorResult<unknown>;
+        try {
+          rawResult = await actionType.executor({
+            actionId,
+            services,
+            params: validatedParams,
+            config: validatedConfig,
+            secrets: validatedSecrets,
+          });
+        } catch (err) {
+          rawResult = {
+            actionId,
+            status: 'error',
+            message: 'an error occurred while running the action executor',
+            serviceMessage: err.message,
+            retry: false,
+          };
+        }
+        eventLogger.stopTiming(event);
 
-    // allow null-ish return to indicate success
-    const result = rawResult || {
-      actionId,
-      status: 'ok',
-    };
+        // allow null-ish return to indicate success
+        const result = rawResult || {
+          actionId,
+          status: 'ok',
+        };
 
-    event.event = event.event || {};
-
-    if (result.status === 'ok') {
-      event.event.outcome = 'success';
-      event.message = `action executed: ${actionLabel}`;
-    } else if (result.status === 'error') {
-      event.event.outcome = 'failure';
-      event.message = `action execution failure: ${actionLabel}`;
-      event.error = event.error || {};
-      event.error.message = actionErrorToMessage(result);
-      logger.warn(`action execution failure: ${actionLabel}: ${event.error.message}`);
-    } else {
-      event.event.outcome = 'failure';
-      event.message = `action execution returned unexpected result: ${actionLabel}: "${result.status}"`;
-      event.error = event.error || {};
-      event.error.message = 'action execution returned unexpected result';
-      logger.warn(
-        `action execution failure: ${actionLabel}: returned unexpected result "${result.status}"`
-      );
-    }
+        event.event = event.event || {};
 
-    eventLogger.logEvent(event);
-    return result;
+        if (result.status === 'ok') {
+          span?.setOutcome('success');
+          event.event.outcome = 'success';
+          event.message = `action executed: ${actionLabel}`;
+        } else if (result.status === 'error') {
+          span?.setOutcome('failure');
+          event.event.outcome = 'failure';
+          event.message = `action execution failure: ${actionLabel}`;
+          event.error = event.error || {};
+          event.error.message = actionErrorToMessage(result);
+          logger.warn(`action execution failure: ${actionLabel}: ${event.error.message}`);
+        } else {
+          span?.setOutcome('failure');
+          event.event.outcome = 'failure';
+          event.message = `action execution returned unexpected result: ${actionLabel}: "${result.status}"`;
+          event.error = event.error || {};
+          event.error.message = 'action execution returned unexpected result';
+          logger.warn(
+            `action execution failure: ${actionLabel}: returned unexpected result "${result.status}"`
+          );
+        }
+
+        eventLogger.logEvent(event);
+        return result;
+      }
+    );
   }
 }
 

diff --git a/x-pack/plugins/task_manager/server/queries/task_claiming.ts b/x-pack/plugins/task_manager/server/queries/task_claiming.ts
@@ -379,34 +379,41 @@ export class TaskClaiming {
       sort.unshift('_score');
     }
 
-    const apmTrans = apm.startTransaction(`taskManager markAvailableTasksAsClaimed`, 'taskManager');
-    const result = await this.taskStore.updateByQuery(
-      {
-        query: matchesClauses(
-          claimTasksById && claimTasksById.length
-            ? mustBeAllOf(asPinnedQuery(claimTasksById, queryForScheduledTasks))
-            : queryForScheduledTasks,
-          filterDownBy(InactiveTasks)
-        ),
-        script: updateFieldsAndMarkAsFailed(
-          {
-            ownerId: this.taskStore.taskManagerId,
-            retryAt: claimOwnershipUntil,
-          },
-          claimTasksById || [],
-          taskTypesToClaim,
-          taskTypesToSkip,
-          pick(this.taskMaxAttempts, taskTypesToClaim)
-        ),
-        sort,
-      },
-      {
-        max_docs: size,
-      }
+    const apmTrans = apm.startTransaction(
+      'markAvailableTasksAsClaimed',
+      `taskManager markAvailableTasksAsClaimed`
     );
-
-    if (apmTrans) apmTrans.end();
-    return result;
+    try {
+      const result = await this.taskStore.updateByQuery(
+        {
+          query: matchesClauses(
+            claimTasksById && claimTasksById.length
+              ? mustBeAllOf(asPinnedQuery(claimTasksById, queryForScheduledTasks))
+              : queryForScheduledTasks,
+            filterDownBy(InactiveTasks)
+          ),
+          script: updateFieldsAndMarkAsFailed(
+            {
+              ownerId: this.taskStore.taskManagerId,
+              retryAt: claimOwnershipUntil,
+            },
+            claimTasksById || [],
+            taskTypesToClaim,
+            taskTypesToSkip,
+            pick(this.taskMaxAttempts, taskTypesToClaim)
+          ),
+          sort,
+        },
+        {
+          max_docs: size,
+        }
+      );
+      apmTrans?.end('success');
+      return result;
+    } catch (err) {
+      apmTrans?.end('failure');
+      throw err;
+    }
   }
 
   /**

diff --git a/x-pack/plugins/task_manager/server/saved_objects/mappings.json b/x-pack/plugins/task_manager/server/saved_objects/mappings.json
@@ -29,6 +29,9 @@
       "status": {
         "type": "keyword"
       },
+      "traceparent": {
+        "type": "text"
+      },
       "params": {
         "type": "text"
       },

diff --git a/x-pack/plugins/task_manager/server/task.ts b/x-pack/plugins/task_manager/server/task.ts
@@ -257,6 +257,8 @@ export interface TaskInstance {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   state: Record<string, any>;
 
+  traceparent?: string;
+
   /**
    * The id of the user who scheduled this task.
    */
@@ -364,6 +366,7 @@ export type SerializedConcreteTaskInstance = Omit<
 > & {
   state: string;
   params: string;
+  traceparent: string;
   scheduledAt: string;
   startedAt: string | null;
   retryAt: string | null;