feat(NODE-3011): Load Balancer Support (#2909)

Author: durran

.evergreen/config.yml

@@ -134,6 +134,44 @@ functions:
           MONGODB_API_VERSION="${MONGODB_API_VERSION}" \
           NODE_VERSION=${NODE_VERSION} SKIP_DEPS=${SKIP_DEPS|1} NO_EXIT=${NO_EXIT|1} \
             bash ${PROJECT_DIRECTORY}/.evergreen/run-tests.sh
+  start-load-balancer:
+    - command: shell.exec
+      params:
+        script: |
+          DRIVERS_TOOLS=${DRIVERS_TOOLS} MONGODB_URI=${MONGODB_URI} \
+            bash ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh start
+    - command: expansions.update
+      params:
+        file: lb-expansion.yml
+  stop-load-balancer:
+    - command: shell.exec
+      params:
+        script: |
+          DRIVERS_TOOLS=${DRIVERS_TOOLS} MONGODB_URI=${MONGODB_URI} \
+          bash ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh stop
+  run-lb-tests:
+    - command: shell.exec
+      type: test
+      params:
+        working_dir: src
+        timeout_secs: 60
+        script: |
+          ${PREPARE_SHELL}
+
+          MONGODB_URI="${MONGODB_URI}" \
+          AUTH=${AUTH} \
+          SSL=${SSL} \
+          UNIFIED=${UNIFIED} \
+          MONGODB_API_VERSION="${MONGODB_API_VERSION}" \
+          NODE_VERSION=${NODE_VERSION} \
+          SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}" \
+          MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}" \
+          TOPOLOGY="${TOPOLOGY}" \
+          SKIP_DEPS=${SKIP_DEPS|1} \
+          NO_EXIT=${NO_EXIT|1} \
+          TEST_NPM_SCRIPT="check:load-balancer" \
+          FAKE_MONGODB_SERVICE_ID="true" \
+            bash ${PROJECT_DIRECTORY}/.evergreen/run-tests.sh
   run checks:
     - command: shell.exec
       type: test
@@ -937,6 +975,20 @@ tasks:
       - func: install dependencies
       - func: bootstrap mongohoused
       - func: run data lake tests
+  - name: test-load-balancer
+    tags:
+      - latest
+      - sharded_cluster
+      - load_balancer
+    commands:
+      - func: install dependencies
+      - func: bootstrap mongo-orchestration
+        vars:
+          VERSION: '5.0'
+          TOPOLOGY: sharded_cluster
+      - func: start-load-balancer
+      - func: run-lb-tests
+      - func: stop-load-balancer
   - name: test-auth-kerberos
     tags:
       - auth
@@ -1760,6 +1812,7 @@ buildvariants:
       - test-latest-server-v1-api
       - test-atlas-connectivity
       - test-atlas-data-lake
+      - test-load-balancer
       - test-auth-kerberos
       - test-auth-ldap
       - test-ocsp-valid-cert-server-staples

.evergreen/config.yml.in

@@ -155,6 +155,47 @@ functions:
           NODE_VERSION=${NODE_VERSION} SKIP_DEPS=${SKIP_DEPS|1} NO_EXIT=${NO_EXIT|1} \
             bash ${PROJECT_DIRECTORY}/.evergreen/run-tests.sh
 
+  "start-load-balancer":
+    - command: shell.exec
+      params:
+        script: |
+          DRIVERS_TOOLS=${DRIVERS_TOOLS} MONGODB_URI=${MONGODB_URI} \
+            bash ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh start
+    - command: expansions.update
+      params:
+        file: lb-expansion.yml
+
+  "stop-load-balancer":
+    - command: shell.exec
+      params:
+        script: |
+          DRIVERS_TOOLS=${DRIVERS_TOOLS} MONGODB_URI=${MONGODB_URI} \
+          bash ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh stop
+
+  "run-lb-tests":
+    - command: shell.exec
+      type: test
+      params:
+        working_dir: src
+        timeout_secs: 60
+        script: |
+          ${PREPARE_SHELL}
+
+          MONGODB_URI="${MONGODB_URI}" \
+          AUTH=${AUTH} \
+          SSL=${SSL} \
+          UNIFIED=${UNIFIED} \
+          MONGODB_API_VERSION="${MONGODB_API_VERSION}" \
+          NODE_VERSION=${NODE_VERSION} \
+          SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}" \
+          MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}" \
+          TOPOLOGY="${TOPOLOGY}" \
+          SKIP_DEPS=${SKIP_DEPS|1} \
+          NO_EXIT=${NO_EXIT|1} \
+          TEST_NPM_SCRIPT="check:load-balancer" \
+          FAKE_MONGODB_SERVICE_ID="true" \
+            bash ${PROJECT_DIRECTORY}/.evergreen/run-tests.sh
+
   "run checks":
     - command: shell.exec
       type: test

.evergreen/generate_evergreen_tasks.js

@@ -42,7 +42,8 @@ const OPERATING_SYSTEMS = [
 }));
 
 // TODO: NODE-3060: enable skipped tests on windows
