Merge pull request #3866 from Infisical/feat/posthogEventBatch

feat(telemetry): Add aggregated events and groups to posthog

backend/src/services/telemetry/telemetry-queue.ts

@@ -7,13 +7,18 @@ import { QueueJobs, QueueName, TQueueServiceFactory } from "@app/queue";
 
 import { getServerCfg } from "../super-admin/super-admin-service";
 import { TTelemetryDALFactory } from "./telemetry-dal";
-import { TELEMETRY_SECRET_OPERATIONS_KEY, TELEMETRY_SECRET_PROCESSED_KEY } from "./telemetry-service";
+import {
+  TELEMETRY_SECRET_OPERATIONS_KEY,
+  TELEMETRY_SECRET_PROCESSED_KEY,
+  TTelemetryServiceFactory
+} from "./telemetry-service";
 import { PostHogEventTypes } from "./telemetry-types";
 
 type TTelemetryQueueServiceFactoryDep = {
   queueService: TQueueServiceFactory;
   keyStore: Pick<TKeyStoreFactory, "getItem" | "deleteItem">;
   telemetryDAL: TTelemetryDALFactory;
+  telemetryService: TTelemetryServiceFactory;
 };
 
 export type TTelemetryQueueServiceFactory = ReturnType<typeof telemetryQueueServiceFactory>;
