review comments

Author: alkatrivedi

handwritten/spanner/google-cloud-spanner-executor/src/cloud-client-executor.ts

@@ -49,9 +49,12 @@ export class ExecutionFlowContext implements ExecutionFlowContextInterface {
    * Sends a response back to the client.
    */
   public onNext(response: SpannerAsyncActionResponse): void {
-
     // Prevent writing if client cancelled the call, or the underlying Node stream is un-writable/destroyed
-    if (this.call.cancelled || this.call.destroyed || this.call.writable === false) {
+    if (
+      this.call.cancelled ||
+      this.call.destroyed ||
+      this.call.writable === false
+    ) {
       console.warn('Attempted to write to a closed or cancelled stream.');
       return;
     }
@@ -84,10 +87,25 @@ export class ExecutionFlowContext implements ExecutionFlowContextInterface {
   }
 }
 
+type ActionHandler = (action: any, sender: OutcomeSender) => Promise<void>;
+
 export class CloudClientExecutor {
   private spanner: Spanner;
   private tracer: Tracer;
 
+  private readonly adminActionRegistry: Record<string, ActionHandler> = {
+    createCloudInstance: (action, sender) =>
+      this.executeCreateCloudInstance(
+        action as ICreateCloudInstanceAction,
+        sender,
+      ),
+  };
+
+  private readonly actionRegistry: Record<string, ActionHandler> = {
+    admin: (action, sender) =>
+      this.executeAdminAction(action as IAdminAction, sender),
+  };
+
   constructor() {
     const spannerOptions = CloudUtil.getSpannerOptions();
     this.spanner = new Spanner(spannerOptions);
@@ -144,8 +162,12 @@ export class CloudClientExecutor {
 
     return context.with(trace.setSpan(context.active(), span), async () => {
       try {
-        if (action.admin) {
-          await this.executeAdminAction(action.admin, outcomeSender);
+        const handler = this.actionRegistry[actionType];
+        if (handler) {
+          await handler(
+            action[actionType as keyof typeof action],
+            outcomeSender,
+          );
           return;
         }
 
@@ -171,16 +193,21 @@ export class CloudClientExecutor {
     sender: OutcomeSender,
   ): Promise<void> {
     try {
-      if (action.createCloudInstance) {
-        await this.executeCreateCloudInstance(
-          action.createCloudInstance,
+      const adminType = Object.keys(action).find(
+        k => action[k as keyof typeof action] !== undefined,
+      );
+
+      if (adminType && this.adminActionRegistry[adminType]) {
+        await this.adminActionRegistry[adminType](
+          action[adminType as keyof typeof action],
           sender,
         );
         return;
       }
+
       sender.finishWithError({
         code: status.UNIMPLEMENTED,
-        message: 'Admin action not implemented',
+        message: `Admin action ${adminType || 'unknown'} not implemented`,
       });
     } catch (e: any) {
       sender.finishWithError(e);
@@ -213,7 +240,6 @@ export class CloudClientExecutor {
       });
 
       console.log('Waiting for instance creation operation to complete...');
-
       await operation.promise();
 
       console.log(`Instance ${instanceId} created successfully.`);

handwritten/spanner/google-cloud-spanner-executor/src/cloud-executor-impl.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import {ServerDuplexStream, status, ServiceError} from '@grpc/grpc-js';
+import {ServerDuplexStream, status} from '@grpc/grpc-js';
 import {trace, context, Tracer} from '@opentelemetry/api';
 import {CloudClientExecutor} from './cloud-client-executor';
 import * as protos from '../../protos/protos';
@@ -64,22 +64,7 @@ export class CloudExecutorImpl {
     call.on('data', (request: SpannerAsyncActionRequest) => {
       context.with(streamContext, () => {
         console.log(`Receiving request: \n${JSON.stringify(request, null, 2)}`);
-
-        // Ensure nested properties exist before attempting to inject configuration overrides.
-        if (!request.action) request.action = {};
-        if (!request.action.spannerOptions) request.action.spannerOptions = {};
-        if (!request.action.spannerOptions.sessionPoolOptions)
-          request.action.spannerOptions.sessionPoolOptions = {};
-
-        // as the multiplexed session is defualt enabled, mutate the incoming request to forcefully enforce multiplexed sessions for this proxy execution.
-        request.action.spannerOptions.sessionPoolOptions.useMultiplexed = true;
-
-        console.log(
-          `Updated request to set multiplexed session flag: \n${JSON.stringify(request, null, 2)}`,
-        );
-
         // TODO: Set requestHasReadOrQueryAction flag here when Read/Query are implemented.
-
         try {
           const reqStatus = this.clientExecutor.startHandlingRequest(
             request,
@@ -107,8 +92,8 @@ export class CloudExecutorImpl {
     });
 
     // Handle the completion of the client stream
-    call.on('end', async () => {
-      await context.with(streamContext, async () => {
+    call.on('end', () => {
+      context.with(streamContext, () => {
         span.end();
         // TODO: Add End-to-End trace verification here once Read/Query actions are implemented.
         console.log('Client called Done, half closed');

handwritten/spanner/google-cloud-spanner-executor/src/cloud-util.ts

@@ -17,12 +17,49 @@
 import * as fs from 'fs';
 import * as grpc from '@grpc/grpc-js';
 import {WorkerProxy} from './worker-proxy';
-import { SpannerOptions } from '../../src';
+import {SpannerOptions} from '../../src';
+
+import {NodeTracerProvider} from '@opentelemetry/sdk-trace-node';
+import {OTLPTraceExporter} from '@opentelemetry/exporter-trace-otlp-grpc';
+import {Resource} from '@opentelemetry/resources';
+import {ATTR_SERVICE_NAME} from '@opentelemetry/semantic-conventions';
+import {
+  BatchSpanProcessor,
+  TraceIdRatioBasedSampler,
+} from '@opentelemetry/sdk-trace-base';
+import {GoogleAuth} from 'google-auth-library';
 
 /**
  * Provides utility methods for configuring the Cloud Spanner client for tests.
  */
 export class CloudUtil {
+  public static async setupOpenTelemetrySdk(): Promise<NodeTracerProvider> {
+    const auth = new GoogleAuth({
+      scopes: 'https://www.googleapis.com/auth/cloud-platform',
+      keyFile: WorkerProxy.serviceKeyFile || undefined,
+    });
+    const authenticatedClient = await auth.getClient();
+
+    const traceExporter = new OTLPTraceExporter({
+      url: 'https://telemetry.googleapis.com',
+      credentials: grpc.credentials.combineChannelCredentials(
+        grpc.credentials.createSsl(),
+        grpc.credentials.createFromGoogleCredential(authenticatedClient as any),
+      ),
+    });
+
+    const provider = new NodeTracerProvider({
+      resource: new Resource({
+        [ATTR_SERVICE_NAME]: 'spanner-node-worker-proxy',
+        'gcp.project_id': WorkerProxy.PROJECT_ID,
+      }) as any,
+      sampler: new TraceIdRatioBasedSampler(1.0),
+      spanProcessors: [new BatchSpanProcessor(traceExporter as any)],
+    });
+
+    provider.register();
+    return provider;
+  }
   // If this is set too low, the peer server may return RESOURCE_EXHAUSTED errors if the response
   // error message causes the trailing headers to exceed this limit.
   private static readonly GRPC_MAX_HEADER_LIST_SIZE_BYTES = 10 * 1024 * 1024; // 10 MB
@@ -59,7 +96,8 @@ export class CloudUtil {
       // - default_authority is used for the HTTP/2 :authority header.
       (grpcOptions as grpc.ChannelOptions)['grpc.ssl_target_name_override'] =
         this.TEST_HOST_IN_CERT;
-      (grpcOptions as grpc.ChannelOptions)['grpc.default_authority'] = this.TEST_HOST_IN_CERT;
+      (grpcOptions as grpc.ChannelOptions)['grpc.default_authority'] =
+        this.TEST_HOST_IN_CERT;
     }
 
     (options as any).grpcOptions = grpcOptions;

handwritten/spanner/google-cloud-spanner-executor/src/worker-proxy.ts

@@ -19,15 +19,8 @@ import * as protoLoader from '@grpc/proto-loader';
 import yargs from 'yargs';
 import * as path from 'path';
 import * as fs from 'fs';
-import {NodeTracerProvider} from '@opentelemetry/sdk-trace-node';
-import {OTLPTraceExporter} from '@opentelemetry/exporter-trace-otlp-grpc';
-import {Resource} from '@opentelemetry/resources';
-import {ATTR_SERVICE_NAME} from '@opentelemetry/semantic-conventions';
-import {
-  BatchSpanProcessor,
-  TraceIdRatioBasedSampler,
-} from '@opentelemetry/sdk-trace-base';
 import {CloudExecutorImpl} from './cloud-executor-impl';
+import {CloudUtil} from './cloud-util';
 import {HealthImplementation} from 'grpc-health-check';
 import {ReflectionService} from '@grpc/reflection';
 
@@ -42,8 +35,6 @@ const OPTION_CERTIFICATE = 'cert';
 const OPTION_SERVICE_KEY_FILE = 'service_key_file';
 const OPTION_USE_PLAIN_TEXT_CHANNEL = 'use_plain_text_channel';
 const OPTION_ENABLE_GRPC_FAULT_INJECTOR = 'enable_grpc_fault_injector';
-const OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO =
-  'multiplexed_session_operations_ratio';
 
 /**
  * Acts as a proxy server that forwards incoming gRPC requests to the underlying
@@ -54,10 +45,10 @@ export class WorkerProxy {
   public static proxyPort = 0;
   public static cert = '';
   public static serviceKeyFile = '';
-  public static multiplexedSessionOperationsRatio = 0.0;
+
   public static usePlainTextChannel = false;
   public static enableGrpcFaultInjector = false;
-  public static openTelemetrySdk: NodeTracerProvider;
+  public static openTelemetrySdk: any;
 
   public static readonly PROJECT_ID = 'spanner-cloud-systest';
   public static readonly CLOUD_TRACE_ENDPOINT =
@@ -67,34 +58,6 @@ export class WorkerProxy {
   private static readonly MAX_PORT = 65535;
   private static readonly TRACE_SAMPLING_RATE = 0.01;
 
-  /**
-   * Sets up the OpenTelemetry SDK
-   */
-  public static async setupOpenTelemetrySdk(): Promise<NodeTracerProvider> {
-    const exporterConfig = {
-      url: this.CLOUD_TRACE_ENDPOINT,
-      headers: {
-        'x-goog-api-key': this.PROJECT_ID,
-      },
-    };
-
-    // TODO: Replace with TraceExporter once the Cloud Trace Exporter is compatible with OpenTelemetry v2.
-    // This will allow us to utilize `exporterConfig` and emit traces to GCP.
-
-    const traceExporter = new OTLPTraceExporter(exporterConfig);
-
-    const provider = new NodeTracerProvider({
-      resource: new Resource({
-        [ATTR_SERVICE_NAME]: 'spanner-node-worker-proxy',
-      }) as any,
-      sampler: new TraceIdRatioBasedSampler(this.TRACE_SAMPLING_RATE),
-      spanProcessors: [new BatchSpanProcessor(traceExporter as any)],
-    });
-
-    provider.register();
-    return provider;
-  }
-
   /**
    * Parses and builds the command line options for the worker proxy.
    */
@@ -126,10 +89,6 @@ export class WorkerProxy {
       type: 'boolean',
       description: 'Enable grpc fault injector in cloud client executor.',
     });
-    parser.option(OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO, {
-      type: 'number',
-      description: 'Ratio of operations to use multiplexed sessions.',
-    });
 
     try {
       return parser.parseSync();
@@ -187,7 +146,7 @@ export class WorkerProxy {
       !!commandLine[OPTION_ENABLE_GRPC_FAULT_INJECTOR];
 
     // Setup the OpenTelemetry for tracing
-    this.openTelemetrySdk = await this.setupOpenTelemetrySdk();
+    this.openTelemetrySdk = await CloudUtil.setupOpenTelemetrySdk();
 
     // Check if proto file exists
     if (!fs.existsSync(PROTO_PATH)) {
@@ -270,6 +229,16 @@ export class WorkerProxy {
 
     process.on('SIGTERM', shutdown);
     process.on('SIGINT', shutdown);
+
+    process.on('uncaughtException', err => {
+      console.error('Uncaught Exception:', err);
+      shutdown();
+    });
+
+    process.on('unhandledRejection', reason => {
+      console.error('Unhandled Rejection:', reason);
+      shutdown();
+    });
   }
 }
 if (require.main === module) {