fix: bound end-to-end refresh discovery latency

src/common/telemetry/errorClassifier.ts

@@ -1,6 +1,7 @@
 import { CancellationError } from 'vscode';
 import * as rpc from 'vscode-jsonrpc/node';
-import { RpcTimeoutError } from '../../managers/common/nativePythonFinder';
+import { RefreshBudgetExceededError, RpcTimeoutError } from '../../managers/common/nativePythonFinder';
+import { QueueTaskExpiredError } from '../utils/workerPool';
 import { BaseError } from '../errors/types';
 
 export type DiscoveryErrorType =
@@ -49,6 +50,11 @@ export function classifyError(ex: unknown): DiscoveryErrorType {
         }
     }
 
+    // Queue-expiry and refresh-budget errors are time-budget exhaustions → generic RPC timeout category.
+    if (ex instanceof QueueTaskExpiredError || ex instanceof RefreshBudgetExceededError) {
+        return 'rpc_timeout';
+    }
+
     // JSON-RPC connection errors (e.g., PET process died mid-request, connection closed/disposed)
     if (ex instanceof rpc.ConnectionError) {
         return 'connection_error';

src/common/utils/workerPool.ts

@@ -4,6 +4,14 @@
 import { traceError } from '../logging';
 import { createDeferred, Deferred } from './deferred';
 
+/** Rejects a queued work item that expired before a worker could dequeue it. */
+export class QueueTaskExpiredError extends Error {
+    constructor(expiresInMs: number) {
+        super(`Queued task expired after ${expiresInMs}ms before it could start`);
+        this.name = this.constructor.name;
+    }
+}
+
 interface Worker {
     /**
      * Start processing of items.
@@ -23,8 +31,15 @@ type PostResult<T, R> = (item: T, result?: R, err?: Error) => void;
 
 interface IWorkItem<T> {
     item: T;
+    running: boolean;
+    expired: boolean;
+    expiryTimer?: ReturnType<typeof setTimeout>;
+    expiresAt?: number;
+    expiresInMs?: number;
 }
 
+export type QueueClock = () => number;
+
 export enum QueuePosition {
     back,
     front,
@@ -36,9 +51,11 @@ export interface WorkerPool<T, R> extends Worker {
      * @method addToQueue
      * @param {T} item: Item to process
      * @param {QueuePosition} position: Add items to the front or back of the queue.
+     * @param {number} expiresInMs: Optional. When set, a still-queued item is rejected with
+     *        {@link QueueTaskExpiredError} after this many ms and never runs; omit to queue unbounded.
      * @returns A promise that when resolved gets the result from running the worker function.
      */
-    addToQueue(item: T, position?: QueuePosition): Promise<R>;
+    addToQueue(item: T, position?: QueuePosition, expiresInMs?: number): Promise<R>;
 }
 
 class WorkerImpl<T, R> implements Worker {
@@ -76,14 +93,17 @@ class WorkerImpl<T, R> implements Worker {
 class WorkQueue<T, R> {
     private readonly items: IWorkItem<T>[] = [];
     private readonly results: Map<IWorkItem<T>, Deferred<R>> = new Map();
-    public add(item: T, position?: QueuePosition): Promise<R> {
+
+    public constructor(private readonly now: QueueClock = Date.now) {}
+
+    public add(item: T, position?: QueuePosition, expiresInMs?: number): Promise<R> {
         // Wrap the user provided item in a wrapper object. This will allow us to track multiple
         // submissions of the same item. For example, addToQueue(2), addToQueue(2). If we did not
         // wrap this, then from the map both submissions will look the same. Since this is a generic
         // worker pool, we do not know if we can resolve both using the same promise. So, a better
         // approach is to ensure each gets a unique promise, and let the worker function figure out
         // how to handle repeat submissions.
-        const workItem: IWorkItem<T> = { item };
+        const workItem: IWorkItem<T> = { item, running: false, expired: false };
         if (position === QueuePosition.front) {
             this.items.unshift(workItem);
         } else {
@@ -96,29 +116,88 @@ class WorkQueue<T, R> {
         const deferred = createDeferred<R>();
         this.results.set(workItem, deferred);
 
+        if (expiresInMs !== undefined) {
+            workItem.expiresInMs = expiresInMs;
+            workItem.expiresAt = this.now() + expiresInMs;
+            workItem.expiryTimer = setTimeout(() => this.expire(workItem), Math.max(0, expiresInMs));
+        }
+
         return deferred.promise;
     }
 
+    private clearExpiry(workItem: IWorkItem<T>): void {
+        if (workItem.expiryTimer !== undefined) {
+            clearTimeout(workItem.expiryTimer);
+            workItem.expiryTimer = undefined;
+        }
+    }
+
+    private settleExpired(workItem: IWorkItem<T>): void {
+        this.clearExpiry(workItem);
+        if (workItem.running || workItem.expired) {
+            return;
+        }
+        workItem.expired = true;
+        const deferred = this.results.get(workItem);
+        if (deferred !== undefined) {
+            this.results.delete(workItem);
+            deferred.reject(new QueueTaskExpiredError(workItem.expiresInMs ?? 0));
+        }
+    }
+
+    private expire(workItem: IWorkItem<T>): void {
+        this.clearExpiry(workItem);
+        if (workItem.running || workItem.expired) {
+            return;
+        }
+        const index = this.items.indexOf(workItem);
+        if (index < 0) {
+            return;
+        }
+        this.items.splice(index, 1);
+        this.settleExpired(workItem);
+    }
+
     public completed(workItem: IWorkItem<T>, result?: R, error?: Error): void {
+        this.clearExpiry(workItem);
         const deferred = this.results.get(workItem);
         if (deferred !== undefined) {
             this.results.delete(workItem);
             if (error !== undefined) {
                 deferred.reject(error);
+            } else {
+                deferred.resolve(result);
             }
-            deferred.resolve(result);
         }
     }
 
     public next(): IWorkItem<T> | undefined {
-        return this.items.shift();
+        let workItem = this.items.shift();
+        while (workItem !== undefined) {
+            if (workItem.expired) {
+                workItem = this.items.shift();
+                continue;
+            }
+            // Absolute-deadline recheck: never start an item past its deadline even if the timer hasn't fired.
+            if (workItem.expiresAt !== undefined && this.now() >= workItem.expiresAt) {
+                this.settleExpired(workItem);
+                workItem = this.items.shift();
+                continue;
+            }
+            workItem.running = true;
+            this.clearExpiry(workItem);
+            return workItem;
+        }
+        return undefined;
     }
 
     public clear(): void {
         this.results.forEach((v: Deferred<R>, k: IWorkItem<T>, map: Map<IWorkItem<T>, Deferred<R>>) => {
+            this.clearExpiry(k);
             v.reject(Error('Queue stopped processing'));
             map.delete(k);
         });
+        this.items.length = 0;
     }
 }
 
@@ -131,7 +210,7 @@ class WorkerPoolImpl<T, R> implements WorkerPool<T, R> {
     private readonly waitingWorkersUnblockQueue: { unblock(w: IWorkItem<T>): void; stop(): void }[] = [];
 
     // A collection that manages the work items.
-    private readonly queue = new WorkQueue<T, R>();
+    private readonly queue: WorkQueue<T, R>;
 
     // State of the pool manages via stop(), start()
     private stopProcessing = false;
@@ -140,16 +219,19 @@ class WorkerPoolImpl<T, R> implements WorkerPool<T, R> {
         private readonly workerFunc: WorkFunc<T, R>,
         private readonly numWorkers: number = 2,
         private readonly name: string = 'Worker',
-    ) {}
+        now?: QueueClock,
+    ) {
+        this.queue = new WorkQueue<T, R>(now);
+    }
 
-    public addToQueue(item: T, position?: QueuePosition): Promise<R> {
+    public addToQueue(item: T, position?: QueuePosition, expiresInMs?: number): Promise<R> {
         if (this.stopProcessing) {
             throw Error('Queue is stopped');
         }
 
         // This promise when resolved should return the processed result of the item
         // being added to the queue.
-        const deferred = this.queue.add(item, position);
+        const deferred = this.queue.add(item, position, expiresInMs);
 
         const worker = this.waitingWorkersUnblockQueue.shift();
         if (worker) {
@@ -160,9 +242,8 @@ class WorkerPoolImpl<T, R> implements WorkerPool<T, R> {
                 // and give the worker the newly added item.
                 worker.unblock(workItem);
             } else {
-                // Something is wrong, we should not be here. we just added an item to
-                // the queue. It should not be empty.
-                traceError('Work queue was empty immediately after adding item.');
+                // next() dropped the just-added item as already expired; re-park the worker.
+                this.waitingWorkersUnblockQueue.unshift(worker);
             }
         }
 
@@ -243,8 +324,9 @@ export function createRunningWorkerPool<T, R>(
     workerFunc: WorkFunc<T, R>,
     numWorkers?: number,
     name?: string,
+    now?: QueueClock,
 ): WorkerPool<T, R> {
-    const pool = new WorkerPoolImpl<T, R>(workerFunc, numWorkers, name);
+    const pool = new WorkerPoolImpl<T, R>(workerFunc, numWorkers, name, now);
     pool.start();
     return pool;
 }