first pass of basic osquery usage stats collection

x-pack/plugins/osquery/kibana.json

@@ -10,6 +10,7 @@
   "kibanaVersion": "kibana",
   "optionalPlugins": [
     "home",
+    "usageCollection",
     "lens"
   ],
   "requiredBundles": [

x-pack/plugins/osquery/server/plugin.ts

@@ -18,6 +18,7 @@ import { OsqueryPluginSetup, OsqueryPluginStart, SetupPlugins, StartPlugins } fr
 import { defineRoutes } from './routes';
 import { osquerySearchStrategyProvider } from './search_strategy/osquery';
 import { initSavedObjects } from './saved_objects';
+import { initUsageCollectors } from './usage';
 import { OsqueryAppContext, OsqueryAppContextService } from './lib/osquery_app_context_services';
 import { ConfigType } from './config';
 
@@ -48,6 +49,11 @@ export class OsqueryPlugin implements Plugin<OsqueryPluginSetup, OsqueryPluginSt
     };
 
     initSavedObjects(core.savedObjects, osqueryContext);
+    initUsageCollectors({
+      core,
+      osqueryContext,
+      usageCollection: plugins.usageCollection,
+    });
     defineRoutes(router, osqueryContext);
 
     core.getStartServices().then(([, depsStart]) => {

x-pack/plugins/osquery/server/routes/action/create_action_route.ts

@@ -18,6 +18,8 @@ import {
   CreateActionRequestBodySchema,
 } from '../../../common/schemas/routes/action/create_action_request_body_schema';
 
+import {getUsageRecorder} from '../usage'
+
 export const createActionRoute = (router: IRouter, osqueryContext: OsqueryAppContext) => {
   router.post(
     {
@@ -39,34 +41,45 @@ export const createActionRoute = (router: IRouter, osqueryContext: OsqueryAppCon
         osqueryContext,
         agentSelection
       );
-
+      const usageRecorder = getUsageRecorder()
+      usageRecorder.incrementCallCount('live_query')
       if (!selectedAgents.length) {
+        usageRecorder.incrementErrorCount('live_query')
         return response.badRequest({ body: new Error('No agents found for selection') });
       }
 
-      const action = {
-        action_id: uuid.v4(),
-        '@timestamp': moment().toISOString(),
-        expiration: moment().add(1, 'days').toISOString(),
-        type: 'INPUT_ACTION',
-        input_type: 'osquery',
-        agents: selectedAgents,
-        data: {
-          id: uuid.v4(),
-          query: request.body.query,
-        },
-      };
-      const actionResponse = await esClient.index<{}, {}>({
-        index: '.fleet-actions',
-        body: action,
-      });
+      try {
+        const action = {
+          action_id: uuid.v4(),
+          '@timestamp': moment().toISOString(),
+          expiration: moment().add(1, 'days').toISOString(),
+          type: 'INPUT_ACTION',
+          input_type: 'osquery',
+          agents: selectedAgents,
+          data: {
+            id: uuid.v4(),
+            query: request.body.query,
+          },
+        };
+        const actionResponse = await esClient.index<{}, {}>({
+          index: '.fleet-actions',
+          body: action,
+        });
 
-      return response.ok({
-        body: {
-          response: actionResponse,
-          actions: [action],
-        },
-      });
+        return response.ok({
+          body: {
+            response: actionResponse,
+            actions: [action],
+          },
+        });
+      } catch (error) {
+        usageRecorder.incrementErrorCount('live_query')
+        return response.customError({
+          statusCode: 500,
+          body: new Error(`Error occurred whlie processing ${error}`) ,
+        });
+
+      }
     }
   );
 };

x-pack/plugins/osquery/server/routes/usage/index.ts

@@ -0,0 +1,8 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export * from './recorder'

x-pack/plugins/osquery/server/routes/usage/recorder.ts

@@ -0,0 +1,48 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export interface RouteUsageMetric {
+    call_count: number;
+    error_count: number;
+}
+
+// TODO: use ES for this recording
+class UsageRecorder {
+    private counts = new Map<string, number>()
+    public incrementCallCount(route: string, increment: number = 1) {
+        const count = this.counts.get(route) ?? 0
+        this.counts.set(route, count + increment)
+    }
+    public getCallCount(route: string): number {
+        return this.counts.get(route) ?? 0;
+    }
+
+    private errors = new Map<string, number>()
+    public incrementErrorCount(route: string, increment: number = 1) {
+        const count = this.errors.get(route) ?? 0
+        this.errors.set(route, count + increment)
+    }
+    public getErrorCount(route: string): number {
+        return this.errors.get(route) ?? 0;
+    }
+
+    public getRouteMetric(route: string): RouteUsageMetric {
+        return {
+            call_count: this.getCallCount(route),
+            error_count: this.getErrorCount(route)
+        }
+    }
+}
+
+let usageRecorder: UsageRecorder;
+
+export const getUsageRecorder = (): UsageRecorder => {
+    if (usageRecorder == null) {
+        usageRecorder = new UsageRecorder()
+    }
+    return usageRecorder
+}

x-pack/plugins/osquery/server/types.ts

@@ -11,6 +11,7 @@ import {
   PluginStart as DataPluginStart,
 } from '../../../../src/plugins/data/server';
 import { FleetStartContract } from '../../fleet/server';
+import { UsageCollectionSetup } from '../../../../src/plugins/usage_collection/server';
 import { PluginSetupContract } from '../../features/server';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
@@ -19,6 +20,7 @@ export interface OsqueryPluginSetup {}
 export interface OsqueryPluginStart {}
 
 export interface SetupPlugins {
+  usageCollection?: UsageCollectionSetup;
   actions: ActionsPlugin['setup'];
   data: DataPluginSetup;
   features: PluginSetupContract;

x-pack/plugins/osquery/server/usage/collector.ts

@@ -0,0 +1,39 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { CoreSetup } from '../../../../../src/core/server';
+import { CollectorFetchContext } from '../../../../../src/plugins/usage_collection/server';
+import { getBeatUsage, getLiveQueryUsage } from './fetchers';
+import { CollectorDependencies, usageSchema } from './types';
+
+export type RegisterCollector = (deps: CollectorDependencies) => void;
+export async function getInternalSavedObjectsClient(core: CoreSetup) {
+  return core.getStartServices().then(async ([coreStart]) => {
+    return coreStart.savedObjects.createInternalRepository();
+  });
+}
+
+export const registerCollector: RegisterCollector = ({ usageCollection }) => {
+  if (!usageCollection) {
+    return;
+  }
+  const collector = usageCollection.makeUsageCollector<any>({
+    type: 'osquery',
+    schema: usageSchema,
+    isReady: () => true,
+    fetch: async ({ esClient }: CollectorFetchContext): Promise<any> => {
+      return {
+        beat_metrics: {
+          usage: await getBeatUsage(esClient),
+        },
+        live_query_usage: getLiveQueryUsage(),
+      };
+    },
+  });
+
+  usageCollection.registerCollector(collector);
+};

x-pack/plugins/osquery/server/usage/constants.ts

@@ -0,0 +1,8 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const METRICS_INDICES = 'logs-elastic_agent.osquerybeat*';

x-pack/plugins/osquery/server/usage/fetchers.ts

@@ -0,0 +1,121 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import {
+    SingleBucketAggregate,
+    TopHitsAggregate,
+    ValueAggregate,
+} from '@elastic/elasticsearch/api/types';
+import {getUsageRecorder} from '../routes/usage'
+import { ElasticsearchClient } from '../../../../../src/core/server';
+import { METRICS_INDICES } from './constants';
+
+export interface MetricEntry {
+  max?: number;
+  latest?: number;
+  avg?: number;
+}
+
+export interface BeatMetricAggregation {
+  rss: MetricEntry;
+  cpuMs: MetricEntry;
+}
+
+// TODO: pipe this through ES
+export function getLiveQueryUsage() {
+    const usageRecorder = getUsageRecorder()
+    return usageRecorder.getRouteMetric('live_query')
+}
+
+export async function getBeatUsage(esClient: ElasticsearchClient) {
+  // is there a better way to get these aggregates?
+  // needs a time window limit to make sure the reports are fresh
+  // XXX: these aggregates conflate agents, they should be broken out by id
+  // XXX: currently cpu is recorded as a duration rather than a load %
+  const { body: metricResponse } = await esClient.search({
+    body: {
+      size: 0,
+      aggs: {
+        lastDay: {
+          filter: {
+            range: {
+              '@timestamp': {
+                gte: 'now-24h',
+                lte: 'now',
+              },
+            },
+          },
+          aggs: {
+            latest: {
+              top_hits: {
+                sort: [
+                  {
+                    '@timestamp': {
+                      order: 'desc',
+                    },
+                  },
+                ],
+                size: 1,
+              },
+            },
+            max_rss: {
+              max: {
+                field: 'monitoring.metrics.beat.memstats.rss',
+              },
+            },
+            avg_rss: {
+              avg: {
+                field: 'monitoring.metrics.beat.memstats.rss',
+              },
+            },
+            max_cpu: {
+              max: {
+                field: 'monitoring.metrics.beat.cpu.total.time.ms',
+              },
+            },
+            avg_cpu: {
+              avg: {
+                field: 'monitoring.metrics.beat.cpu.total.time.ms',
+              },
+            },
+          },
+        },
+      },
+    },
+    index: METRICS_INDICES,
+  });
+  const lastDayAggs = metricResponse.aggregations?.lastDay as SingleBucketAggregate;
+  const result: BeatMetricAggregation = {
+    rss: {},
+    cpuMs: {},
+  };
+
+  // XXX: discrimating the union types gets hairy when attempting to genericize, figure out a fix!
+  if ('max_rss' in lastDayAggs) {
+      result.rss.max = (lastDayAggs.max_rss as ValueAggregate).value
+  }
+
+  if ('avg_rss' in lastDayAggs) {
+      result.rss.avg = (lastDayAggs.max_rss as ValueAggregate).value
+  }
+
+  if ('max_cpu' in lastDayAggs) {
+      result.cpuMs.max = (lastDayAggs.max_cpu as ValueAggregate).value
+  }
+
+  if ('avg_cpu' in lastDayAggs) {
+      result.cpuMs.avg = (lastDayAggs.max_cpu as ValueAggregate).value
+  }
+
+  if ('latest' in lastDayAggs) {
+    const latest = (lastDayAggs.latest as TopHitsAggregate).hits.hits[0]?._source?.monitoring.metrics.beat;
+    result.cpuMs.latest = latest.cpu.total.time.ms;
+    result.rss.latest = latest.memstats.rss;
+  }
+
+  return result;
+}

x-pack/plugins/osquery/server/usage/index.ts

@@ -0,0 +1,15 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { CollectorDependencies } from './types';
+import { registerCollector } from './collector';
+
+export type InitUsageCollectors = (deps: CollectorDependencies) => void;
+
+export const initUsageCollectors: InitUsageCollectors = (dependencies) => {
+  registerCollector(dependencies);
+};