feat(loadBalancer)!: add support for load balancers

Author: durran

src/cmap/connect.ts

@@ -11,7 +11,7 @@ import {
   MIN_SUPPORTED_WIRE_VERSION,
   MIN_SUPPORTED_SERVER_VERSION
 } from './wire_protocol/constants';
-import type { Document } from '../bson';
+import type { Document, ObjectId } from '../bson';
 
 import type { Socket, SocketConnectOpts } from 'net';
 import type { TLSSocket, ConnectionOptions as TLSConnectionOpts } from 'tls';
@@ -149,6 +149,7 @@ export interface HandshakeDocument extends Document {
   client: ClientMetadata;
   compression: string[];
   saslSupportedMechs?: string;
+  loadBalanced?: boolean;
 }
 
 function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<HandshakeDocument>) {
@@ -158,7 +159,8 @@ function prepareHandshakeDocument(authContext: AuthContext, callback: Callback<H
   const handshakeDoc: HandshakeDocument = {
     ismaster: true,
     client: options.metadata || makeClientMetadata(options),
-    compression: compressors
+    compression: compressors,
+    loadBalanced: options.loadBalanced
   };
 
   const credentials = authContext.credentials;
@@ -295,7 +297,7 @@ function makeConnection(options: ConnectionOptions, _callback: CallbackWithType<
   socket.setNoDelay(noDelay);
 
   const connectEvent = useTLS ? 'secureConnect' : 'connect';
-  let cancellationHandler: (err: Error) => void;
+  let cancellationHandler: (err: Error, serviceId?: ObjectId) => void;
   function errorHandler(eventName: ErrorHandlerEventName) {
     return (err: Error) => {
       SOCKET_ERROR_EVENTS.forEach(event => socket.removeAllListeners(event));

src/cmap/connection.ts

@@ -34,7 +34,7 @@ import {
   OpQueryOptions,
   Msg
 } from './commands';
-import { BSONSerializeOptions, Document, Long, pluckBSONSerializeOptions } from '../bson';
+import { BSONSerializeOptions, Document, Long, pluckBSONSerializeOptions, ObjectId } from '../bson';
 import type { AutoEncrypter } from '../deps';
 import type { MongoCredentials } from './auth/mongo_credentials';
 import type { Stream } from './connect';
@@ -220,6 +220,10 @@ export class Connection extends EventEmitter {
     this[kIsMaster] = response;
   }
 
+  get serviceIdd(): ObjectId {
+    return this.ismaster.serviceId;
+  }
+
   get generation(): number {
     return this[kGeneration] || 0;
   }
@@ -632,8 +636,10 @@ export class CryptoConnection extends Connection {
   }
 }
 
-function hasSessionSupport(conn: Connection) {
-  return conn.description.logicalSessionTimeoutMinutes != null;
+/** @public */
+export function hasSessionSupport(conn: Connection): boolean {
+  const description = conn.description;
+  return description.logicalSessionTimeoutMinutes != null || !!description.loadBalanced;
 }
 
 function supportsOpMsg(conn: Connection) {

src/cmap/connection_pool.ts

@@ -18,6 +18,7 @@ import {
   ConnectionCheckedInEvent,
   ConnectionPoolClearedEvent
 } from './connection_pool_events';
+import type { ObjectId } from '../bson';
 
 const kLogger = Symbol('logger');
 const kConnections = Symbol('connections');
@@ -270,6 +271,17 @@ export class ConnectionPool extends EventEmitter {
     this.emit('connectionPoolCleared', new ConnectionPoolClearedEvent(this));
   }
 
+  /**
+   * Close all connections in the pool for the provided serviceId.
+   */
+  closeConnections(serviceId: ObjectId): void {
+    // cancel in flight matching connections.
+    this[kCancellationToken].emit('cancel', serviceId);
+
+    // destroy each matching connection
+    // this.destroyConnections('serverError', {}, serviceId, callback);
+  }
+
   /** Close the pool */
   close(callback: Callback<void>): void;
   close(options: CloseOptions, callback: Callback<void>): void;
@@ -314,19 +326,25 @@ export class ConnectionPool extends EventEmitter {
 
     // mark the pool as closed immediately
     this.closed = true;
-
     eachAsync<Connection>(
       this[kConnections].toArray(),
       (conn, cb) => {
+        // Destroy the connection in the case of closing the entire pool
+        // or if the connection matches the server id.
+        //if (!serviceId || serviceId === conn.serviceId) {
         this.emit(
           ConnectionPool.CONNECTION_CLOSED,
           new ConnectionClosedEvent(this, conn, 'poolClosed')
         );
         conn.destroy(options, cb);
+        //}
       },
       err => {
+        // Dont close the entire pool for error on single server.
+        //if (!serviceId) {
         this[kConnections].clear();
         this.emit(ConnectionPool.CONNECTION_POOL_CLOSED, new ConnectionPoolClosedEvent(this));
+        //}
         callback(err);
       }
     );
@@ -382,6 +400,7 @@ function connectionIsIdle(pool: ConnectionPool, connection: Connection) {
   return !!(pool.options.maxIdleTimeMS && connection.idleTime > pool.options.maxIdleTimeMS);
 }
 
+// TODO: Durran: In LB mode set the server id on the connection.
 function createConnection(pool: ConnectionPool, callback?: Callback<Connection>) {
   const connectOptions: ConnectionOptions = {
     ...pool.options,

src/cmap/stream_description.ts

@@ -15,6 +15,8 @@ const RESPONSE_FIELDS = [
 /** @public */
 export interface StreamDescriptionOptions {
   compressors?: CompressorName[];
+  logicalSessionTimeoutMinutes?: number;
+  loadBalanced?: boolean;
 }
 
 /** @public */
@@ -29,6 +31,7 @@ export class StreamDescription {
   compressors: CompressorName[];
   compressor?: CompressorName;
   logicalSessionTimeoutMinutes?: number;
+  loadBalanced?: boolean;
 
   __nodejs_mock_server__?: boolean;
 
@@ -42,6 +45,8 @@ export class StreamDescription {
     this.maxBsonObjectSize = 16777216;
     this.maxMessageSizeBytes = 48000000;
     this.maxWriteBatchSize = 100000;
+    this.logicalSessionTimeoutMinutes = options ? options.logicalSessionTimeoutMinutes : undefined;
+    this.loadBalanced = options ? !!options.loadBalanced : false;
     this.compressors =
       options && options.compressors && Array.isArray(options.compressors)
         ? options.compressors

src/connection_string.ts

@@ -31,6 +31,12 @@ import { Logger, LoggerLevelId } from './logger';
 import { PromiseProvider } from './promise_provider';
 import { Encrypter } from './encrypter';
 
+const VALID_TXT_RECORDS = ['authSource', 'replicaSet', 'loadBalanced'];
+
+const LB_SINGLE_HOST = 'loadBalanced option only supported with a single host in the URI';
+const LB_REPLICA_SET = 'loadBalanced option not supported with a replicaSet option';
+const LB_DIRECT_CONNECTION = 'loadBalanced option not supported when directConnection is provided';
+
 /**
  * Determines whether a provided address matches the provided parent domain in order
  * to avoid certain attack vectors.
@@ -83,6 +89,11 @@ export function resolveSRVRecord(options: MongoOptions, callback: Callback<HostA
       HostAddress.fromString(`${r.name}:${r.port ?? 27017}`)
     );
 
+    const lbError = checkLoadBalancerOptions(options, hostAddresses);
+    if (lbError) {
+      return callback(lbError);
+    }
+
     // Resolve TXT record and add options from there if they exist.
     dns.resolveTxt(lookupAddress, (err, record) => {
       if (err) {
@@ -96,14 +107,15 @@ export function resolveSRVRecord(options: MongoOptions, callback: Callback<HostA
 
         const txtRecordOptions = new URLSearchParams(record[0].join(''));
         const txtRecordOptionKeys = [...txtRecordOptions.keys()];
-        if (txtRecordOptionKeys.some(key => key !== 'authSource' && key !== 'replicaSet')) {
+        if (txtRecordOptionKeys.some(key => !VALID_TXT_RECORDS.includes(key))) {
           return callback(
-            new MongoParseError('Text record must only set `authSource` or `replicaSet`')
+            new MongoParseError(`Text record must only set one of: ${VALID_TXT_RECORDS.join(', ')}`)
           );
         }
 
         const source = txtRecordOptions.get('authSource') ?? undefined;
         const replicaSet = txtRecordOptions.get('replicaSet') ?? undefined;
+        const loadBalanced = txtRecordOptions.get('loadBalanced') ?? undefined;
 
         if (source === '' || replicaSet === '') {
           return callback(new MongoParseError('Cannot have empty URI params in DNS TXT Record'));
@@ -116,13 +128,49 @@ export function resolveSRVRecord(options: MongoOptions, callback: Callback<HostA
         if (!options.userSpecifiedReplicaSet && replicaSet) {
           options.replicaSet = replicaSet;
         }
+
+        if (loadBalanced === 'true') {
+          options.loadBalanced = true;
+        }
+
+        const lbError = checkLoadBalancerOptions(options, hostAddresses);
+        if (lbError) {
+          return callback(lbError);
+        }
+
+        if (options.loadBalanced) {
+          if (options.replicaSet) {
+            return callback(new MongoParseError(LB_REPLICA_SET));
+          }
+
+          if (hostAddresses.length > 1) {
+            return callback(new MongoParseError(LB_SINGLE_HOST));
+          }
+        }
       }
 
       callback(undefined, hostAddresses);
     });
   });
 }
 
+function checkLoadBalancerOptions(
+  options: MongoOptions,
+  addresses: HostAddress[]
+): MongoParseError | null {
+  if (options.loadBalanced) {
+    if (options.replicaSet) {
+      return new MongoParseError(LB_REPLICA_SET);
+    }
+
+    if (addresses.length > 1) {
+      return new MongoParseError(LB_SINGLE_HOST);
+    }
+    return null;
+  }
+  return null;
+}
+
 /**
  * Checks if TLS options are valid
  *
@@ -431,6 +479,8 @@ export function parseOptions(
     throw new MongoParseError('directConnection not supported with SRV URI');
   }
 
+  validateLoadBalancedOptions(hosts, mongoOptions);
+
   // Potential SRV Overrides
   mongoOptions.userSpecifiedAuthSource =
     objectOptions.has('authSource') || urlOptions.has('authSource');
@@ -446,6 +496,20 @@ export function parseOptions(
   return mongoOptions;
 }
 
+function validateLoadBalancedOptions(hosts: HostAddress[] | string[], mongoOptions: MongoOptions) {
+  if (mongoOptions.loadBalanced) {
+    if (hosts.length > 1) {
+      throw new MongoParseError(LB_SINGLE_HOST);
+    }
+    if (mongoOptions.replicaSet) {
+      throw new MongoParseError(LB_REPLICA_SET);
+    }
+    if (mongoOptions.directConnection) {
+      throw new MongoParseError(LB_DIRECT_CONNECTION);
+    }
+  }
+}
+
 function setOption(
   mongoOptions: any,
   key: string,
@@ -706,6 +770,10 @@ export const OPTIONS = {
     default: 120000,
     type: 'uint'
   },
+  loadBalanced: {
+    default: false,
+    type: 'boolean'
+  },
   localThresholdMS: {
     default: 15,
     type: 'uint'

src/mongo_client.ts

@@ -174,6 +174,8 @@ export interface MongoClientOptions extends BSONSerializeOptions, SupportedNodeC
   retryWrites?: boolean;
   /** Allow a driver to force a Single topology type with a connection string containing one host */
   directConnection?: boolean;
+  /** Instruct the driver it is connecting to a load balancer fronting a mongos like service */
+  loadBalanced?: boolean;
 
   /** The write concern w value */
   w?: W;
@@ -623,6 +625,7 @@ export interface MongoOptions
         | 'heartbeatFrequencyMS'
         | 'keepAlive'
         | 'keepAliveInitialDelay'
+        | 'loadBalanced'
         | 'localThresholdMS'
         | 'logger'
         | 'maxIdleTimeMS'

src/sdam/common.ts

@@ -17,7 +17,8 @@ export const TopologyType = Object.freeze({
   ReplicaSetNoPrimary: 'ReplicaSetNoPrimary',
   ReplicaSetWithPrimary: 'ReplicaSetWithPrimary',
   Sharded: 'Sharded',
-  Unknown: 'Unknown'
+  Unknown: 'Unknown',
+  LoadBalanced: 'LoadBalanced'
 } as const);
 
 /** @public */
@@ -36,7 +37,8 @@ export const ServerType = Object.freeze({
   RSArbiter: 'RSArbiter',
   RSOther: 'RSOther',
   RSGhost: 'RSGhost',
-  Unknown: 'Unknown'
+  Unknown: 'Unknown',
+  LoadBalancer: 'LoadBalancer'
 } as const);
 
 /** @public */

src/sdam/monitor.ts

@@ -53,6 +53,7 @@ export interface MonitorOptions
   connectTimeoutMS: number;
   heartbeatFrequencyMS: number;
   minHeartbeatFrequencyMS: number;
+  loadBalanced?: boolean;
 }
 
 /** @public */
@@ -61,7 +62,10 @@ export class Monitor extends EventEmitter {
   s: MonitorPrivate;
   address: string;
   options: Readonly<
-    Pick<MonitorOptions, 'connectTimeoutMS' | 'heartbeatFrequencyMS' | 'minHeartbeatFrequencyMS'>
+    Pick<
+      MonitorOptions,
+      'connectTimeoutMS' | 'heartbeatFrequencyMS' | 'minHeartbeatFrequencyMS' | 'loadBalanced'
+    >
   >;
   connectOptions: ConnectionOptions;
   [kServer]: Server;
@@ -87,7 +91,8 @@ export class Monitor extends EventEmitter {
     this.options = Object.freeze({
       connectTimeoutMS: options.connectTimeoutMS ?? 10000,
       heartbeatFrequencyMS: options.heartbeatFrequencyMS ?? 10000,
-      minHeartbeatFrequencyMS: options.minHeartbeatFrequencyMS ?? 500
+      minHeartbeatFrequencyMS: options.minHeartbeatFrequencyMS ?? 500,
+      loadBalanced: options.loadBalanced ?? false
     });
 
     const cancellationToken = this[kCancellationToken];
@@ -211,14 +216,20 @@ function checkServer(monitor: Monitor, callback: Callback<Document>) {
   const connection = monitor[kConnection];
   if (connection && !connection.closed) {
     const connectTimeoutMS = monitor.options.connectTimeoutMS;
+    const loadBalanced = monitor.options.loadBalanced;
     const maxAwaitTimeMS = monitor.options.heartbeatFrequencyMS;
     const topologyVersion = monitor[kServer].description.topologyVersion;
     const isAwaitable = topologyVersion != null;
 
     const cmd =
       isAwaitable && topologyVersion
-        ? { ismaster: true, maxAwaitTimeMS, topologyVersion: makeTopologyVersion(topologyVersion) }
-        : { ismaster: true };
+        ? {
+            ismaster: true,
+            maxAwaitTimeMS,
+            topologyVersion: makeTopologyVersion(topologyVersion),
+            loadBalanced: loadBalanced
+          }
+        : { ismaster: true, loadBalanced: loadBalanced };
 
     const options = isAwaitable
       ? {