feat(spanner): worker proxy for executor framework

Author: alkatrivedi

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

@@ -0,0 +1,318 @@
+/*!
+ * Copyright 2026 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 grpc from '@grpc/grpc-js';
+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 {TraceExporter} from '@google-cloud/opentelemetry-cloud-trace-exporter';
+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 {HealthImplementation} from 'grpc-health-check';
+import {ReflectionService} from '@grpc/reflection';
+
+const PROTO_PATH = path.join(
+  __dirname,
+  '../../protos/google/spanner/executor/v1/cloud_executor.proto',
+);
+
+const OPTION_SPANNER_PORT = 'spanner_port';
+const OPTION_PROXY_PORT = 'proxy_port';
+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';
+
+/**
+ * class WorkerProxy which acts as a proxy server that forwards requests to the Spanner executor
+ */
+export class WorkerProxy {
+  public static spannerPort = 0;
+  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 readonly PROJECT_ID = 'spanner-cloud-systest';
+  public static readonly CLOUD_TRACE_ENDPOINT =
+    'staging-cloudtrace.sandbox.googleapis.com:443';
+
+  private static readonly MIN_PORT = 0;
+  private static readonly MAX_PORT = 65535;
+  private static readonly MIN_RATIO = 0.0;
+  private static readonly MAX_RATIO = 1.0;
+  private static readonly TRACE_SAMPLING_RATE = 0.01;
+
+  /**
+   * Sets up the OpenTelemetry SDK
+   */
+  public static async setupOpenTelemetrySdk(): Promise<NodeTracerProvider> {
+    const exporterConfig: any = {
+      projectId: WorkerProxy.PROJECT_ID,
+      apiEndpoint: WorkerProxy.CLOUD_TRACE_ENDPOINT,
+    };
+
+    if (WorkerProxy.serviceKeyFile) {
+      exporterConfig.keyFilename = WorkerProxy.serviceKeyFile;
+    }
+
+    const traceExporter = new TraceExporter(exporterConfig);
+
+    const provider = new NodeTracerProvider({
+      resource: new Resource({
+        [ATTR_SERVICE_NAME]: 'spanner-node-worker-proxy',
+      }) as any,
+      sampler: new TraceIdRatioBasedSampler(WorkerProxy.TRACE_SAMPLING_RATE),
+      spanProcessors: [new BatchSpanProcessor(traceExporter as any)],
+    });
+
+    provider.register();
+    return provider;
+  }
+
+  /**
+   * Builds the command line options for the worker proxy
+   */
+  public static buildOptions(args: string[]): any {
+    const parser = yargs(args);
+
+    parser.option(OPTION_SPANNER_PORT, {
+      type: 'number',
+      description: 'Port of Spanner Frontend to which to send requests.',
+    });
+    parser.option(OPTION_PROXY_PORT, {
+      type: 'number',
+      description: 'Proxy port to start worker proxy on.',
+    });
+    parser.option(OPTION_CERTIFICATE, {
+      type: 'string',
+      description: 'Certificate used to connect to Spanner GFE.',
+    });
+    parser.option(OPTION_SERVICE_KEY_FILE, {
+      type: 'string',
+      description: 'Service key file used to set authentication.',
+    });
+    parser.option(OPTION_USE_PLAIN_TEXT_CHANNEL, {
+      type: 'boolean',
+      description:
+        'Use a plain text gRPC channel (intended for the Cloud Spanner Emulator).',
+    });
+    parser.option(OPTION_ENABLE_GRPC_FAULT_INJECTOR, {
+      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();
+    } catch (e: any) {
+      throw new Error(e.message);
+    }
+  }
+
+  /**
+   * main method to spin up the server and start the worker proxy
+   */
+  public static async main(args: string[]) {
+    const commandLineString = WorkerProxy.buildOptions(args);
+    const commandLine = commandLineString as any;
+
+    if (commandLine[OPTION_SPANNER_PORT] === undefined) {
+      throw new Error(
+        'Spanner proxyPort need to be assigned in order to start worker proxy.',
+      );
+    }
+    WorkerProxy.spannerPort = commandLine[OPTION_SPANNER_PORT];
+    if (
+      WorkerProxy.spannerPort < WorkerProxy.MIN_PORT ||
+      WorkerProxy.spannerPort > WorkerProxy.MAX_PORT
+    ) {
+      throw new Error(
+        'Spanner proxyPort must be between ' +
+          WorkerProxy.MIN_PORT +
+          ' and ' +
+          WorkerProxy.MAX_PORT,
+      );
+    }
+
+    if (commandLine[OPTION_PROXY_PORT] === undefined) {
+      throw new Error(
+        'Proxy port need to be assigned in order to start worker proxy.',
+      );
+    }
+    WorkerProxy.proxyPort = commandLine[OPTION_PROXY_PORT];
+    if (
+      WorkerProxy.proxyPort < WorkerProxy.MIN_PORT ||
+      WorkerProxy.proxyPort > WorkerProxy.MAX_PORT
+    ) {
+      throw new Error(
+        'Proxy port must be between ' +
+          WorkerProxy.MIN_PORT +
+          ' and ' +
+          WorkerProxy.MAX_PORT,
+      );
+    }
+
+    if (!commandLine[OPTION_CERTIFICATE]) {
+      throw new Error(
+        'Certificate need to be assigned in order to start worker proxy.',
+      );
+    }
+    WorkerProxy.cert = commandLine[OPTION_CERTIFICATE];
+
+    if (commandLine[OPTION_SERVICE_KEY_FILE]) {
+      WorkerProxy.serviceKeyFile = commandLine[OPTION_SERVICE_KEY_FILE];
+    }
+
+    WorkerProxy.usePlainTextChannel =
+      !!commandLine[OPTION_USE_PLAIN_TEXT_CHANNEL];
+    WorkerProxy.enableGrpcFaultInjector =
+      !!commandLine[OPTION_ENABLE_GRPC_FAULT_INJECTOR];
+
+    if (
+      commandLine[OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO] !== undefined
+    ) {
+      WorkerProxy.multiplexedSessionOperationsRatio = Number(
+        commandLine[OPTION_MULTIPLEXED_SESSION_OPERATIONS_RATIO],
+      );
+      console.info(
+        `Multiplexed session ratio from commandline arg: \n${WorkerProxy.multiplexedSessionOperationsRatio}`,
+      );
+      if (
+        WorkerProxy.multiplexedSessionOperationsRatio < WorkerProxy.MIN_RATIO ||
+        WorkerProxy.multiplexedSessionOperationsRatio > WorkerProxy.MAX_RATIO
+      ) {
+        throw new Error(
+          'Spanner multiplexedSessionOperationsRatio must be between ' +
+            WorkerProxy.MIN_RATIO +
+            ' and ' +
+            WorkerProxy.MAX_RATIO,
+        );
+      }
+    }
+
+    // Setup the OpenTelemetry for tracing
+    WorkerProxy.openTelemetrySdk = await WorkerProxy.setupOpenTelemetrySdk();
+
+    // Check if proto file exists
+    if (!fs.existsSync(PROTO_PATH)) {
+      throw new Error(`Proto file not found at ${PROTO_PATH}`);
+    }
+
+    const packageDefinition = protoLoader.loadSync(PROTO_PATH, {
+      keepCase: true,
+      longs: String,
+      enums: String,
+      defaults: true,
+      oneofs: true,
+      includeDirs: [
+        path.join(__dirname, '../../protos'),
+        path.join(__dirname, '../../node_modules/google-proto-files'),
+        path.join(__dirname, '../../node_modules/google-gax/build/protos'),
+      ],
+    });
+    const protoDescriptor = grpc.loadPackageDefinition(
+      packageDefinition,
+    ) as any;
+    const spannerExecutorProxy =
+      protoDescriptor.google.spanner.executor.v1.SpannerExecutorProxy;
+
+    let server: grpc.Server;
+    for (;;) {
+      try {
+        const cloudExecutorImpl = new CloudExecutorImpl(
+          WorkerProxy.enableGrpcFaultInjector,
+          WorkerProxy.multiplexedSessionOperationsRatio,
+        );
+
+        server = new grpc.Server();
+        server.addService(
+          spannerExecutorProxy.service,
+          cloudExecutorImpl as any,
+        );
+
+        const healthImpl = new HealthImplementation({
+          '': 'SERVING',
+        });
+        healthImpl.addToServer(server);
+        const reflection = new ReflectionService(packageDefinition);
+        reflection.addToServer(server);
+
+        const bindAddr = `0.0.0.0:${WorkerProxy.proxyPort}`;
+        const port = await new Promise<number>((resolve, reject) => {
+          server.bindAsync(
+            bindAddr,
+            grpc.ServerCredentials.createInsecure(),
+            (err, port) => {
+              if (err) {
+                return reject(err);
+              }
+              resolve(port);
+            },
+          );
+        });
+
+        console.info(`Server started on proxyPort: ${port}`);
+        break;
+      } catch (e) {
+        console.warn(
+          `Failed to start server on proxyPort ${WorkerProxy.proxyPort}`,
+          e,
+        );
+        // Wait briefly before retrying to avoid tight loop
+        await new Promise(resolve => setTimeout(resolve, 1000));
+      }
+    }
+
+    const shutdown = () => {
+      // eslint-disable-next-line n/no-process-exit
+      setTimeout(() => process.exit(1), 2000).unref();
+      server.tryShutdown(() => {
+        WorkerProxy.openTelemetrySdk
+          .shutdown()
+          .then(() => console.info('Tracing terminated'))
+          .catch(error => console.error('Error terminating tracing', error))
+          // eslint-disable-next-line n/no-process-exit
+          .finally(() => process.exit(0));
+      });
+    };
+
+    process.on('SIGTERM', shutdown);
+    process.on('SIGINT', shutdown);
+  }
+}
+
+if (require.main === module) {
+  WorkerProxy.main(process.argv.slice(2)).catch(err => {
+    console.error('Failed to start worker proxy:', err);
+    throw err;
+  });
+}

handwritten/spanner/package.json

@@ -98,6 +98,8 @@
     "uuid": "^11.1.0"
   },
   "devDependencies": {
+    "@google-cloud/opentelemetry-cloud-trace-exporter": "^3.0.0",
+    "@grpc/reflection": "^1.0.4",
     "@opentelemetry/sdk-trace-base": "^2.0.0",
     "@opentelemetry/sdk-trace-node": "^2.0.0",
     "@types/concat-stream": "^2.0.3",
@@ -112,13 +114,15 @@
     "@types/request": "^2.48.12",
     "@types/sinon": "^21.0.0",
     "@types/through2": "^2.0.41",
+    "@types/yargs": "^17.0.35",
     "binary-search-bounds": "^2.0.5",
     "c8": "^10.1.3",
     "codecov": "^3.8.3",
     "concat-stream": "^2.0.0",
     "dedent": "^1.5.3",
     "execa": "^5.0.0",
     "gapic-tools": "^1.0.1",
+    "grpc-health-check": "^2.1.0",
     "gts": "^6.0.2",
     "jsdoc": "^4.0.4",
     "jsdoc-fresh": "^5.0.0",
@@ -140,5 +144,8 @@
     "typescript": "^5.8.2",
     "yargs": "^17.7.2"
   },
+  "overrides": {
+    "@opentelemetry/resources": "^1.8.0"
+  },
   "homepage": "https://github.com/googleapis/google-cloud-node/tree/main/handwritten/spanner"
-}
+}