@@ -21,7 +26,8 @@ export type TTelemetryQueueServiceFactory = ReturnType<typeof telemetryQueueServ
 export const telemetryQueueServiceFactory = ({
   queueService,
   keyStore,
-  telemetryDAL
+  telemetryDAL,
+  telemetryService
 }: TTelemetryQueueServiceFactoryDep) => {
   const appCfg = getConfig();
   const postHog =
@@ -48,6 +54,10 @@ export const telemetryQueueServiceFactory = ({
     await keyStore.deleteItem(TELEMETRY_SECRET_OPERATIONS_KEY);
   });
 
+  queueService.start(QueueName.TelemetryAggregatedEvents, async () => {
+    await telemetryService.processAggregatedEvents();
+  });
+
   // every day at midnight a telemetry job executes on self-hosted instances
   // this sends some telemetry information like instance id secrets operated etc
   const startTelemetryCheck = async () => {
@@ -60,11 +70,26 @@ export const telemetryQueueServiceFactory = ({
       { pattern: "0 0 * * *", utc: true },
       QueueName.TelemetryInstanceStats // just a job id
     );
+
+    // clear previous aggregated events job
+    await queueService.stopRepeatableJob(
+      QueueName.TelemetryAggregatedEvents,
+      QueueJobs.TelemetryAggregatedEvents,
+      { pattern: "*/5 * * * *", utc: true },
+      QueueName.TelemetryAggregatedEvents // just a job id
+    );
+
     if (postHog) {
       await queueService.queue(QueueName.TelemetryInstanceStats, QueueJobs.TelemetryInstanceStats, undefined, {
         jobId: QueueName.TelemetryInstanceStats,
         repeat: { pattern: "0 0 * * *", utc: true }
       });
+
+      // Start aggregated events job (runs every five minutes)
+      await queueService.queue(QueueName.TelemetryAggregatedEvents, QueueJobs.TelemetryAggregatedEvents, undefined, {
+        jobId: QueueName.TelemetryAggregatedEvents,
+        repeat: { pattern: "*/5 * * * *", utc: true }
+      });
     }
   };
 
@@ -72,6 +97,10 @@ export const telemetryQueueServiceFactory = ({
     logger.error(err?.failedReason, `${QueueName.TelemetryInstanceStats}: failed`);
   });
 
+  queueService.listen(QueueName.TelemetryAggregatedEvents, "failed", (err) => {
+    logger.error(err?.failedReason, `${QueueName.TelemetryAggregatedEvents}: failed`);
+  });
+
   return {
     startTelemetryCheck
   };

backend/src/services/telemetry/telemetry-service.ts

@@ -1,3 +1,4 @@
+import { createHash, randomUUID } from "crypto";
 import { PostHog } from "posthog-node";
 
 import { TLicenseServiceFactory } from "@app/ee/services/license/license-service";
@@ -12,12 +13,49 @@ import { PostHogEventTypes, TPostHogEvent, TSecretModifiedEvent } from "./teleme
 export const TELEMETRY_SECRET_PROCESSED_KEY = "telemetry-secret-processed";
 export const TELEMETRY_SECRET_OPERATIONS_KEY = "telemetry-secret-operations";
 
+export const POSTHOG_AGGREGATED_EVENTS = [PostHogEventTypes.SecretPulled];
+const TELEMETRY_AGGREGATED_KEY_EXP = 900; // 15mins
+
+// Bucket configuration
+const TELEMETRY_BUCKET_COUNT = 30;
+const TELEMETRY_BUCKET_NAMES = Array.from(
+  { length: TELEMETRY_BUCKET_COUNT },
+  (_, i) => `bucket-${i.toString().padStart(2, "0")}`
+);
+
+type AggregatedEventData = Record<string, unknown>;
+type SingleEventData = {
+  distinctId: string;
+  event: string;
+  properties: unknown;
+  organizationId: string;
+};
+
 export type TTelemetryServiceFactory = ReturnType<typeof telemetryServiceFactory>;
 export type TTelemetryServiceFactoryDep = {
-  keyStore: Pick<TKeyStoreFactory, "getItem" | "incrementBy">;
+  keyStore: Pick<
+    TKeyStoreFactory,
+    "incrementBy" | "deleteItemsByKeyIn" | "setItemWithExpiry" | "getKeysByPattern" | "getItems"
+  >;
   licenseService: Pick<TLicenseServiceFactory, "getInstanceType">;
 };
 
+const getBucketForDistinctId = (distinctId: string): string => {
+  // Use SHA-256 hash for consistent distribution
+  const hash = createHash("sha256").update(distinctId).digest("hex");
+
+  // Take first 8 characters and convert to number for better distribution
+  const hashNumber = parseInt(hash.substring(0, 8), 16);
+  const bucketIndex = hashNumber % TELEMETRY_BUCKET_COUNT;
+
+  return TELEMETRY_BUCKET_NAMES[bucketIndex];
+};
+
+export const createTelemetryEventKey = (event: string, distinctId: string): string => {
+  const bucketId = getBucketForDistinctId(distinctId);
+  return `telemetry-event-${event}-${bucketId}-${distinctId}-${randomUUID()}`;
+};
+
 export const telemetryServiceFactory = ({ keyStore, licenseService }: TTelemetryServiceFactoryDep) => {
   const appCfg = getConfig();
 
@@ -64,11 +102,33 @@ To opt into telemetry, you can set "TELEMETRY_ENABLED=true" within the environme
       const instanceType = licenseService.getInstanceType();
       // capture posthog only when its cloud or signup event happens in self-hosted
       if (instanceType === InstanceType.Cloud || event.event === PostHogEventTypes.UserSignedUp) {
-        postHog.capture({
-          event: event.event,
-          distinctId: event.distinctId,
-          properties: event.properties
-        });
+        if (event.organizationId) {
+          try {
+            postHog.groupIdentify({ groupType: "organization", groupKey: event.organizationId });
+          } catch (error) {
+            logger.error(error, "Failed to identify PostHog organization");
+          }
+        }
+        if (POSTHOG_AGGREGATED_EVENTS.includes(event.event)) {
+          const eventKey = createTelemetryEventKey(event.event, event.distinctId);
+          await keyStore.setItemWithExpiry(
+            eventKey,
+            TELEMETRY_AGGREGATED_KEY_EXP,
+            JSON.stringify({
+              distinctId: event.distinctId,
+              event: event.event,
+              properties: event.properties,
+              organizationId: event.organizationId
+            })
+          );
+        } else {
+          postHog.capture({
+            event: event.event,
+            distinctId: event.distinctId,
+            properties: event.properties,
+            ...(event.organizationId ? { groups: { organization: event.organizationId } } : {})
+          });
+        }
         return;
       }
 
@@ -89,6 +149,160 @@ To opt into telemetry, you can set "TELEMETRY_ENABLED=true" within the environme
     }
   };
 
+  const aggregateGroupProperties = (events: SingleEventData[]): AggregatedEventData => {
+    const aggregatedData: AggregatedEventData = {};
+
+    // Set the total count
+    aggregatedData.count = events.length;
+
+    events.forEach((event) => {
+      if (!event.properties) return;
+
+      Object.entries(event.properties as Record<string, unknown>).forEach(([key, value]: [string, unknown]) => {
+        if (Array.isArray(value)) {
+          // For arrays, count occurrences of each item
+          const existingCounts =
+            aggregatedData[key] &&
+            typeof aggregatedData[key] === "object" &&
+            aggregatedData[key]?.constructor === Object
+              ? (aggregatedData[key] as Record<string, number>)
+              : {};
+
+          value.forEach((item) => {
+            const itemKey = typeof item === "object" ? JSON.stringify(item) : String(item);
+            existingCounts[itemKey] = (existingCounts[itemKey] || 0) + 1;
+          });
+
+          aggregatedData[key] = existingCounts;
+        } else if (typeof value === "object" && value?.constructor === Object) {
+          // For objects, count occurrences of each field value
+          const existingCounts =
+            aggregatedData[key] &&
+            typeof aggregatedData[key] === "object" &&
+            aggregatedData[key]?.constructor === Object
+              ? (aggregatedData[key] as Record<string, number>)
+              : {};
+
+          if (value) {
+            Object.values(value).forEach((fieldValue) => {
+              const valueKey = typeof fieldValue === "object" ? JSON.stringify(fieldValue) : String(fieldValue);
+              existingCounts[valueKey] = (existingCounts[valueKey] || 0) + 1;
+            });
+          }
+          aggregatedData[key] = existingCounts;
+        } else if (typeof value === "number") {
+          // For numbers, add to existing sum
+          aggregatedData[key] = ((aggregatedData[key] as number) || 0) + value;
+        } else if (value !== undefined && value !== null) {
+          // For other types (strings, booleans, etc.), count occurrences
+          const stringValue = String(value);
+          const existingValue = aggregatedData[key];
+
+          if (!existingValue) {
+            aggregatedData[key] = { [stringValue]: 1 };
+          } else if (existingValue && typeof existingValue === "object" && existingValue.constructor === Object) {
+            const countObject = existingValue as Record<string, number>;
+            countObject[stringValue] = (countObject[stringValue] || 0) + 1;
+          } else {
+            const oldValue = String(existingValue);
+            aggregatedData[key] = {
+              [oldValue]: 1,
+              [stringValue]: 1
+            };
+          }
+        }
+      });
+    });
+
+    return aggregatedData;
+  };
+
+  const processBucketEvents = async (eventType: string, bucketId: string) => {
+    if (!postHog) return 0;
+
+    try {
+      const bucketPattern = `telemetry-event-${eventType}-${bucketId}-*`;
+      const bucketKeys = await keyStore.getKeysByPattern(bucketPattern);
+
+      if (bucketKeys.length === 0) return 0;
+
+      const bucketEvents = await keyStore.getItems(bucketKeys);
+      let bucketEventsParsed: SingleEventData[] = [];
+
+      try {
+        bucketEventsParsed = bucketEvents
+          .filter((event) => event !== null)
+          .map((event) => JSON.parse(event as string) as SingleEventData);
+      } catch (error) {
+        logger.error(error, `Failed to parse bucket events for ${eventType} in ${bucketId}`);
+        return 0;
+      }
+
+      const eventsGrouped = new Map<string, SingleEventData[]>();
+
+      bucketEventsParsed.forEach((event) => {
+        const key = JSON.stringify({ id: event.distinctId, org: event.organizationId });
+        if (!eventsGrouped.has(key)) {
+          eventsGrouped.set(key, []);
+        }
+        eventsGrouped.get(key)!.push(event);
+      });
+
+      if (eventsGrouped.size === 0) return 0;
+
+      for (const [eventsKey, events] of eventsGrouped) {
+        const key = JSON.parse(eventsKey) as { id: string; org?: string };
+        if (key.org) {
+          try {
+            postHog.groupIdentify({ groupType: "organization", groupKey: key.org });
+          } catch (error) {
+            logger.error(error, "Failed to identify PostHog organization");
+          }
+        }
+        const properties = aggregateGroupProperties(events);
+
+        postHog.capture({
+          event: `${eventType} aggregated`,
+          distinctId: key.id,
+          properties,
+          ...(key.org ? { groups: { organization: key.org } } : {})
+        });
+      }
+
+      // Clean up processed data for this bucket
+      await keyStore.deleteItemsByKeyIn(bucketKeys);
+
+      logger.info(`Processed ${bucketEventsParsed.length} events from bucket ${bucketId} for ${eventType}`);
+      return bucketEventsParsed.length;
+    } catch (error) {
+      logger.error(error, `Failed to process bucket ${bucketId} for ${eventType}`);
+      return 0;
+    }
+  };
+
+  const processAggregatedEvents = async () => {
+    if (!postHog) return;
+
+    for (const eventType of POSTHOG_AGGREGATED_EVENTS) {
+      let totalProcessed = 0;
+
+      logger.info(`Starting bucket processing for ${eventType}`);
+
+      // Process each bucket sequentially to control memory usage
+      for (const bucketId of TELEMETRY_BUCKET_NAMES) {
+        try {
+          // eslint-disable-next-line no-await-in-loop
+          const processed = await processBucketEvents(eventType, bucketId);
+          totalProcessed += processed;
+        } catch (error) {
+          logger.error(error, `Failed to process bucket ${bucketId} for ${eventType}`);
+        }
+      }
+
+      logger.info(`Completed processing ${totalProcessed} total events for ${eventType}`);
+    }
+  };
+
   const flushAll = async () => {
     if (postHog) {
       await postHog.shutdownAsync();
@@ -98,6 +312,8 @@ To opt into telemetry, you can set "TELEMETRY_ENABLED=true" within the environme
   return {
     sendLoopsEvent,
     sendPostHogEvents,
-    flushAll
+    processAggregatedEvents,
+    flushAll,
+    getBucketForDistinctId
   };
 };

backend/src/services/telemetry/telemetry-types.ts

@@ -1,3 +1,13 @@
+import {
+  IdentityActor,
+  KmipClientActor,
+  PlatformActor,
+  ScimClientActor,
+  ServiceActor,
+  UnknownUserActor,
+  UserActor
+} from "@app/ee/services/audit-log/audit-log-types";
+
 export enum PostHogEventTypes {
   SecretPush = "secrets pushed",
   SecretPulled = "secrets pulled",
@@ -40,6 +50,14 @@ export type TSecretModifiedEvent = {
     secretPath: string;
     channel?: string;
     userAgent?: string;
+    actor?:
+      | UserActor
+      | IdentityActor
+      | ServiceActor
+      | ScimClientActor
+      | PlatformActor
+      | UnknownUserActor
+      | KmipClientActor;
   };
 };
 
@@ -214,7 +232,7 @@ export type TInvalidateCacheEvent = {
   };
 };
 
-export type TPostHogEvent = { distinctId: string } & (
+export type TPostHogEvent = { distinctId: string; organizationId?: string } & (
   | TSecretModifiedEvent
   | TAdminInitEvent
   | TUserSignedUpEvent