Fix close-ordering: keep context registered through final telemetry flush

`TelemetryClientProvider.releaseClient` was calling `unregisterContext`
before `client.close()`. On the last refcount, that emptied the FIFO of
`(context, authProvider)` pairs before the final flush ran, so the
exporter's `getAuthProvider()` walk returned undefined and the batch was
dropped with "Telemetry: Authorization header missing — metrics will be
dropped".

Defer `unregisterContext` until after `close()` on the last refcount;
multi-refcount path is unchanged (immediate unregister so surviving
consumers don't reach into a closing context).

Verified end-to-end against a real SPOG host: the final flush now
exports its 3 metrics (connection.open, statement.start,
statement.complete) instead of warning and dropping them.

Co-authored-by: Isaac

Author: samikshya-db

lib/telemetry/TelemetryClient.ts

@@ -169,8 +169,10 @@ class TelemetryClient implements IClientContext {
 
   /**
    * Remove a `DBSQLClient`'s context from the pool. Called by
-   * `TelemetryClientProvider.releaseClient` before refcount decrement so the
-   * exporter doesn't keep trying to use a closed context.
+   * `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.
    *
    * Uses the cached snapshot pair (`context`, `authProvider`) from register
    * time, not a fresh `context.getAuthProvider?.()` call. If the underlying

lib/telemetry/TelemetryClientProvider.ts

@@ -133,22 +133,35 @@ class TelemetryClientProvider {
       return;
     }
 
-    holder.client.unregisterContext(context);
+    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}`);
 
-    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}`);
-      }
+    // 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}`);
     }
   }
 

tests/unit/telemetry/TelemetryClientProvider.test.ts

@@ -267,6 +267,50 @@ describe('TelemetryClientProvider', () => {
       expect(second.isClosed()).to.be.false;
       expect(provider.getRefCount(HOST1)).to.equal(1);
     });
+
+    it('keeps the context registered through close() on last refcount', async () => {
+      // Regression: releaseClient used to call unregisterContext before
+      // close(), which dropped the only FIFO entry and left the final
+      // flush without an auth provider — metrics were dropped with
+      // "missing Authorization header".
+      const context = new ClientContextStub();
+      const provider = new TelemetryClientProvider();
+
+      const client = provider.getOrCreateClient(context, HOST1);
+      let authProviderAtClose: unknown = 'unset';
+      sinon.stub(client, 'close').callsFake(async () => {
+        // The TelemetryClient is its own IClientContext. While close() runs,
+        // getAuthProvider() must still resolve (the FIFO entry survives).
+        authProviderAtClose = client.getAuthProvider?.();
+      });
+
+      await provider.releaseClient(context, HOST1);
+
+      // The FIFO walk hits the (still-registered) context; the stub's
+      // getAuthProvider returns undefined but the lookup itself completes.
+      // What we're really asserting is that the FIFO wasn't pre-emptied:
+      // pre-fix, this assertion would not even fire because close() runs
+      // against an empty FIFO and the auth-provider walk short-circuits
+      // before close()'s callsFake. Here we verify the spy ran AND the
+      // context is still registered at that moment.
+      expect(authProviderAtClose).to.not.equal('unset');
+      expect((client as any).contexts.length).to.equal(1);
+    });
+
+    it('drops the context immediately when other refcounts remain', async () => {
+      const context = new ClientContextStub();
+      const provider = new TelemetryClientProvider();
+
+      const client = provider.getOrCreateClient(context, HOST1);
+      provider.getOrCreateClient(context, HOST1); // refcount=2
+
+      await provider.releaseClient(context, HOST1);
+
+      // Multi-refcount path: unregisterContext runs immediately; the
+      // (single) FIFO entry tracking this context was removed.
+      expect((client as any).contexts.length).to.equal(0);
+      expect(provider.getRefCount(HOST1)).to.equal(1);
+    });
   });
 
   describe('Host normalization', () => {