[SSE] Event stream TTL refresh (#17337)

This PR update the redis subscription iterator wrapper with an heartbeat
system. Heartbeat interval are based on event stream TTL. Those refresh
the event stream TTL and active streams in redis.

I also cleaned a few functions that were not used anymore.

---------

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>

packages/twenty-server/src/engine/core-modules/cache-storage/services/cache-storage.service.ts

@@ -245,6 +245,25 @@ export class CacheStorageService {
     return newValue;
   }
 
+  async expire(key: string, ttlMs: Milliseconds): Promise<boolean> {
+    if (this.isRedisCache()) {
+      return (this.cache as RedisCache).store.client.expire(
+        this.getKey(key),
+        ttlMs / 1000,
+      );
+    }
+
+    const existing = await this.get(key);
+
+    if (existing !== undefined) {
+      await this.set(key, existing, ttlMs);
+
+      return true;
+    }
+
+    return false;
+  }
+
   private isRedisCache() {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     return (this.cache.store as any)?.name === 'redis';

packages/twenty-server/src/engine/subscriptions/event-stream.service.ts

@@ -15,7 +15,6 @@ import {
   EventStreamExceptionCode,
 } from 'src/engine/subscriptions/event-stream.exception';
 import { type EventStreamData } from 'src/engine/subscriptions/types/event-stream-data.type';
-import { type ObjectRecordSubscriptionEvent } from 'src/engine/subscriptions/types/object-record-subscription-event.type';
 
 @Injectable()
 export class EventStreamService {
@@ -131,15 +130,6 @@ export class EventStreamService {
     return result;
   }
 
-  async getStreamData(
-    workspaceId: string,
-    eventStreamChannelId: string,
-  ): Promise<EventStreamData | undefined> {
-    const key = this.getEventStreamKey(workspaceId, eventStreamChannelId);
-
-    return this.cacheStorageService.get<EventStreamData>(key);
-  }
-
   async isAuthorized({
     workspaceId,
     eventStreamChannelId,
@@ -214,40 +204,25 @@ export class EventStreamService {
     }
   }
 
-  async getQueries(
-    workspaceId: string,
-    eventStreamId: string,
-  ): Promise<Map<string, RecordGqlOperationSignature>> {
-    const streamData = await this.getStreamData(workspaceId, eventStreamId);
+  async refreshEventStreamTTL({
+    workspaceId,
+    eventStreamChannelId,
+  }: {
+    workspaceId: string;
+    eventStreamChannelId: string;
+  }): Promise<boolean> {
+    const eventStreamKey = this.getEventStreamKey(
+      workspaceId,
+      eventStreamChannelId,
+    );
+    const activeStreamsKey = this.getActiveStreamsKey(workspaceId);
 
-    if (!isDefined(streamData)) {
-      return new Map();
-    }
+    const [eventStreamRefreshed, activeStreamsRefreshed] = await Promise.all([
+      this.cacheStorageService.expire(eventStreamKey, EVENT_STREAM_TTL_MS),
+      this.cacheStorageService.expire(activeStreamsKey, EVENT_STREAM_TTL_MS),
+    ]);
 
-    return new Map(Object.entries(streamData.queries));
-  }
-
-  matchQueriesWithEvent(
-    queries: Record<string, RecordGqlOperationSignature>,
-    event: ObjectRecordSubscriptionEvent,
-  ): string[] {
-    const matchedQueryIds: string[] = [];
-
-    for (const [queryId, operationSignature] of Object.entries(queries)) {
-      if (this.isQueryMatchingEvent(operationSignature, event)) {
-        matchedQueryIds.push(queryId);
-      }
-    }
-
-    return matchedQueryIds;
-  }
-
-  private isQueryMatchingEvent(
-    operationSignature: RecordGqlOperationSignature,
-    event: ObjectRecordSubscriptionEvent,
-  ): boolean {
-    // to be improved
-    return operationSignature.objectNameSingular === event.objectNameSingular;
+    return eventStreamRefreshed && activeStreamsRefreshed;
   }
 
   private getEventStreamKey(
@@ -260,4 +235,13 @@ export class EventStreamService {
   private getActiveStreamsKey(workspaceId: string): string {
     return `workspace:${workspaceId}:activeStreams`;
   }
+
+  private async getStreamData(
+    workspaceId: string,
+    eventStreamChannelId: string,
+  ): Promise<EventStreamData | undefined> {
+    const key = this.getEventStreamKey(workspaceId, eventStreamChannelId);
+
+    return this.cacheStorageService.get<EventStreamData>(key);
+  }
 }

packages/twenty-server/src/engine/workspace-event-emitter/utils/wrap-async-iterator-with-cleanup.ts

@@ -1,32 +0,0 @@
-export function wrapAsyncIteratorWithCleanup<T>(
-  iterator: AsyncIterableIterator<T>,
-  onClose: () => void | Promise<void>,
-): AsyncIterableIterator<T> {
-  return {
-    next: () => iterator.next(),
-    return: async () => {
-      let result: IteratorResult<T>;
-
-      try {
-        await onClose();
-      } finally {
-        result = (await iterator.return?.()) ?? {
-          done: true,
-          value: undefined,
-        };
-      }
-
-      return result;
-    },
-    throw: async (error) => {
-      if (iterator.throw) {
-        return iterator.throw(error);
-      }
-
-      throw error;
-    },
-    [Symbol.asyncIterator]() {
-      return this;
-    },
-  };
-}

packages/twenty-server/src/engine/workspace-event-emitter/utils/wrap-async-iterator-with-lifecycle.ts

@@ -0,0 +1,85 @@
+import { isDefined } from 'twenty-shared/utils';
+
+type AsyncIteratorLifecycleOptions = {
+  onHeartbeat?: () => Promise<boolean>;
+  heartbeatIntervalMs?: number;
+  onCleanup?: () => Promise<void>;
+};
+
+export function wrapAsyncIteratorWithLifecycle<T>(
+  iterator: AsyncIterableIterator<T>,
+  options: AsyncIteratorLifecycleOptions,
+): AsyncIterableIterator<T> {
+  const { onHeartbeat, heartbeatIntervalMs, onCleanup } = options;
+  let heartbeatInterval: NodeJS.Timeout | null = null;
+
+  const startHeartbeat = () => {
+    if (onHeartbeat && heartbeatIntervalMs) {
+      heartbeatInterval = setInterval(async () => {
+        try {
+          await onHeartbeat();
+        } catch {
+          // Heartbeat failure shouldn't crash the stream
+        }
+      }, heartbeatIntervalMs);
+    }
+  };
+
+  const cleanup = async () => {
+    if (heartbeatInterval) {
+      clearInterval(heartbeatInterval);
+      heartbeatInterval = null;
+    }
+    if (onCleanup) {
+      await onCleanup();
+    }
+  };
+
+  return {
+    next: async () => {
+      if (!isDefined(heartbeatInterval)) {
+        startHeartbeat();
+      }
+
+      let result: IteratorResult<T>;
+
+      try {
+        result = await iterator.next();
+      } catch (error) {
+        await cleanup();
+
+        throw error;
+      }
+
+      if (result.done) {
+        await cleanup();
+      }
+
+      return result;
+    },
+    return: async () => {
+      let result: IteratorResult<T>;
+
+      try {
+        await cleanup();
+      } finally {
+        result = (await iterator.return?.()) ?? {
+          done: true,
+          value: undefined,
+        };
+      }
+
+      return result;
+    },
+    throw: async (error) => {
+      await cleanup();
+      if (iterator.throw) {
+        return iterator.throw(error);
+      }
+      throw error;
+    },
+    [Symbol.asyncIterator]() {
+      return this;
+    },
+  };
+}

packages/twenty-server/src/engine/workspace-event-emitter/workspace-event-emitter.resolver.ts

@@ -15,6 +15,7 @@ import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorat
 import { NoPermissionGuard } from 'src/engine/guards/no-permission.guard';
 import { UserAuthGuard } from 'src/engine/guards/user-auth.guard';
 import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
+import { EVENT_STREAM_TTL_MS } from 'src/engine/subscriptions/constants/event-stream-ttl.constant';
 import { AddQuerySubscriptionInput } from 'src/engine/subscriptions/dtos/add-query-subscription.input';
 import {
   EventSubscriptionDTO,
@@ -30,7 +31,7 @@ import {
 } from 'src/engine/subscriptions/event-stream.exception';
 import { EventStreamService } from 'src/engine/subscriptions/event-stream.service';
 import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
-import { wrapAsyncIteratorWithCleanup } from 'src/engine/workspace-event-emitter/utils/wrap-async-iterator-with-cleanup';
+import { wrapAsyncIteratorWithLifecycle } from 'src/engine/workspace-event-emitter/utils/wrap-async-iterator-with-lifecycle';
 
 import { eventStreamIdToChannelId } from './utils/get-channel-id-from-event-stream-id';
 
@@ -123,11 +124,18 @@ export class WorkspaceEventEmitterResolver {
       throw error;
     }
 
-    return wrapAsyncIteratorWithCleanup(iterator, async () => {
-      await this.eventStreamService.destroyEventStream({
-        workspaceId: workspace.id,
-        eventStreamChannelId,
-      });
+    return wrapAsyncIteratorWithLifecycle(iterator, {
+      onHeartbeat: () =>
+        this.eventStreamService.refreshEventStreamTTL({
+          workspaceId: workspace.id,
+          eventStreamChannelId,
+        }),
+      heartbeatIntervalMs: EVENT_STREAM_TTL_MS / 5,
+      onCleanup: () =>
+        this.eventStreamService.destroyEventStream({
+          workspaceId: workspace.id,
+          eventStreamChannelId,
+        }),
     });
   }