[status] Improve reconnection logic for SSE implementation

Using an explicit keepalive mechanism allows us to detect when the
connection has been broken and re-establish it.

Bug: b/542641998
Change-Id: I912ea294dbefa6df2def52749bfd427b556f158a
Reviewed-on: https://skia-review.googlesource.com/c/buildbot/+/1329896
Commit-Queue: Eric Boren <borenet@google.com>
Reviewed-by: Alexis Cruz-Ayala <alexisdavidc@google.com>
diff --git a/status/go/sse/sse.go b/status/go/sse/sse.go
index dc3440a..fe11d77 100644
--- a/status/go/sse/sse.go
+++ b/status/go/sse/sse.go
@@ -30,6 +30,12 @@
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
+const (
+	commitPollInterval  = 10 * time.Second
+	commentPollInterval = 10 * time.Second
+	keepaliveInterval   = 30 * time.Second
+)
+
 type Config struct {
 	Repos                repograph.Map
 	TaskDb               db.RemoteDB
@@ -267,10 +273,12 @@
 }
 
 func (s *SSEServer) broadcastLoop(ctx context.Context) {
-	commitsTicker := time.NewTicker(10 * time.Second)
+	commitsTicker := time.NewTicker(commitPollInterval)
 	defer commitsTicker.Stop()
-	commentsTicker := time.NewTicker(10 * time.Second)
+	commentsTicker := time.NewTicker(commentPollInterval)
 	defer commentsTicker.Stop()
+	keepaliveTicker := time.NewTicker(keepaliveInterval)
+	defer keepaliveTicker.Stop()
 
 	for {
 		select {
@@ -282,6 +290,8 @@
 			s.checkForNewCommits()
 		case <-commentsTicker.C:
 			s.checkForNewComments(ctx)
+		case <-keepaliveTicker.C:
+			s.broadcastKeepalives()
 		}
 	}
 }
@@ -298,6 +308,31 @@
 	}
 }
 
+func (s *SSEServer) broadcastKeepalives() {
+	s.clientsMtx.Lock()
+	var eg errgroup.Group
+	for client := range s.clients {
+		client := client // https://golang.org/doc/faq#closures_and_goroutines
+		eg.Go(func() error {
+			select {
+			case <-client.Done():
+				s.unregister(client)
+				return nil
+			default:
+			}
+			if err := client.SendKeepalive(); err != nil {
+				s.unregister(client)
+				return skerr.Wrapf(err, "failed to send keepalive; unregistered the client")
+			}
+			return nil
+		})
+	}
+	s.clientsMtx.Unlock()
+	if err := eg.Wait(); err != nil {
+		sklog.Error(err)
+	}
+}
+
 func (s *SSEServer) checkForNewCommits() {
 	s.ancestryCacheMtx.Lock()
 	defer s.ancestryCacheMtx.Unlock()
@@ -662,6 +697,20 @@
 	return c.branchRegex == nil || c.branchRegex.MatchString(branchName)
 }
 
