Only send heartbeat when data hasn't been written in the given interval.

Author: taeold

spec/common/providers/https.spec.ts

@@ -864,7 +864,37 @@ describe("onCallHandler", () => {
         const handlerPromise = runHandler(fn, mockReq as any);
         await clock.tickAsync(11_000);
         const resp = await handlerPromise;
-        expect(resp.body).to.include(': ping\n: ping\ndata: {"result":"done"}');
+        const data = [": ping", ": ping", `data: {"result":"done"}`];
+        expect(resp.body).to.equal([...data, ""].join("\n"));
+      });
+
+      it("doesn't send heartbeat messages if user writes data", async () => {
+        const mockReq = mockRequest(
+          { message: "test heartbeat" },
+          "application/json",
+          {},
+          { accept: "text/event-stream" }
+        );
+
+        const fn = https.onCallHandler(
+          {
+            cors: { origin: true, methods: "POST" },
+            heartbeatSeconds: 5,
+          },
+          async (resp, res) => {
+            await new Promise((resolve) => setTimeout(resolve, 3_000));
+            res.write("hello");
+            await new Promise((resolve) => setTimeout(resolve, 3_000));
+            return "done";
+          },
+          "gcfv2"
+        );
+
+        const handlerPromise = runHandler(fn, mockReq as any);
+        await clock.tickAsync(10_000);
+        const resp = await handlerPromise;
+        const data = [`data: {"message":"hello"}`, `data: {"result":"done"}`];
+        expect(resp.body).to.equal([...data, ""].join("\n"));
       });
 
       it("respects null heartbeatSeconds option", async () => {
@@ -890,7 +920,7 @@ describe("onCallHandler", () => {
         const handlerPromise = runHandler(fn, mockReq as any);
         await clock.tickAsync(31_000);
         const resp = await handlerPromise;
-        expect(resp.body).to.include('data: {"result":"done"}');
+        expect(resp.body).to.equal('data: {"result":"done"}\n');
       });
     });
   });

spec/helper.ts

@@ -86,6 +86,7 @@ export function runHandler(
 
       public write(writeBody: any) {
         this.sentBody += typeof writeBody === "object" ? JSON.stringify(writeBody) : writeBody;
+        return true;
       }
 
       public end() {

src/common/providers/https.ts

@@ -747,15 +747,27 @@ function wrapOnCallHandler<Req = any, Res = any>(
     const abortController = new AbortController();
     let heartbeatInterval: NodeJS.Timeout | null = null;
 
-    const clearHeartbeatInterval = () => {
+    const clearScheduledHeartbeat = () => {
       if (heartbeatInterval) {
-        clearInterval(heartbeatInterval);
+        clearTimeout(heartbeatInterval);
         heartbeatInterval = null;
       }
     };
 
+    const scheduleHeartbeat = (heartbeatSeconds: number) => {
+      clearScheduledHeartbeat(); // Clear any existing timeout
+      if (!abortController.signal.aborted) {
+        heartbeatInterval = setTimeout(() => {
+          if (!abortController.signal.aborted) {
+            res.write(": ping\n");
+            scheduleHeartbeat(heartbeatSeconds);
+          }
+        }, heartbeatSeconds * 1000);
+      }
+    };
+
     req.on("close", () => {
-      clearHeartbeatInterval();
+      clearScheduledHeartbeat();
       abortController.abort();
     });
 
@@ -828,6 +840,12 @@ function wrapOnCallHandler<Req = any, Res = any>(
           ...context,
           data,
         };
+
+        const heartbeatSeconds =
+          options.heartbeatSeconds === undefined
+            ? DEFAULT_HEARTBEAT_SECONDS
+            : options.heartbeatSeconds;
+
         const responseProxy: CallableProxyResponse = {
           write(chunk): boolean {
             // if client doesn't accept sse-protocol, response.write() is no-op.
@@ -839,23 +857,29 @@ function wrapOnCallHandler<Req = any, Res = any>(
               return false;
             }
             const formattedData = encodeSSE({ message: chunk });
-            return res.write(formattedData);
+            const wrote = res.write(formattedData);
+            //
+            // Reset heartbeat timer after successful write
+            if (wrote && heartbeatInterval !== null && heartbeatSeconds > 0) {
+              scheduleHeartbeat(heartbeatSeconds);
+            }
+            return wrote;
           },
           acceptsStreaming,
           signal: abortController.signal,
         };
         if (acceptsStreaming) {
           // SSE always responds with 200
           res.status(200);
-          const heartbeatSeconds = options.heartbeatSeconds ?? DEFAULT_HEARTBEAT_SECONDS;
+
           if (heartbeatSeconds !== null && heartbeatSeconds > 0) {
-            heartbeatInterval = setInterval(() => res.write(": ping\n"), heartbeatSeconds * 1000);
+            scheduleHeartbeat(heartbeatSeconds);
           }
         }
         // For some reason the type system isn't picking up that the handler
         // is a one argument function.
         result = await (handler as any)(arg, responseProxy);
-        clearHeartbeatInterval();
+        clearScheduledHeartbeat();
       }
 
       if (!abortController.signal.aborted) {
@@ -894,7 +918,7 @@ function wrapOnCallHandler<Req = any, Res = any>(
         res.end();
       }
     } finally {
-      clearHeartbeatInterval();
+      clearScheduledHeartbeat();
     }
   };
 }