feat(loadBalancer)!: add support for load balancers

Author: durran

.evergreen/config.yml

@@ -131,6 +131,40 @@ 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} 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} \
+            bash ${PROJECT_DIRECTORY}/.evergreen/run-tests.sh
   run checks:
     - command: shell.exec
       type: test

.evergreen/run-load-balancer.sh

@@ -0,0 +1,83 @@
+#!/bin/bash
+
+set -o errexit  # Exit the script with error if any of the commands fail
+
+start() {
+  echo "Starting HAProxy..."
+
+  cat <<EOF_HAPROXY_CONFIG >> $DRIVERS_TOOLS/haproxy.conf
+  defaults
+      mode tcp
+      timeout connect 10s
+      timeout client 30m
+      timeout server 30m
+
+  frontend mongos_frontend
+      bind *:8000
+      use_backend mongos_backend
+
+  frontend mongoses_frontend
+      bind *:8001
+      use_backend mongoses_backend
+
+  backend mongos_backend
+      mode tcp
+      server mongos 127.0.0.1:27017 check
+
+  backend mongoses_backend
+      mode tcp
+      server mongos_one 127.0.0.1:27017 check
+      server mongos_two 127.0.0.1:27018 check
+EOF_HAPROXY_CONFIG
+
+  PREFIX=$(echo $MONGODB_URI | grep -Eo "(.*?)@" | cat)
+  SUFFIX=$(echo $MONGODB_URI | grep -Eo "\?(.*)" | cat)
+
+  if [[ $PREFIX = "" ]]
+  then
+    # No auth then just set the URI
+    SINGLE_MONGOS_LB_URI="mongodb://127.0.0.1:8000/"
+    MULTI_MONGOS_LB_URI="mongodb://127.0.0.1:8001/"
+  else
+    # We have auth so append the lb host:port
+    SINGLE_MONGOS_LB_URI="${PREFIX}127.0.0.1:8000/"
+    MULTI_MONGOS_LB_URI="${PREFIX}127.0.0.1:8001/"
+  fi
+
+  if [[ $SUFFIX = "" ]]
+  then
+    # If there are no query params then add only the load balanced option.
+    SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}?loadBalanced=true"
+    MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}?loadBalanced=true"
+  else
+    # If there are query params then append the load balanced option to them.
+    SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}${SUFFIX}&loadBalanced=true"
+    MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}${SUFFIX}&loadBalanced=true"
+  fi
+
+  echo "Single Mongos LB: $SINGLE_MONGOS_LB_URI"
+  echo "Multiple Mongos LB: $MULTI_MONGOS_LB_URI"
+
+  /usr/sbin/haproxy -D -f $DRIVERS_TOOLS/haproxy.conf -p $DRIVERS_TOOLS/haproxy.pid
+
+  echo 'SINGLE_MONGOS_LB_URI: "'$SINGLE_MONGOS_LB_URI'"' > lb-expansion.yml
+  echo 'MULTI_MONGOS_LB_URI: "'$MULTI_MONGOS_LB_URI'"' >> lb-expansion.yml
+}
+
+stop() {
+  echo "Stopping HAProxy..."
+  kill -USR1 $(cat $DRIVERS_TOOLS/haproxy.pid)
+  rm $DRIVERS_TOOLS/haproxy.conf $DRIVERS_TOOLS/haproxy.pid
+}
+
+case "$1" in
+  start)
+    start
+    ;;
+  stop)
+    stop
+    ;;
+  *)
+    echo "Usage: load-balancer.sh (start|stop)"
+    exit 1
+esac

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