Simplify close-ordering fix

Tighten the previous commit to the minimal diff: guard the existing
unregisterContext call instead of restructuring the function. Restores
the unchanged comment on TelemetryClient.unregisterContext.

Co-authored-by: Isaac

Author: samikshya-db

lib/telemetry/TelemetryClient.ts

@@ -169,10 +169,8 @@ class TelemetryClient implements IClientContext {
 
   /**
    * Remove a `DBSQLClient`'s context from the pool. Called by
-   * `TelemetryClientProvider.releaseClient` in the multi-refcount case so the
-   * exporter doesn't keep trying to use a closing context. Deliberately
-   * NOT called on the last refcount release: `close()` needs the snapshot
-   * pair to resolve auth/connection providers for the final flush.
+   * `TelemetryClientProvider.releaseClient` before refcount decrement so the
+   * exporter doesn't keep trying to use a closed context.
    *
    * Uses the cached snapshot pair (`context`, `authProvider`) from register
    * time, not a fresh `context.getAuthProvider?.()` call. If the underlying

lib/telemetry/TelemetryClientProvider.ts

@@ -133,35 +133,26 @@ class TelemetryClientProvider {
       return;
     }
 
+    // Skip unregister on the last release so close()'s final flush can still
+    // resolve auth/connection providers from the FIFO snapshot.
     if (holder.refCount > 1) {
-      // Other registrants remain — drop this context now so subsequent
-      // flushes by surviving consumers don't try to authenticate via a
-      // context whose underlying DBSQLClient is closing.
       holder.client.unregisterContext(context);
-      holder.refCount -= 1;
-      logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`);
-      return;
     }
-
-    // Last refcount holder. Keep the context registered through close()
-    // so the final flush can still resolve `getAuthProvider()` and
-    // `getConnectionProvider()` from the FIFO snapshot — otherwise the
-    // exporter drops the final batch with "missing Authorization header".
-    // The TelemetryClient is fully closed below, so the lingering FIFO
-    // entry has no further effect.
     holder.refCount -= 1;
     logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`);
 
-    // Remove from map BEFORE awaiting close so a concurrent
-    // getOrCreateClient creates a fresh instance rather than receiving
-    // this closing one.
-    this.clients.delete(key);
-    try {
-      await holder.client.close();
-      logger.log(LogLevel.debug, `Closed and removed TelemetryClient for host: ${host}`);
-    } catch (error) {
-      const msg = error instanceof Error ? error.message : String(error);
-      logger.log(LogLevel.debug, `Error releasing TelemetryClient: ${msg}`);
+    if (holder.refCount <= 0) {
+      // Remove from map BEFORE awaiting close so a concurrent
+      // getOrCreateClient creates a fresh instance rather than receiving
+      // this closing one.
+      this.clients.delete(key);
+      try {
+        await holder.client.close();
+        logger.log(LogLevel.debug, `Closed and removed TelemetryClient for host: ${host}`);
+      } catch (error) {
+        const msg = error instanceof Error ? error.message : String(error);
+        logger.log(LogLevel.debug, `Error releasing TelemetryClient: ${msg}`);
+      }
     }
   }