Onboard dashboard_telemetry task type to use stateSchemaByVersion for task state validation (#161855)

Part of https://github.com/elastic/kibana/issues/159342.

In this PR, I'm preparing the `dashboard_telemetry` for serverless by
defining an explicit task state schema. This schema is used to validate
the task's state before saving the task but also when reading the task.
In the scenario an older Kibana node runs a task after a newer Kibana
node has stored additional task state, the unknown state properties will
be dropped. Additionally, this will prompt developers to be aware that
adding required fields to the task state is a breaking change that must
be handled with care. (see
https://github.com/elastic/kibana/issues/155764).

For more information on how to use `stateSchemaByVersion`, see
https://github.com/elastic/kibana/pull/159048 and
https://github.com/elastic/kibana/blob/main/x-pack/plugins/task_manager/README.md.

---------

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

src/plugins/dashboard/server/usage/dashboard_telemetry.ts

@@ -14,7 +14,11 @@ import { EmbeddablePersistableStateService } from '@kbn/embeddable-plugin/common
 import { type ControlGroupTelemetry, CONTROL_GROUP_TYPE } from '@kbn/controls-plugin/common';
 
 import { DashboardAttributes, SavedDashboardPanel } from '../../common/content_management';
-import { TASK_ID, DashboardTelemetryTaskState } from './dashboard_telemetry_collection_task';
+import { TASK_ID } from './dashboard_telemetry_collection_task';
+import { emptyState, type LatestTaskStateSchema } from './task_state';
+
+// TODO: Merge with LatestTaskStateSchema. Requires a refactor of collectPanelsByType() because
+// LatestTaskStateSchema doesn't allow mutations (uses ReadOnly<..>).
 export interface DashboardCollectorData {
   panels: {
     total: number;
@@ -127,9 +131,9 @@ export async function collectDashboardTelemetry(taskManager: TaskManagerStartCon
   const latestTaskState = await getLatestTaskState(taskManager);
 
   if (latestTaskState !== null) {
-    const state = latestTaskState[0].state as DashboardTelemetryTaskState;
+    const state = latestTaskState[0].state as LatestTaskStateSchema;
     return state.telemetry;
   }
 
-  return getEmptyDashboardData();
+  return emptyState.telemetry;
 }

src/plugins/dashboard/server/usage/dashboard_telemetry_collection_task.ts

@@ -15,12 +15,12 @@ import {
 } from '@kbn/task-manager-plugin/server';
 import { EmbeddableSetup } from '@kbn/embeddable-plugin/server';
 import { CoreSetup, Logger, SavedObjectReference } from '@kbn/core/server';
+import { stateSchemaByVersion, emptyState, type LatestTaskStateSchema } from './task_state';
 
 import {
   controlsCollectorFactory,
   collectPanelsByType,
   getEmptyDashboardData,
-  DashboardCollectorData,
 } from './dashboard_telemetry';
 import { injectReferences } from '../../common';
 import { DashboardAttributesAndReferences } from '../../common/types';
@@ -32,11 +32,6 @@ import { DashboardAttributes, SavedDashboardPanel } from '../../common/content_m
 const TELEMETRY_TASK_TYPE = 'dashboard_telemetry';
 export const TASK_ID = `Dashboard-${TELEMETRY_TASK_TYPE}`;
 
-export interface DashboardTelemetryTaskState {
-  runs: number;
-  telemetry: DashboardCollectorData;
-}
-
 export function initializeDashboardTelemetryTask(
   logger: Logger,
   core: CoreSetup,
@@ -60,6 +55,7 @@ function registerDashboardTelemetryTask(
     [TELEMETRY_TASK_TYPE]: {
       title: 'Dashboard telemetry collection task',
       timeout: '2m',
+      stateSchemaByVersion,
       createTaskRunner: dashboardTaskRunner(logger, core, embeddable),
     },
   });
@@ -70,7 +66,7 @@ async function scheduleTasks(logger: Logger, taskManager: TaskManagerStartContra
     return await taskManager.ensureScheduled({
       id: TASK_ID,
       taskType: TELEMETRY_TASK_TYPE,
-      state: { byDate: {}, suggestionsByDate: {}, saved: {}, runs: 0 },
+      state: emptyState,
       params: {},
     });
   } catch (e) {
@@ -80,7 +76,7 @@ async function scheduleTasks(logger: Logger, taskManager: TaskManagerStartContra
 
 export function dashboardTaskRunner(logger: Logger, core: CoreSetup, embeddable: EmbeddableSetup) {
   return ({ taskInstance }: RunContext) => {
-    const { state } = taskInstance;
+    const state = taskInstance.state as LatestTaskStateSchema;
 
     const getEsClient = async () => {
       const [coreStart] = await core.getStartServices();
@@ -172,11 +168,12 @@ export function dashboardTaskRunner(logger: Logger, core: CoreSetup, embeddable:
             );
           }
 
+          const updatedState: LatestTaskStateSchema = {
+            runs: state.runs + 1,
+            telemetry: dashboardData,
+          };
           return {
-            state: {
-              runs: (state.runs || 0) + 1,
-              telemetry: dashboardData,
-            },
+            state: updatedState,
             runAt: getNextMidnight(),
           };
         } catch (e) {

src/plugins/dashboard/server/usage/task_state.test.ts

@@ -0,0 +1,94 @@
+/*
+ * 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 { cloneDeep } from 'lodash';
+import { stateSchemaByVersion } from './task_state';
+
+describe('telemetry task state', () => {
+  describe('v1', () => {
+    const v1 = stateSchemaByVersion[1];
+    it('should work on empty object when running the up migration', () => {
+      const result = v1.up({});
+      expect(result).toMatchInlineSnapshot(`
+        Object {
+          "runs": 0,
+          "telemetry": Object {
+            "controls": Object {
+              "by_type": Object {},
+              "chaining_system": Object {},
+              "ignore_settings": Object {},
+              "label_position": Object {},
+              "total": 0,
+            },
+            "panels": Object {
+              "by_reference": 0,
+              "by_type": Object {},
+              "by_value": 0,
+              "total": 0,
+            },
+          },
+        }
+      `);
+    });
+
+    it(`shouldn't overwrite properties when running the up migration`, () => {
+      const state = {
+        runs: 1,
+        telemetry: {
+          panels: {
+            total: 2,
+            by_reference: 3,
+            by_value: 4,
+            by_type: {
+              foo: 5,
+            },
+          },
+          controls: {
+            total: 6,
+            chaining_system: { foo: 7 },
+            label_position: { foo: 8 },
+            ignore_settings: { foo: 9 },
+            by_type: { foo: 10 },
+          },
+        },
+      };
+      const result = v1.up(cloneDeep(state));
+      expect(result).toEqual(state);
+    });
+
+    it(`should migrate the old default state that didn't match the schema`, () => {
+      const result = v1.up({ byDate: {}, suggestionsByDate: {}, saved: {}, runs: 0 });
+      expect(result).toMatchInlineSnapshot(`
+        Object {
+          "runs": 0,
+          "telemetry": Object {
+            "controls": Object {
+              "by_type": Object {},
+              "chaining_system": Object {},
+              "ignore_settings": Object {},
+              "label_position": Object {},
+              "total": 0,
+            },
+            "panels": Object {
+              "by_reference": 0,
+              "by_type": Object {},
+              "by_value": 0,
+              "total": 0,
+            },
+          },
+        }
+      `);
+    });
+
+    it('should drop unknown properties when running the up migration', () => {
+      const state = { foo: true };
+      const result = v1.up(state);
+      expect(result).not.toHaveProperty('foo');
+    });
+  });
+});

src/plugins/dashboard/server/usage/task_state.ts

@@ -0,0 +1,95 @@
+/*
+ * 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 { schema, type TypeOf } from '@kbn/config-schema';
+
+/**
+ * WARNING: Do not modify the existing versioned schema(s) below, instead define a new version (ex: 2, 3, 4).
+ * This is required to support zero-downtime upgrades and rollbacks. See https://github.com/elastic/kibana/issues/155764.
+ *
+ * As you add a new schema version, don't forget to change latestTaskStateSchema variable to reference the latest schema.
+ * For example, changing stateSchemaByVersion[1].schema to stateSchemaByVersion[2].schema.
+ */
+export const stateSchemaByVersion = {
+  1: {
+    // A task that was created < 8.10 will go through this "up" migration
+    // to ensure it matches the v1 schema.
+    up: (state: Record<string, any>) => ({
+      runs: typeof state.runs === 'number' ? state.runs : 0,
+      telemetry: {
+        panels: {
+          total: state.telemetry?.panels?.total || 0,
+          by_reference: state.telemetry?.panels?.by_reference || 0,
+          by_value: state.telemetry?.panels?.by_value || 0,
+          by_type: state.telemetry?.panels?.by_type || {},
+        },
+        controls: {
+          total: state.telemetry?.controls?.total || 0,
+          chaining_system: state.telemetry?.controls?.chaining_system || {},
+          label_position: state.telemetry?.controls?.label_position || {},
+          ignore_settings: state.telemetry?.controls?.ignore_settings || {},
+          by_type: state.telemetry?.controls?.by_type || {},
+        },
+      },
+    }),
+    schema: schema.object({
+      runs: schema.number(),
+      telemetry: schema.object({
+        panels: schema.object({
+          total: schema.number(),
+          by_reference: schema.number(),
+          by_value: schema.number(),
+          by_type: schema.recordOf(
+            schema.string(),
+            schema.object({
+              total: schema.number(),
+              by_reference: schema.number(),
+              by_value: schema.number(),
+              details: schema.recordOf(schema.string(), schema.number()),
+            })
+          ),
+        }),
+        controls: schema.object({
+          total: schema.number(),
+          chaining_system: schema.recordOf(schema.string(), schema.number()),
+          label_position: schema.recordOf(schema.string(), schema.number()),
+          ignore_settings: schema.recordOf(schema.string(), schema.number()),
+          by_type: schema.recordOf(
+            schema.string(),
+            schema.object({
+              total: schema.number(),
+              details: schema.recordOf(schema.string(), schema.number()),
+            })
+          ),
+        }),
+      }),
+    }),
+  },
+};
+
+const latestTaskStateSchema = stateSchemaByVersion[1].schema;
+export type LatestTaskStateSchema = TypeOf<typeof latestTaskStateSchema>;
+
+export const emptyState: LatestTaskStateSchema = {
+  runs: 0,
+  telemetry: {
+    panels: {
+      total: 0,
+      by_reference: 0,
+      by_value: 0,
+      by_type: {},
+    },
+    controls: {
+      total: 0,
+      chaining_system: {},
+      ignore_settings: {},
+      label_position: {},
+      by_type: {},
+    },
+  },
+};