From 809b02cb263d60a9c9637802cc6760acecefaf7d Mon Sep 17 00:00:00 2001
From: Manav Rathi <manav@mrmr.io>
Date: Wed, 3 Jul 2024 20:33:39 +0530
Subject: [PATCH] Last trace

---
 web/apps/photos/src/services/searchService.ts |  4 +---
 web/packages/new/photos/services/ml/index.ts  | 19 ++++++++-------
 web/packages/new/photos/services/ml/worker.ts | 24 ++++++++++++-------
 3 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/web/apps/photos/src/services/searchService.ts b/web/apps/photos/src/services/searchService.ts
index 9bf34ee57a..98e5c5e842 100644
--- a/web/apps/photos/src/services/searchService.ts
+++ b/web/apps/photos/src/services/searchService.ts
@@ -1,6 +1,5 @@
 import { FILE_TYPE } from "@/media/file-type";
 import { faceIndexingStatus, isMLEnabled } from "@/new/photos/services/ml";
-import mlWorkManager from "@/new/photos/services/ml/mlWorkManager";
 import type { Person } from "@/new/photos/services/ml/people";
 import { EnteFile } from "@/new/photos/types/file";
 import { isDesktop } from "@/next/app";
@@ -178,8 +177,7 @@ export async function getAllPeopleSuggestion(): Promise<Array<Suggestion>> {
 
 export async function getIndexStatusSuggestion(): Promise<Suggestion> {
     try {
-        const isSyncing = mlWorkManager.isSyncing;
-        const indexStatus = await faceIndexingStatus(isSyncing);
+        const indexStatus = await faceIndexingStatus();
 
         let label: string;
         switch (indexStatus.phase) {
diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts
index 685a056ecc..63a6aa20e7 100644
--- a/web/packages/new/photos/services/ml/index.ts
+++ b/web/packages/new/photos/services/ml/index.ts
@@ -201,18 +201,21 @@ export interface FaceIndexingStatus {
     nTotalFiles: number;
 }
 
-export const faceIndexingStatus = async (
-    isSyncing: boolean,
-): Promise<FaceIndexingStatus> => {
+/**
+ * Return the current state of the face indexing pipeline.
+ *
+ * Precondition: ML must be enabled.
+ */
+export const faceIndexingStatus = async (): Promise<FaceIndexingStatus> => {
+    if (!isMLEnabled())
+        throw new Error("Cannot get indexing status when ML is not enabled");
+
     const { indexedCount, indexableCount } = await indexedAndIndexableCounts();
+    const isIndexing = await (await worker()).isIndexing();
 
     let phase: FaceIndexingStatus["phase"];
     if (indexableCount > 0) {
-        if (!isSyncing) {
-            phase = "scheduled";
-        } else {
-            phase = "indexing";
-        }
+        phase = !isIndexing ? "scheduled" : "indexing";
     } else {
         phase = "done";
     }
diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts
index c8c5752fa4..1ad874d74c 100644
--- a/web/packages/new/photos/services/ml/worker.ts
+++ b/web/packages/new/photos/services/ml/worker.ts
@@ -47,7 +47,7 @@ export class MLWorker {
     private userAgent: string | undefined;
     private shouldSync = false;
     private liveQ: EnteFile[] = [];
-    private haveStarted = false;
+    private state: "idle" | "pull" | "indexing" = "idle";
     private idleTimeout: ReturnType<typeof setTimeout> | undefined;
     private idleDuration = idleDurationStart; /* unit: seconds */
 
@@ -84,13 +84,9 @@ export class MLWorker {
 
     /** Invoked in response to external events. */
     private wakeUp() {
-        if (!this.haveStarted) {
-            // First time something happened.
-            this.haveStarted = true;
-            void this.tick();
-        } else if (this.idleTimeout) {
+        if (this.state == "idle") {
             // Currently paused. Get back to work.
-            clearTimeout(this.idleTimeout);
+            if (this.idleTimeout) clearTimeout(this.idleTimeout);
             this.idleTimeout = undefined;
             void this.tick();
         } else {
@@ -125,9 +121,17 @@ export class MLWorker {
         }
     }
 
+    /**
+     * Return true if we're currently indexing.
+     */
+    isIndexing() {
+        return this.state == "indexing";
+    }
+
     private async tick() {
         log.debug(() => ({
             t: "ml-tick",
+            state: this.state,
             shouldSync: this.shouldSync,
             liveQ: this.liveQ,
             idleDuration: this.idleDuration,
@@ -138,6 +142,7 @@ export class MLWorker {
         // If we've been asked to sync, do that irrespective of anything else.
         if (this.shouldSync) {
             this.shouldSync = false;
+            this.state = "pull";
             void pull().then((didPull) => {
                 // Reset the idle duration if we did pull something.
                 if (didPull) this.idleDuration = idleDurationStart;
@@ -150,6 +155,7 @@ export class MLWorker {
 
         const liveQ = this.liveQ;
         this.liveQ = [];
+        this.state = "indexing";
         const allSuccess = await indexNextBatch(ensure(this.userAgent), liveQ);
         if (allSuccess) {
             // Everything is running smoothly. Reset the idle duration.
@@ -169,6 +175,7 @@ export class MLWorker {
         // So in all cases, we pause for exponentially longer durations of time
         // (limited to some maximum).
 
+        this.state = "idle";
         this.idleDuration = Math.min(this.idleDuration * 2, idleDurationMax);
         this.idleTimeout = setTimeout(scheduleTick, this.idleDuration * 1000);
     }
@@ -208,8 +215,7 @@ const indexNextBatch = async (userAgent: string, liveQ: EnteFile[]) => {
     for (const file of files) {
         try {
             await index(file, undefined, userAgent);
-            // Let the event loop run so that other events (like onUpload) can
-            // be acknowledged and noted down.
+            // Possibly unnecessary, but let us drain the microtask queue.
             await wait(0);
         } catch {
             allSuccess = false;