+// SendKeepalive sends a named "keepalive" event keeping the connection alive.
+func (s *clientStream) SendKeepalive() error {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+	_, err := fmt.Fprint(s.w, "event: keepalive\ndata: {}\n\n")
+	if err != nil {
+		return skerr.Wrapf(err, "failed to send SSE keepalive")
+	}
+	if err := s.w.Flush(); err != nil {
+		return skerr.Wrapf(err, "failed to flush SSE stream")
+	}
+	return nil
+}
+
 // updateAncestryCache updates the given sub-map of the ancestryCache. The
 // caller MUST hold a lock on ancestryCache.
 func updateAncestryCache(cache map[string][]string, repo *repograph.Graph, oldBranchHeads []*git.Branch) {
diff --git a/status/modules/commits-table-experimental-sk/commits-table-experimental-sk.ts b/status/modules/commits-table-experimental-sk/commits-table-experimental-sk.ts
index 90acb94..8ac7f58 100644
--- a/status/modules/commits-table-experimental-sk/commits-table-experimental-sk.ts
+++ b/status/modules/commits-table-experimental-sk/commits-table-experimental-sk.ts
@@ -66,6 +66,10 @@
 const TASK_STATUS_FAILURE = 'FAILURE';
 const TASK_STATUS_MISHAP = 'MISHAP';
 
+export const KEEPALIVE_TIMEOUT_MS = 2 * 60 * 1000;
+export const KEEPALIVE_CHECK_INTERVAL_MS = 5000;
+export const RECONNECT_DELAY_MS = 5000;
+
 // Makes some maps more self-documenting.
 export type CommitHash = string;
 export type TaskSpec = string;
@@ -207,10 +211,72 @@
 
   private taskSearch: string = '';
 
+  private keepaliveReconnectTimeout?: number;
+
+  private lastKeepaliveTime: number = Date.now();
+
+  private keepaliveInterval?: number;
+
   get lastLoaded() {
     return this._lastLoaded;
   }
 
+  private startKeepaliveCheck(onUpdate: () => void) {
+    this.lastKeepaliveTime = Date.now();
+    this.stopKeepaliveCheck();
+    this.keepaliveInterval = window.setInterval(() => {
+      const elapsed = Date.now() - this.lastKeepaliveTime;
+      if (elapsed > KEEPALIVE_TIMEOUT_MS) {
+        console.warn(
+          `No SSE activity for ${KEEPALIVE_TIMEOUT_MS / 1000}s. Connection has stalled. Reconnecting...`
+        );
+        this.resetEventSource(
+          this.repo,
+          this.numCommits,
+          this.branchFilter,
+          this.cursor,
+          this.taskFilter,
+          this.taskSearch,
+          onUpdate,
+          true
+        );
+      }
+    }, KEEPALIVE_CHECK_INTERVAL_MS);
+  }
+
+  private stopKeepaliveCheck() {
+    if (this.keepaliveInterval) {
+      window.clearInterval(this.keepaliveInterval);
+      this.keepaliveInterval = undefined;
+    }
+  }
+
+  private scheduleKeepaliveReconnect(onUpdate: () => void) {
+    // Debounce concurrent connection triggers by clearing any existing
+    // scheduled reconnect.
+    if (this.keepaliveReconnectTimeout) {
+      window.clearTimeout(this.keepaliveReconnectTimeout);
+    }
+    // Delay reconnection by RECONNECT_DELAY_MS. This ensures that we don't
+    // overwhelm the server with requests or create a tight client-side loop
+    // which wastes CPU unnecessarily.
+    this.keepaliveReconnectTimeout = window.setTimeout(() => {
+      if (this.eventSource && this.eventSource.readyState === EventSource.CLOSED) {
+        console.log('EventSource is CLOSED. Forcing reconnection...');
+        this.resetEventSource(
+          this.repo,
+          this.numCommits,
+          this.branchFilter,
+          this.cursor,
+          this.taskFilter,
+          this.taskSearch,
+          onUpdate,
+          true
+        );
+      }
+    }, RECONNECT_DELAY_MS);
+  }
+
   resetEventSource(
     repo: string,
     numCommits: number,
@@ -218,9 +284,11 @@
     cursor: string,
     taskFilter: TaskFilter,
     taskSearch: string,
-    onUpdate: () => void
+    onUpdate: () => void,
+    force: boolean = false
   ): Promise<void> {
     if (
+      force ||
       this.repo !== repo ||
       this.numCommits !== numCommits ||
       this.branchFilter !== branchFilter ||
@@ -228,6 +296,11 @@
       this.taskFilter !== taskFilter ||
       (taskFilter === 'Search' && this.taskSearch !== taskSearch)
     ) {
+      if (this.keepaliveReconnectTimeout) {
+        window.clearTimeout(this.keepaliveReconnectTimeout);
+        this.keepaliveReconnectTimeout = undefined;
+      }
+      this.stopKeepaliveCheck();
       if (this.eventSource) {
         this.eventSource.close();
         this.eventSource = null;
@@ -257,6 +330,10 @@
       const url = `/sse?${queryParams.toString()}`;
       try {
         this.eventSource = new EventSource(url);
+        this.eventSource.addEventListener('keepalive', () => {
+          console.log('Received SSE keepalive event');
+          this.lastKeepaliveTime = Date.now();
+        });
       } catch (err) {
         console.error('Failed to create event source');
         console.error(err);
@@ -266,8 +343,14 @@
       return new Promise<void>((resolve, reject) => {
         let firstMessage = true;
         let disconnected = false;
+        this.startKeepaliveCheck(onUpdate);
         this.eventSource!.onmessage = (event) => {
           console.log('Received message:');
+          this.lastKeepaliveTime = Date.now();
+          if (this.keepaliveReconnectTimeout) {
+            window.clearTimeout(this.keepaliveReconnectTimeout);
+            this.keepaliveReconnectTimeout = undefined;
+          }
           try {
             const json: GetIncrementalCommitsResponse = JSON.parse(event.data);
             console.log(json);
@@ -304,6 +387,7 @@
             firstMessage = false;
             reject(err);
           }
+          this.scheduleKeepaliveReconnect(onUpdate);
         };
       });
     }
@@ -612,8 +696,6 @@
 
   private mishapTasks: Array<Task> = [];
 
-  private refreshHandle?: number;
-
   private requestLimiter: RequestLimiter = new RequestLimiter();
 
   private stateHasChanged: () => void = () => {};
@@ -642,7 +724,7 @@
           <select
             id="repoSelector"
             @change=${(e: Event) => {
-              el.repo = (e.target as any).value;
+              el.repo = (e.target as HTMLSelectElement).value;
             }}>
             ${repos().map((r) => html`<option value=${r}>${r}</option>`)}
           </select>
@@ -1425,8 +1507,6 @@
       return;
     }
     const numCommits = Number((<HTMLInputElement>$$('#commitsInput', this)).value);
-    window.clearTimeout(this.refreshHandle);
-    this.refreshHandle = undefined;
     this.dispatchEvent(new CustomEvent('begin-task', { bubbles: true }));
     this.loading = true;
     this.draw();
diff --git a/status/modules/commits-table-experimental-sk/commits-table-experimental-sk_test.ts b/status/modules/commits-table-experimental-sk/commits-table-experimental-sk_test.ts
index b8e42d3..94947d1 100644
--- a/status/modules/commits-table-experimental-sk/commits-table-experimental-sk_test.ts
+++ b/status/modules/commits-table-experimental-sk/commits-table-experimental-sk_test.ts
@@ -20,7 +20,12 @@
   incrementalResponse1,
 } from '../rpc-mock/test_data';
 import { GetIncrementalCommitsRequest, GetIncrementalCommitsResponse } from '../rpc';
-import { CommitsTableExperimentalSk } from './commits-table-experimental-sk';
+import {
+  CommitsTableExperimentalSk,
+  KEEPALIVE_TIMEOUT_MS,
+  KEEPALIVE_CHECK_INTERVAL_MS,
+  RECONNECT_DELAY_MS,
+} from './commits-table-experimental-sk';
 import { MockStatusService, SetupMocks } from '../rpc-mock';
 import { SetTestSettings } from '../settings';
 
@@ -57,10 +62,27 @@
 
       url: string;
 
+      readyState: number = EventSource.OPEN;
+
       onmessage: ((ev: MessageEvent) => void) | null = null;
 
       onerror: ((ev: Event) => void) | null = null;
 
+      private listeners: Record<string, Array<() => void>> = {};
+
+      addEventListener(type: string, listener: () => void) {
+        if (!this.listeners[type]) {
+          this.listeners[type] = [];
+        }
+        this.listeners[type].push(listener);
+      }
+
+      triggerEvent(type: string) {
+        if (this.listeners[type]) {
+          this.listeners[type].forEach((l) => l());
+        }
+      }
+
       constructor(url: string) {
         this.url = url;
         MockEventSource.activeInstance = this;
@@ -94,6 +116,7 @@
       }
 
       close() {
+        this.readyState = EventSource.CLOSED;
         if (MockEventSource.activeInstance === this) {
           MockEventSource.activeInstance = null;
         }
@@ -687,5 +710,99 @@
         commitsData.comments.get(commentTask.commit)!.get(commentTask.taskSpecName)![0]
       ).to.deep.include({ message: commentTask.message });
     });
