Improve the stream throttling

Author: ericallam

packages/react-hooks/src/hooks/useRealtime.ts

@@ -2,9 +2,9 @@
 
 import { AnyTask, ApiClient, InferRunTypes, RealtimeRun } from "@trigger.dev/core/v3";
 import { useCallback, useEffect, useId, useRef, useState } from "react";
-import { throttle } from "../utils/throttle.js";
 import { KeyedMutator, useSWR } from "../utils/trigger-swr.js";
 import { useApiClient, UseApiClientOptions } from "./useApiClient.js";
+import { createThrottledQueue } from "../utils/throttle.js";
 
 export type UseRealtimeRunOptions = UseApiClientOptions & {
   id?: string;
@@ -78,12 +78,7 @@ export function useRealtimeRun<TTask extends AnyTask>(
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
 
-      await processRealtimeRun(
-        runId,
-        apiClient,
-        throttle(mutateRun, options?.experimental_throttleInMs),
-        abortControllerRef
-      );
+      await processRealtimeRun(runId, apiClient, mutateRun, abortControllerRef);
     } catch (err) {
       // Ignore abort errors as they are expected.
       if ((err as any).name === "AbortError") {
@@ -199,10 +194,11 @@ export function useRealtimeRunWithStreams<
       await processRealtimeRunWithStreams(
         runId,
         apiClient,
-        throttle(mutateRun, options?.experimental_throttleInMs),
-        throttle(mutateStreams, options?.experimental_throttleInMs),
+        mutateRun,
+        mutateStreams,
         streamsRef,
-        abortControllerRef
+        abortControllerRef,
+        options?.experimental_throttleInMs
       );
     } catch (err) {
       // Ignore abort errors as they are expected.
@@ -285,13 +281,7 @@ export function useRealtimeRunsWithTag<TTask extends AnyTask>(
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
 
-      await processRealtimeRunsWithTag(
-        tag,
-        apiClient,
-        throttle(mutateRuns, options?.experimental_throttleInMs),
-        runsRef,
-        abortControllerRef
-      );
+      await processRealtimeRunsWithTag(tag, apiClient, mutateRuns, runsRef, abortControllerRef);
     } catch (err) {
       // Ignore abort errors as they are expected.
       if ((err as any).name === "AbortError") {
@@ -372,13 +362,7 @@ export function useRealtimeBatch<TTask extends AnyTask>(
       const abortController = new AbortController();
       abortControllerRef.current = abortController;
 
-      await processRealtimeBatch(
-        batchId,
-        apiClient,
-        throttle(mutateRuns, options?.experimental_throttleInMs),
-        runsRef,
-        abortControllerRef
-      );
+      await processRealtimeBatch(batchId, apiClient, mutateRuns, runsRef, abortControllerRef);
     } catch (err) {
       // Ignore abort errors as they are expected.
       if ((err as any).name === "AbortError") {
@@ -486,23 +470,51 @@ async function processRealtimeRunWithStreams<
   mutateRunData: KeyedMutator<RealtimeRun<TTask>>,
   mutateStreamData: KeyedMutator<StreamResults<TStreams>>,
   existingDataRef: React.MutableRefObject<StreamResults<TStreams>>,
-  abortControllerRef: React.MutableRefObject<AbortController | null>
+  abortControllerRef: React.MutableRefObject<AbortController | null>,
+  throttleInMs?: number
 ) {
   const subscription = apiClient.subscribeToRun<InferRunTypes<TTask>>(runId, {
     signal: abortControllerRef.current?.signal,
   });
 
+  type StreamUpdate = {
+    type: keyof TStreams;
+    chunk: any;
+  };
+
+  const streamQueue = createThrottledQueue<StreamUpdate>(async (updates) => {
+    const nextStreamData = { ...existingDataRef.current };
+
+    // Group updates by type
+    const updatesByType = updates.reduce(
+      (acc, update) => {
+        if (!acc[update.type]) {
+          acc[update.type] = [];
+        }
+        acc[update.type].push(update.chunk);
+        return acc;
+      },
+      {} as Record<keyof TStreams, any[]>
+    );
+
+    // Apply all updates
+    for (const [type, chunks] of Object.entries(updatesByType)) {
+      // @ts-ignore
+      nextStreamData[type] = [...(existingDataRef.current[type] || []), ...chunks];
+    }
+
+    await mutateStreamData(nextStreamData);
+  }, throttleInMs);
+
   for await (const part of subscription.withStreams<TStreams>()) {
     if (part.type === "run") {
       mutateRunData(part.run);
     } else {
-      const nextStreamData = {
-        ...existingDataRef.current,
+      streamQueue.add({
+        type: part.type,
         // @ts-ignore
-        [part.type]: [...(existingDataRef.current[part.type] || []), part.chunk],
-      };
-
-      mutateStreamData(nextStreamData);
+        chunk: part.chunk,
+      });
     }
   }
 }

packages/react-hooks/src/utils/throttle.ts

@@ -1,5 +1,58 @@
-import throttleFunction from "throttleit";
+// Reusable throttle utility
+export type ThrottledQueue<T> = {
+  add: (item: T) => void;
+  flush: () => Promise<void>;
+  isEmpty: () => boolean;
+};
 
-export function throttle<T extends (...args: any[]) => any>(fn: T, waitMs: number | undefined): T {
-  return waitMs != null ? throttleFunction(fn, waitMs) : fn;
+export function createThrottledQueue<T>(
+  onFlush: (items: T[]) => Promise<void>,
+  throttleInMs?: number
+): ThrottledQueue<T> {
+  let queue: T[] = [];
+  let lastFlushTime = 0;
+  let flushPromise: Promise<void> | null = null;
+
+  const scheduleFlush = async () => {
+    // If no throttle specified or there's already a flush in progress, return
+    if (!throttleInMs) {
+      // Immediately flush when no throttling is specified
+      const itemsToFlush = [...queue];
+      queue = [];
+      await onFlush(itemsToFlush);
+      return;
+    }
+
+    if (queue.length === 0 || flushPromise) return;
+
+    const now = Date.now();
+    const timeUntilNextFlush = Math.max(0, lastFlushTime + throttleInMs - now);
+
+    if (timeUntilNextFlush === 0) {
+      const itemsToFlush = [...queue];
+      queue = [];
+      lastFlushTime = now;
+      flushPromise = onFlush(itemsToFlush).finally(() => {
+        flushPromise = null;
+        // Check if more items accumulated during flush
+        scheduleFlush();
+      });
+    } else {
+      setTimeout(scheduleFlush, timeUntilNextFlush);
+    }
+  };
+
+  return {
+    add: (item: T) => {
+      queue.push(item);
+      scheduleFlush();
+    },
+    flush: async () => {
+      if (queue.length === 0) return;
+      const itemsToFlush = [...queue];
+      queue = [];
+      await onFlush(itemsToFlush);
+    },
+    isEmpty: () => queue.length === 0,
+  };
 }

references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx

@@ -23,6 +23,7 @@ export default function TriggerButton({ accessToken }: { accessToken: string })
   >("openai-streaming", {
     accessToken,
     baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
+    experimental_throttleInMs: 100,
   });
 
   const openWeatherReport = useCallback(() => {