-const WINDOWS_SKIP_TAGS = new Set(['atlas-connect', 'auth']);
+const WINDOWS_SKIP_TAGS = new Set(['atlas-connect', 'auth', 'load_balancer']);
+const MACOS_SKIP_TAGS = new Set(['load_balancer']);
 
 const TASKS = [];
 const SINGLETON_TASKS = [];
@@ -107,6 +108,23 @@ TASKS.push(...[
       { func: 'run data lake tests' }
     ]
   },
+  {
+    name: 'test-load-balancer',
+    tags: ['latest', 'sharded_cluster', 'load_balancer'],
+    commands: [
+      { func: 'install dependencies' },
+      {
+        func: 'bootstrap mongo-orchestration',
+        vars: {
+          VERSION: '5.0',
+          TOPOLOGY: 'sharded_cluster'
+        }
+      },
+      { func: 'start-load-balancer' },
+      { func: 'run-lb-tests' },
+      { func: 'stop-load-balancer' }
+    ]
+  },
   {
     name: 'test-auth-kerberos',
     tags: ['auth', 'kerberos'],
@@ -429,11 +447,12 @@ const getTaskList = (() => {
       .filter(task => {
         if (task.name.match(/^aws/)) return false;
 
-        // skip unsupported tasks on windows
+        // skip unsupported tasks on windows or macos
         if (
-          os.match(/^windows/) &&
-          task.tags &&
-          task.tags.filter(tag => WINDOWS_SKIP_TAGS.has(tag)).length
+          task.tags && (
+            (os.match(/^windows/) && task.tags.filter(tag => WINDOWS_SKIP_TAGS.has(tag)).length) ||
+            (os.match(/^macos/) && task.tags.filter(tag => MACOS_SKIP_TAGS.has(tag)).length)
+          )
         ) {
           return false;
         }

.evergreen/run-tests.sh

@@ -49,4 +49,4 @@ else
   . $DRIVERS_TOOLS/.evergreen/csfle/set-temp-creds.sh
 fi
 
-MONGODB_API_VERSION=${MONGODB_API_VERSION} MONGODB_UNIFIED_TOPOLOGY=${UNIFIED} MONGODB_URI=${MONGODB_URI} npm run ${TEST_NPM_SCRIPT}
+SINGLE_MONGOS_LB_URI=${SINGLE_MONGOS_LB_URI} MULTI_MONGOS_LB_URI=${MULTI_MONGOS_LB_URI} MONGODB_API_VERSION=${MONGODB_API_VERSION} MONGODB_UNIFIED_TOPOLOGY=${UNIFIED} MONGODB_URI=${MONGODB_URI} npm run ${TEST_NPM_SCRIPT}

package.json

@@ -121,6 +121,7 @@
     "check:ts": "tsc -v && tsc --noEmit",
     "check:atlas": "mocha --config \"test/manual/mocharc.json\" test/manual/atlas_connectivity.test.js",
     "check:adl": "mocha test/manual/data_lake.test.js",
+    "check:load-balancer": "mocha test/manual/load-balancer.test.js",
     "check:ocsp": "mocha --config \"test/manual/mocharc.json\" test/manual/ocsp_support.test.js",
     "check:kerberos": "mocha --config \"test/manual/mocharc.json\" test/manual/kerberos.test.js",
     "check:tls": "mocha --config \"test/manual/mocharc.json\" test/manual/tls_support.test.js",

src/cmap/command_monitoring_events.ts

@@ -1,8 +1,7 @@
 import { GetMore, KillCursor, Msg, WriteProtocolMessageType } from './commands';
 import { calculateDurationInMs, deepCopy } from '../utils';
-import type { ConnectionPool } from './connection_pool';
 import type { Connection } from './connection';
-import type { Document } from '../bson';
+import type { Document, ObjectId } from '../bson';
 
 /**
  * An event indicating the start of a given
@@ -17,6 +16,7 @@ export class CommandStartedEvent {
   command: Document;
   address: string;
   connectionId?: string | number;
+  serviceId?: ObjectId;
 
   /**
    * Create a started event
@@ -25,10 +25,10 @@ export class CommandStartedEvent {
    * @param pool - the pool that originated the command
    * @param command - the command
    */
-  constructor(pool: Connection | ConnectionPool, command: WriteProtocolMessageType) {
+  constructor(connection: Connection, command: WriteProtocolMessageType) {
     const cmd = extractCommand(command);
     const commandName = extractCommandName(cmd);
-    const { address, connectionId } = extractConnectionDetails(pool);
+    const { address, connectionId, serviceId } = extractConnectionDetails(connection);
 
     // TODO: remove in major revision, this is not spec behavior
     if (SENSITIVE_COMMANDS.has(commandName)) {
@@ -38,11 +38,17 @@ export class CommandStartedEvent {
 
     this.address = address;
     this.connectionId = connectionId;
+    this.serviceId = serviceId;
     this.requestId = command.requestId;
     this.databaseName = databaseName(command);
     this.commandName = commandName;
     this.command = maybeRedact(commandName, cmd, cmd);
   }
+
+  /* @internal */
+  get hasServiceId(): boolean {
+    return !!this.serviceId;
+  }
 }
 
 /**
@@ -57,6 +63,7 @@ export class CommandSucceededEvent {
   duration: number;
   commandName: string;
   reply: unknown;
+  serviceId?: ObjectId;
 
   /**
    * Create a succeeded event
@@ -68,22 +75,28 @@ export class CommandSucceededEvent {
    * @param started - a high resolution tuple timestamp of when the command was first sent, to calculate duration
    */
   constructor(
-    pool: Connection | ConnectionPool,
+    connection: Connection,
     command: WriteProtocolMessageType,
     reply: Document | undefined,
     started: number
   ) {
     const cmd = extractCommand(command);
     const commandName = extractCommandName(cmd);
-    const { address, connectionId } = extractConnectionDetails(pool);
+    const { address, connectionId, serviceId } = extractConnectionDetails(connection);
 
     this.address = address;
     this.connectionId = connectionId;
+    this.serviceId = serviceId;
     this.requestId = command.requestId;
     this.commandName = commandName;
     this.duration = calculateDurationInMs(started);
     this.reply = maybeRedact(commandName, cmd, extractReply(command, reply));
   }
+
+  /* @internal */
+  get hasServiceId(): boolean {
+    return !!this.serviceId;
+  }
 }
 
 /**
@@ -98,6 +111,8 @@ export class CommandFailedEvent {
   duration: number;
   commandName: string;
   failure: Error;
+  serviceId?: ObjectId;
+
   /**
    * Create a failure event
    *
@@ -108,23 +123,29 @@ export class CommandFailedEvent {
    * @param started - a high resolution tuple timestamp of when the command was first sent, to calculate duration
    */
   constructor(
-    pool: Connection | ConnectionPool,
+    connection: Connection,
     command: WriteProtocolMessageType,
     error: Error | Document,
     started: number
   ) {
     const cmd = extractCommand(command);
     const commandName = extractCommandName(cmd);
-    const { address, connectionId } = extractConnectionDetails(pool);
+    const { address, connectionId, serviceId } = extractConnectionDetails(connection);
 
     this.address = address;
     this.connectionId = connectionId;
+    this.serviceId = serviceId;
 
     this.requestId = command.requestId;
     this.commandName = commandName;
     this.duration = calculateDurationInMs(started);
     this.failure = maybeRedact(commandName, cmd, error) as Error;
   }
+
+  /* @internal */
+  get hasServiceId(): boolean {
+    return !!this.serviceId;
+  }
 }
 
 /** Commands that we want to redact because of the sensitive nature of their contents */
@@ -300,13 +321,14 @@ function extractReply(command: WriteProtocolMessageType, reply?: Document) {
   return deepCopy(reply.result ? reply.result : reply);
 }
 
-function extractConnectionDetails(connection: Connection | ConnectionPool) {
+function extractConnectionDetails(connection: Connection) {
   let connectionId;
   if ('id' in connection) {
     connectionId = connection.id;
   }
   return {
     address: connection.address,
+    serviceId: connection.serviceId,
     connectionId
   };
 }

src/cmap/connect.ts

@@ -20,10 +20,14 @@ import {
   MIN_SUPPORTED_SERVER_VERSION
 } from './wire_protocol/constants';
 import type { Document } from '../bson';
+import { Int32 } from '../bson';
 
 import type { Socket, SocketConnectOpts } from 'net';
 import type { TLSSocket, ConnectionOptions as TLSConnectionOpts } from 'tls';
-import { Int32 } from '../bson';
+
+const FAKE_MONGODB_SERVICE_ID =
+  typeof process.env.FAKE_MONGODB_SERVICE_ID === 'string' &&
+  process.env.FAKE_MONGODB_SERVICE_ID.toLowerCase() === 'true';
 
 /** @public */
 export type Stream = Socket | TLSSocket;
@@ -133,6 +137,21 @@ function performInitialHandshake(
         return;
       }
 
+      if (options.loadBalanced) {
+        // TODO: Durran: Remove when server support exists. (NODE-3431)
+        if (FAKE_MONGODB_SERVICE_ID) {
+          response.serviceId = response.topologyVersion.processId;
+        }
+        if (!response.serviceId) {
+          return callback(
+            new MongoDriverError(
+              'Driver attempted to initialize in load balancing mode, ' +
+                'but the server does not support this mode.'
+            )
+          );
+        }
+      }
+
       // NOTE: This is metadata attached to the connection while porting away from
       //       handshake being done in the `Server` class. Likely, it should be
       //       relocated, or at very least restructured.
@@ -172,6 +191,7 @@ export interface HandshakeDocument extends Document {
   client: ClientMetadata;
   compression: string[];
   saslSupportedMechs?: string;
+  loadBalanced: boolean;
 }
 
 function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<HandshakeDocument>) {
@@ -183,7 +203,8 @@ function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<H
     [serverApi?.version ? 'hello' : 'ismaster']: true,
     helloOk: true,
     client: options.metadata || makeClientMetadata(options),
-    compression: compressors
+    compression: compressors,
+    loadBalanced: options.loadBalanced
   };
 
   const credentials = authContext.credentials;