+
+    it('automatically reconnects when EventSource errors and goes CLOSED', async () => {
+      const _ = await setupWithResponse(incrementalResponse0);
+
+      // Get the active MockEventSource instance.
+      const firstES = (window.EventSource as any).activeInstance;
+      expect(firstES).to.not.be.null;
+
+      // Mock window.setTimeout to trigger immediately for the reconnection logic.
+      const originalTimeout = window.setTimeout;
+      let triggerImmediately: (() => void) | null = null;
+      (window as any).setTimeout = (cb: () => void, delay: number) => {
+        if (delay === RECONNECT_DELAY_MS) {
+          triggerImmediately = cb;
+          return 0;
+        }
+        return originalTimeout(cb, delay);
+      };
+
+      try {
+        // Set up the mock expectation for the reconnection query.
+        mocks.expectGetIncrementalCommits(incrementalResponse0);
+
+        // Simulate an error which closes the connection.
+        firstES.readyState = EventSource.CLOSED;
+        if (firstES.onerror) {
+          firstES.onerror(new Event('error'));
+        }
+
+        // Trigger the scheduled reconnection instantly.
+        expect(triggerImmediately).to.not.be.null;
+        triggerImmediately!();
+
+        // Check that a new EventSource was created.
+        const secondES = (window.EventSource as any).activeInstance;
+        expect(secondES).to.not.be.null;
+        expect(secondES).to.not.equal(firstES);
+      } finally {
+        window.setTimeout = originalTimeout;
+      }
+    });
+
+    it('automatically reconnects when no keepalive is received for over the timeout threshold', async () => {
+      const table = await setupWithResponse(incrementalResponse0);
+      const data = (table as any).data;
+
+      // Get the active MockEventSource instance.
+      const firstES = (window.EventSource as any).activeInstance;
+      expect(firstES).to.not.be.null;
+
+      // Mock window.setInterval to trigger instantly for our keepalive check.
+      const originalInterval = window.setInterval;
+      const originalTimeout = window.setTimeout;
+      let keepaliveCallback: (() => void) | null = null;
+
+      (window as any).setInterval = (cb: () => void, delay: number) => {
+        if (delay === KEEPALIVE_CHECK_INTERVAL_MS) {
+          keepaliveCallback = cb;
+          return 123; // Dummy ID
+        }
+        return originalInterval(cb, delay);
+      };
+
+      (window as any).setTimeout = (cb: () => void, delay: number) => {
+        if (delay === RECONNECT_DELAY_MS) {
+          return 456; // Dummy ID
+        }
+        return originalTimeout(cb, delay);
+      };
+
+      try {
+        // Start the keepalive check so our interval mock grabs its callback.
+        data.startKeepaliveCheck(() => {});
+        expect(keepaliveCallback).to.not.be.null;
+
+        // Set up the mock expectation for the reconnection query.
+        mocks.expectGetIncrementalCommits(incrementalResponse0);
+
+        // Advance fake time: set lastKeepaliveTime to > KEEPALIVE_TIMEOUT_MS ago.
+        data.lastKeepaliveTime = Date.now() - (KEEPALIVE_TIMEOUT_MS + 1000);
+
+        // Trigger the keepalive check tick.
+        keepaliveCallback!();
+
+        // Check that a new EventSource was created.
+        const secondES = (window.EventSource as any).activeInstance;
+        expect(secondES).to.not.be.null;
+        expect(secondES).to.not.equal(firstES);
+      } finally {
+        window.setInterval = originalInterval;
+        window.setTimeout = originalTimeout;
+        data.stopKeepaliveCheck();
+      }
+    });
   });
 });