[Synthtrace] Introducing teardown for scenarios (elastic#209739)

### Background In some scenarios we need to perform some setup at bootstrap time, this setup could affect following scenarios. Take for example [failed_logs](https://github.com/elastic/kibana/blob/main/packages/kbn-apm-synthtrace/src/scenarios/failed_logs.ts) scenario where we create a pipeline that will do some checks in `log.level` property, if we try to run an scenario after that one we will enter into some issues. ### Changes This PR aims to introduce a `teardown` setup for scenarios where we could undo the changes done at `bootstrap` time.
TinaHeiligers · Feb 5, 2025 · c56d7ea · c56d7ea
1 parent bc5bff8
commit c56d7ea
Show file tree

Hide file tree

Showing 12 changed files with 83 additions and 7 deletions.
diff --git a/packages/kbn-apm-synthtrace/README.md b/packages/kbn-apm-synthtrace/README.md
@@ -119,7 +119,7 @@ Scenario files accept 3 arguments, 2 of them optional and 1 mandatory
 |-------------|:----------|------------------------------------------------------------------------------------------------------------------------------------------------------|
 | `generate`  | mandatory | This is the main function responsible for returning the events which will be indexed                                                                 |
 | `bootstrap`  | optional  | In case some setup needs to be done, before the data is generated, this function provides access to all available ES Clients to play with            |
-| `setClient` | optional  | By default the apmEsClient used to generate data. If anyother client like logsEsClient needs to be used instead, this is where it should be returned |
+| `teardown` | optional  | In case some setup needs to be done, after all data is generated, this function provides access to all available ES Clients to play with |
 
 The following options are supported:
 

diff --git a/packages/kbn-apm-synthtrace/src/cli/scenario.ts b/packages/kbn-apm-synthtrace/src/cli/scenario.ts
@@ -42,4 +42,5 @@ type Generate<TFields> = (options: {
 export type Scenario<TFields> = (options: RunOptions & { logger: Logger }) => Promise<{
   bootstrap?: (options: EsClients & KibanaClients) => Promise<void>;
   generate: Generate<TFields>;
+  teardown?: (options: EsClients & KibanaClients) => Promise<void>;
 }>;
diff --git a/packages/kbn-apm-synthtrace/src/cli/utils/start_live_data_upload.ts b/packages/kbn-apm-synthtrace/src/cli/utils/start_live_data_upload.ts
@@ -38,7 +38,11 @@ export async function startLiveDataUpload({
   } = await bootstrap(runOptions);
 
   const scenario = await getScenario({ file, logger });
-  const { generate, bootstrap: scenarioBootsrap } = await scenario({ ...runOptions, logger });
+  const {
+    generate,
+    bootstrap: scenarioBootsrap,
+    teardown: scenarioTearDown,
+  } = await scenario({ ...runOptions, logger });
 
   if (scenarioBootsrap) {
     await scenarioBootsrap({
@@ -59,11 +63,27 @@ export async function startLiveDataUpload({
   // @ts-expect-error upgrade typescript v4.9.5
   const cachedStreams: WeakMap<SynthtraceEsClient, PassThrough> = new WeakMap();
 
-  process.on('SIGINT', () => closeStreams());
-  process.on('SIGTERM', () => closeStreams());
-  process.on('SIGQUIT', () => closeStreams());
+  process.on('SIGINT', () => closeStreamsAndTeardown());
+  process.on('SIGTERM', () => closeStreamsAndTeardown());
+  process.on('SIGQUIT', () => closeStreamsAndTeardown());
+
+  async function closeStreamsAndTeardown() {
+    if (scenarioTearDown) {
+      try {
+        await scenarioTearDown({
+          apmEsClient,
+          logsEsClient,
+          infraEsClient,
+          otelEsClient,
+          syntheticsEsClient,
+          entitiesEsClient,
+          entitiesKibanaClient,
+        });
+      } catch (error) {
+        logger.error('Error during scenario teardown', error);
+      }
+    }
 
-  function closeStreams() {
     currentStreams.forEach((stream) => {
       stream.end(() => {
         process.exit(0);

diff --git a/packages/kbn-apm-synthtrace/src/cli/utils/synthtrace_worker.ts b/packages/kbn-apm-synthtrace/src/cli/utils/synthtrace_worker.ts
@@ -84,7 +84,7 @@ async function start() {
 
   logger.info(`Running scenario from ${bucketFrom.toISOString()} to ${bucketTo.toISOString()}`);
 
-  const { generate, bootstrap } = await scenario({ ...runOptions, logger });
+  const { generate, bootstrap, teardown } = await scenario({ ...runOptions, logger });
 
   if (bootstrap) {
     await bootstrap({
@@ -142,6 +142,18 @@ async function start() {
 
     await Promise.all(promises);
   });
+
+  if (teardown) {
+    await teardown({
+      apmEsClient,
+      logsEsClient,
+      infraEsClient,
+      syntheticsEsClient,
+      otelEsClient,
+      entitiesEsClient,
+      entitiesKibanaClient,
+    });
+  }
 }
 
 parentPort!.on('message', (message) => {

diff --git a/packages/kbn-apm-synthtrace/src/lib/logs/logs_synthtrace_es_client.ts b/packages/kbn-apm-synthtrace/src/lib/logs/logs_synthtrace_es_client.ts
@@ -52,6 +52,15 @@ export class LogsSynthtraceEsClient extends SynthtraceEsClient<LogDocument> {
     }
   }
 
+  async deleteIndexTemplate(name: IndexTemplateName) {
+    try {
+      await this.client.indices.deleteIndexTemplate({ name });
+      this.logger.info(`Index template successfully deleted: ${name}`);
+    } catch (err) {
+      this.logger.error(`Index template deletion failed: ${name} - ${err.message}`);
+    }
+  }
+
   async createComponentTemplate({
     name,
     mappings,
@@ -150,6 +159,17 @@ export class LogsSynthtraceEsClient extends SynthtraceEsClient<LogDocument> {
     }
   }
 
+  async deleteCustomPipeline(id = LogsCustom) {
+    try {
+      this.client.ingest.deletePipeline({
+        id,
+      });
+      this.logger.info(`Custom pipeline deleted: ${id}`);
+    } catch (err) {
+      this.logger.error(`Custom pipeline deletion failed: ${id} - ${err.message}`);
+    }
+  }
+
   getDefaultPipeline({ includeSerialization }: Pipeline = { includeSerialization: true }) {
     return logsPipeline({ includeSerialization });
   }

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/degraded_logs.ts b/packages/kbn-apm-synthtrace/src/scenarios/degraded_logs.ts
@@ -36,6 +36,9 @@ const scenario: Scenario<LogDocument> = async (runOptions) => {
     bootstrap: async ({ logsEsClient }) => {
       if (isLogsDb) await logsEsClient.createIndexTemplate(IndexTemplateName.LogsDb);
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient } }) => {
       const { logger } = runOptions;
 

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/failed_logs.ts b/packages/kbn-apm-synthtrace/src/scenarios/failed_logs.ts
@@ -75,6 +75,11 @@ const scenario: Scenario<LogDocument> = async (runOptions) => {
         },
       });
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteComponentTemplate(LogsCustom);
+      await logsEsClient.deleteCustomPipeline();
+      if (isLogsDb) await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient } }) => {
       const { logger } = runOptions;
 

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/logs_and_metrics.ts b/packages/kbn-apm-synthtrace/src/scenarios/logs_and_metrics.ts
@@ -38,6 +38,9 @@ const scenario: Scenario<LogDocument> = async (runOptions) => {
     bootstrap: async ({ logsEsClient }) => {
       if (isLogsDb) await logsEsClient.createIndexTemplate(IndexTemplateName.LogsDb);
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient, apmEsClient } }) => {
       const { numServices = 3 } = runOptions.scenarioOpts || {};
       const { logger } = runOptions;

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/logs_traces_hosts.ts b/packages/kbn-apm-synthtrace/src/scenarios/logs_traces_hosts.ts
@@ -52,6 +52,9 @@ const scenario: Scenario<LogDocument | InfraDocument | ApmFields> = async (runOp
     bootstrap: async ({ logsEsClient }) => {
       if (isLogsDb) await logsEsClient.createIndexTemplate(IndexTemplateName.LogsDb);
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient, infraEsClient, apmEsClient } }) => {
       const {
         numSpaces,

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/simple_logs.ts b/packages/kbn-apm-synthtrace/src/scenarios/simple_logs.ts
@@ -70,6 +70,9 @@ const scenario: Scenario<LogDocument> = async (runOptions) => {
     bootstrap: async ({ logsEsClient }) => {
       if (isLogsDb) await logsEsClient.createIndexTemplate(IndexTemplateName.LogsDb);
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient } }) => {
       const { logger } = runOptions;
 

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/simple_non_ecs_logs.ts b/packages/kbn-apm-synthtrace/src/scenarios/simple_non_ecs_logs.ts
@@ -78,6 +78,9 @@ const scenario: Scenario<LogDocument> = async (runOptions) => {
       await logsEsClient.createIndex('cloud-logs-synth.2-default');
       if (isLogsDb) await logsEsClient.createIndexTemplate(IndexTemplateName.LogsDb);
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient } }) => {
       const { logger } = runOptions;
 

diff --git a/packages/kbn-apm-synthtrace/src/scenarios/unstructured_logs.ts b/packages/kbn-apm-synthtrace/src/scenarios/unstructured_logs.ts
@@ -20,6 +20,9 @@ const scenario: Scenario<LogDocument> = async (runOptions) => {
     bootstrap: async ({ logsEsClient }) => {
       if (isLogsDb) await logsEsClient.createIndexTemplate(IndexTemplateName.LogsDb);
     },
+    teardown: async ({ logsEsClient }) => {
+      await logsEsClient.deleteIndexTemplate(IndexTemplateName.LogsDb);
+    },
     generate: ({ range, clients: { logsEsClient } }) => {
       const { logger } = runOptions;