feat(loadBalancer)!: add support for load balancers

Author: durran

src/cmap/connect.ts

@@ -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;

src/cmap/connection.ts

@@ -30,7 +30,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';
@@ -233,6 +233,10 @@ export class Connection extends EventEmitter {
     return this[kIsMaster];
   }
 
+  get serverId(): ObjectId {
+    return ismaster().serverId;
+  }
+
   // the `connect` method stores the result of the handshake ismaster on the connection
   set ismaster(response: Document) {
     this[kDescription].receiveResponse(response);
@@ -617,8 +621,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,7 +18,7 @@ import {
   ConnectionCheckedInEvent,
   ConnectionPoolClearedEvent
 } from './events';
-import type { Document } from '../bson';
+import type { Document, ObjectId } from '../bson';
 
 const kLogger = Symbol('logger');
 const kConnections = Symbol('connections');
@@ -271,6 +271,17 @@ export class ConnectionPool extends EventEmitter {
     this.emit('connectionPoolCleared', new ConnectionPoolClearedEvent(this));
   }
 
+  /**
+   * Close all connections in the pool for the provided serverId.
+   */
+  closeConnnections(serverId: ObjectId): void {
+    const indexes = this[kConnections].toArray().forEach((connection) => {
+      if (connection.serverId === serverId) {
+        destroyConnection(this, connection, 'serverError')
+      }
+    });
+  }
+
   /** Close the pool */
   close(callback: Callback<void>): void;
   close(options: CloseOptions, callback: Callback<void>): void;
@@ -410,6 +421,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

@@ -29,6 +29,8 @@ import { Logger, LoggerLevel } from './logger';
 import { PromiseProvider } from './promise_provider';
 import { createAutoEncrypter } from './operations/connect';
 
+const VALID_TXT_RECORDS = ['authSource', 'replicaSet', 'loadBalanced'];
+
 /**
  * Determines whether a provided address matches the provided parent domain in order
  * to avoid certain attack vectors.
@@ -94,9 +96,9 @@ 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(', ')}`)
           );
         }
 
@@ -116,6 +118,8 @@ export function resolveSRVRecord(options: MongoOptions, callback: Callback<HostA
         }
       }
 
+      validateLoadBalancedOptions(hostAddresses, options);
+
       callback(undefined, hostAddresses);
     });
   });
@@ -426,6 +430,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');
@@ -435,6 +441,20 @@ export function parseOptions(
   return mongoOptions;
 }
 
+function validateLoadBalancedOptions(hosts: HostAddress[] | string[], mongoOptions: MongoOptions) {
+  if (mongoOptions.loadBalanced) {
+    if (hosts.length > 1) {
+      throw new MongoParseError('loadBalanced option only supported with a single host in the URI');
+    }
+    if (mongoOptions.replicaSet) {
+      throw new MongoParseError('loadBalanced option not supported with a replicaSet option');
+    }
+    if (mongoOptions.directConnection) {
+      throw new MongoParseError('loadBalanced option not supported when directConnection is true');
+    }
+  }
+}
+
 function setOption(
   mongoOptions: any,
   key: string,
@@ -696,6 +716,10 @@ export const OPTIONS = {
     default: 120000,
     type: 'uint'
   },
+  loadBalanced: {
+    default: false,
+    type: 'boolean'
+  },
   localThresholdMS: {
     default: 15,
     type: 'uint'

src/mongo_client.ts

@@ -170,6 +170,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;
@@ -600,6 +602,7 @@ export interface MongoOptions
         | 'heartbeatFrequencyMS'
         | 'keepAlive'
         | 'keepAliveInitialDelay'
+        | 'loadBalanced'
         | 'localThresholdMS'
         | 'logger'
         | 'maxIdleTimeMS'

src/sdam/common.ts

@@ -17,7 +17,8 @@ export enum TopologyType {
   ReplicaSetNoPrimary = 'ReplicaSetNoPrimary',
   ReplicaSetWithPrimary = 'ReplicaSetWithPrimary',
   Sharded = 'Sharded',
-  Unknown = 'Unknown'
+  Unknown = 'Unknown',
+  LoadBalanced = 'LoadBalanced'
 }
 
 /**
@@ -33,7 +34,8 @@ export enum ServerType {
   RSArbiter = 'RSArbiter',
   RSOther = 'RSOther',
   RSGhost = 'RSGhost',
-  Unknown = 'Unknown'
+  Unknown = 'Unknown',
+  LoadBalancer = 'LoadBalancer'
 }
 
 /** @internal */

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
       ? {

src/sdam/server.ts

@@ -172,7 +172,8 @@ export class Server extends EventEmitter {
       this.emit(
         Server.DESCRIPTION_RECEIVED,
         new ServerDescription(this.description.hostAddress, event.reply, {
-          roundTripTime: calculateRoundTripTime(this.description.roundTripTime, event.duration)
+          roundTripTime: calculateRoundTripTime(this.description.roundTripTime, event.duration),
+          loadBalanced: options.loadBalanced
         })
       );
 
@@ -297,6 +298,7 @@ export class Server extends EventEmitter {
       return;
     }
 
+    // TODO: Durran: Alert the connection pool to the type of operation.
     this.s.pool.withConnection((err, conn, cb) => {
       if (err || !conn) {
         markServerUnknown(this, err);
@@ -322,6 +324,7 @@ export class Server extends EventEmitter {
       return;
     }
 
+    // TODO: Durran: Alert the connection pool to the type of operation.
     this.s.pool.withConnection((err, conn, cb) => {
       if (err || !conn) {
         markServerUnknown(this, err);
@@ -347,6 +350,7 @@ export class Server extends EventEmitter {
       return;
     }
 
+    // TODO: Durran: Alert the connection pool to the type of operation.
     this.s.pool.withConnection((err, conn, cb) => {
       if (err || !conn) {
         markServerUnknown(this, err);
@@ -380,6 +384,7 @@ export class Server extends EventEmitter {
       return;
     }
 
+    // TODO: Durran: Alert the connection pool to the type of operation.
     this.s.pool.withConnection((err, conn, cb) => {
       if (err || !conn) {
         markServerUnknown(this, err);

src/sdam/server_description.ts

@@ -2,18 +2,21 @@ import { arrayStrictEqual, errorStrictEqual, now, HostAddress } from '../utils';
 import { ServerType } from './common';
 import type { ObjectId, Long, Document } from '../bson';
 import type { ClusterTime } from './common';
+import { MongoError } from '../error';
 
 const WRITABLE_SERVER_TYPES = new Set([
   ServerType.RSPrimary,
   ServerType.Standalone,
-  ServerType.Mongos
+  ServerType.Mongos,
+  ServerType.LoadBalancer
 ]);
 
 const DATA_BEARING_SERVER_TYPES = new Set([
   ServerType.RSPrimary,
   ServerType.RSSecondary,
   ServerType.Mongos,
-  ServerType.Standalone
+  ServerType.Standalone,
+  ServerType.LoadBalancer
 ]);
 
 /** @public */
@@ -35,6 +38,9 @@ export interface ServerDescriptionOptions {
 
   /** The topologyVersion */
   topologyVersion?: TopologyVersion;
+
+  /** If the client is in load balancing mode. */
+  loadBalanced?: boolean;
 }
 
 /**
@@ -89,7 +95,7 @@ export class ServerDescription {
       this._hostAddress = address;
       this.address = this._hostAddress.toString();
     }
-    this.type = parseServerType(ismaster);
+    this.type = parseServerType(ismaster, options);
     this.hosts = ismaster?.hosts?.map((host: string) => host.toLowerCase()) ?? [];
     this.passives = ismaster?.passives?.map((host: string) => host.toLowerCase()) ?? [];
     this.arbiters = ismaster?.arbiters?.map((host: string) => host.toLowerCase()) ?? [];
@@ -205,11 +211,23 @@ export class ServerDescription {
 }
 
 // Parses an `ismaster` message and determines the server type
-export function parseServerType(ismaster?: Document): ServerType {
+export function parseServerType(
+  ismaster?: Document,
+  options?: ServerDescriptionOptions
+): ServerType {
   if (!ismaster || !ismaster.ok) {
     return ServerType.Unknown;
   }
 
+  if (options && options.loadBalanced) {
+    if (!ismaster.serverId) {
+      throw new MongoError(
+        'Driver attempted to initialize in load balancing mode, but the server does not support it.'
+      );
+    }
+    return ServerType.LoadBalancer;
+  }
+
   if (ismaster.isreplicaset) {
     return ServerType.RSGhost;
   }