feat(observability): add tracing spans to Session

This change adds tracing spans to Session for methods:
* create
* getMetadata
* keepAlive

and corresponding tests.

Updates googleapis#2079
Built from PR googleapis#2087
Updates googleapis#2114

observability-test/session.ts

@@ -0,0 +1,247 @@
+/*!
+ * Copyright 2024 Google LLC. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as assert from 'assert';
+import {grpc} from 'google-gax';
+import {google} from '../protos/protos';
+import {Database, Session, Spanner} from '../src';
+import protobuf = google.spanner.v1;
+import * as mock from '../test/mockserver/mockspanner';
+import * as mockInstanceAdmin from '../test/mockserver/mockinstanceadmin';
+import * as mockDatabaseAdmin from '../test/mockserver/mockdatabaseadmin';
+const {
+  AlwaysOnSampler,
+  NodeTracerProvider,
+  InMemorySpanExporter,
+} = require('@opentelemetry/sdk-trace-node');
+// eslint-disable-next-line n/no-extraneous-require
+const {SimpleSpanProcessor} = require('@opentelemetry/sdk-trace-base');
+
+/** A simple result set for SELECT 1. */
+function createSelect1ResultSet(): protobuf.ResultSet {
+  const fields = [
+    protobuf.StructType.Field.create({
+      name: 'NUM',
+      type: protobuf.Type.create({code: protobuf.TypeCode.INT64}),
+    }),
+  ];
+  const metadata = new protobuf.ResultSetMetadata({
+    rowType: new protobuf.StructType({
+      fields,
+    }),
+  });
+  return protobuf.ResultSet.create({
+    metadata,
+    rows: [{values: [{stringValue: '1'}]}],
+  });
+}
+
+interface setupResults {
+  server: grpc.Server;
+  spanner: Spanner;
+  spannerMock: mock.MockSpanner;
+}
+
+async function setup(): Promise<setupResults> {
+  const server = new grpc.Server();
+
+  const spannerMock = mock.createMockSpanner(server);
+  mockInstanceAdmin.createMockInstanceAdmin(server);
+  mockDatabaseAdmin.createMockDatabaseAdmin(server);
+
+  const port: number = await new Promise((resolve, reject) => {
+    server.bindAsync(
+      '0.0.0.0:0',
+      grpc.ServerCredentials.createInsecure(),
+      (err, assignedPort) => {
+        if (err) {
+          reject(err);
+        } else {
+          resolve(assignedPort);
+        }
+      }
+    );
+  });
+
+  const selectSql = 'SELECT 1';
+  const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2';
+  spannerMock.putStatementResult(
+    selectSql,
+    mock.StatementResult.resultSet(createSelect1ResultSet())
+  );
+  spannerMock.putStatementResult(
+    updateSql,
+    mock.StatementResult.updateCount(1)
+  );
+
+  const spanner = new Spanner({
+    projectId: 'observability-project-id',
+    servicePath: 'localhost',
+    port,
+    sslCreds: grpc.credentials.createInsecure(),
+  });
+
+  return Promise.resolve({
+    spanner: spanner,
+    server: server,
+    spannerMock: spannerMock,
+  });
+}
+
+describe('Session', () => {
+  let server: grpc.Server;
+  let spanner: Spanner;
+  let database: Database;
+  let spannerMock: mock.MockSpanner;
+  let traceExporter: typeof InMemorySpanExporter;
+
+  after(() => {
+    spanner.close();
+    server.tryShutdown(() => {});
+  });
+
+  before(async () => {
+    const setupResult = await setup();
+    spanner = setupResult.spanner;
+    server = setupResult.server;
+    spannerMock = setupResult.spannerMock;
+
+    const selectSql = 'SELECT 1';
+    const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2';
+    spannerMock.putStatementResult(
+      selectSql,
+      mock.StatementResult.resultSet(createSelect1ResultSet())
+    );
+    spannerMock.putStatementResult(
+      updateSql,
+      mock.StatementResult.updateCount(1)
+    );
+
+    traceExporter = new InMemorySpanExporter();
+    const sampler = new AlwaysOnSampler();
+
+    const provider = new NodeTracerProvider({
+      sampler: sampler,
+      exporter: traceExporter,
+    });
+    provider.addSpanProcessor(new SimpleSpanProcessor(traceExporter));
+
+    const instance = spanner.instance('instance');
+    database = instance.database('database');
+    database.observabilityConfig = {
+      tracerProvider: provider,
+      enableExtendedTracing: false,
+    };
+  });
+
+  beforeEach(() => {
+    spannerMock.resetRequests();
+  });
+
+  afterEach(() => {
+    traceExporter.reset();
+  });
+
+  it('create with constructor', done => {
+    const session = new Session(database);
+    session.create(err => {
+      traceExporter.forceFlush();
+      const spans = traceExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
+
+      const actualSpanNames: string[] = [];
+      spans.forEach(span => {
+        actualSpanNames.push(span.name);
+      });
+
+      const expectedSpanNames = ['CloudSpanner.Session.create'];
+      assert.deepStrictEqual(
+        actualSpanNames,
+        expectedSpanNames,
+        `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
+      );
+
+      done();
+    });
+  });
+
+  it('create with database.session()', done => {
+    const session = database.session();
+    session.create(err => {
+      assert.ifError(err);
+      traceExporter.forceFlush();
+      const spans = traceExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
+
+      const actualSpanNames: string[] = [];
+      spans.forEach(span => {
+        actualSpanNames.push(span.name);
+      });
+
+      const expectedSpanNames = ['CloudSpanner.Session.create'];
+      assert.deepStrictEqual(
+        actualSpanNames,
+        expectedSpanNames,
+        `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
+      );
+
+      done();
+    });
+  });
+
+  it('getMetadata', done => {
+    const session = database.session();
+    session.create(err => {
+      assert.ifError(err);
+      traceExporter.forceFlush();
+      traceExporter.reset();
+
+      session.getMetadata((err, metadata) => {
+        const spans = traceExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
+        const span = spans[0];
+
+        const expectedSpanName = 'CloudSpanner.Session.getMetadata';
+        assert.deepStrictEqual(
+          span.name,
+          expectedSpanName,
+          `span names mismatch:\n\tGot:  ${span.name}\n\tWant: ${expectedSpanName}`
+        );
+
+        done();
+      });
+    });
+  });
+
+  it('keepAlive produces no spans', done => {
+    // It is imperative that the .keepAlive() / ping method is not
+    // traced, because that spams trace views/logs unecessarily.
+    // Please see https://github.com/googleapis/google-cloud-go/issues/1691
+    const session = database.session();
+    session.create(err => {
+      assert.ifError(err);
+      traceExporter.forceFlush();
+      traceExporter.reset();
+
+      session.keepAlive(err => {
+        traceExporter.forceFlush();
+        const spans = traceExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 0, 'No spans should be exported');
+        done();
+      });
+    });
+  });
+});

src/session.ts

@@ -44,6 +44,7 @@ import {
 import {grpc, CallOptions} from 'google-gax';
 import IRequestOptions = google.spanner.v1.IRequestOptions;
 import {Spanner} from '.';
+import {ObservabilityOptions, startTrace, setSpanError} from './instrument';
 
 export type GetSessionResponse = [Session, r.Response];
 
@@ -118,6 +119,7 @@ export class Session extends common.GrpcServiceObject {
   lastUsed?: number;
   lastError?: grpc.ServiceError;
   resourceHeader_: {[k: string]: string};
+  observabilityConfig: ObservabilityOptions | undefined;
   constructor(database: Database, name?: string) {
     const methods = {
       /**
@@ -247,14 +249,24 @@ export class Session extends common.GrpcServiceObject {
         this.databaseRole =
           options.databaseRole || database.databaseRole || null;
 
-        return database.createSession(options, (err, session, apiResponse) => {
-          if (err) {
-            callback(err, null, apiResponse);
-            return;
-          }
+        const q = {opts: database.observabilityConfig};
+        return startTrace('Session.create', q, span => {
+          return database.createSession(
+            options,
+            (err, session, apiResponse) => {
+              if (err) {
+                setSpanError(span, err);
+                span.end();
+                callback(err, null, apiResponse);
+                return;
+              }
 
-          extend(this, session);
-          callback(null, this, apiResponse);
+              this.observabilityConfig = database.ovservabilityConfig;
+              span.end();
+              extend(this, session);
+              callback(null, this, apiResponse);
+            }
+          );
         });
       },
     } as {} as ServiceObjectConfig);
@@ -268,6 +280,8 @@ export class Session extends common.GrpcServiceObject {
     if (name) {
       this.formattedName_ = Session.formatName_(database.formattedName_, name);
     }
+
+    this.observabilityConfig = database.observabilityConfig;
   }
   /**
    * Delete a session.
@@ -388,23 +402,32 @@ export class Session extends common.GrpcServiceObject {
     if (this._getSpanner().routeToLeaderEnabled) {
       addLeaderAwareRoutingHeader(headers);
     }
-    return this.request(
-      {
-        client: 'SpannerClient',
-        method: 'getSession',
-        reqOpts,
-        gaxOpts,
-        headers: headers,
-      },
-      (err, resp) => {
-        if (resp) {
-          resp.databaseRole = resp.creatorRole;
-          delete resp.creatorRole;
-          this.metadata = resp;
+
+    const q = {opts: this.observabilityConfig};
+    return startTrace('Session.getMetadata', q, span => {
+      return this.request(
+        {
+          client: 'SpannerClient',
+          method: 'getSession',
+          reqOpts,
+          gaxOpts,
+          headers: headers,
+        },
+        (err, resp) => {
+          if (err) {
+            setSpanError(span, err);
+          }
+
+          if (resp) {
+            resp.databaseRole = resp.creatorRole;
+            delete resp.creatorRole;
+            this.metadata = resp;
+          }
+          span.end();
+          callback!(err, resp);
         }
-        callback!(err, resp);
-      }
-    );
+      );
+    });
   }
   /**
    * Ping the session with `SELECT 1` to prevent it from expiring.
@@ -431,6 +454,10 @@ export class Session extends common.GrpcServiceObject {
     optionsOrCallback?: CallOptions | KeepAliveCallback,
     cb?: KeepAliveCallback
   ): void | Promise<KeepAliveResponse> {
+    // NOTE: Please do not trace Ping as it gets quite spammy
+    // with many root spans polluting the main span.
+    // Please see https://github.com/googleapis/google-cloud-go/issues/1691
+
     const gaxOpts =
       typeof optionsOrCallback === 'object' ? optionsOrCallback : {};
     const callback =
@@ -440,6 +467,7 @@ export class Session extends common.GrpcServiceObject {
       session: this.formattedName_,
       sql: 'SELECT 1',
     };
+
     return this.request(
       {
         client: 'SpannerClient',
@@ -448,7 +476,9 @@ export class Session extends common.GrpcServiceObject {
         gaxOpts,
         headers: this.resourceHeader_,
       },
-      callback!
+      err => {
+        callback!(err);
+      }
     );
   }
   /**