diff --git a/src/connection_string.ts b/src/connection_string.ts index e42a43dd92..e54d3483da 100644 --- a/src/connection_string.ts +++ b/src/connection_string.ts @@ -345,6 +345,12 @@ export function parseOptions( ); } + if (urlOptions.has('serverSelector')) { + throw new MongoParseError( + 'URI cannot contain `serverSelector`, it can only be passed to the client' + ); + } + if (urlOptions.has('srvHostValidator')) { throw new MongoParseError( 'URI cannot contain `srvHostValidator`, it can only be passed to the client' @@ -1187,6 +1193,14 @@ export const OPTIONS = { default: 30000, type: 'uint' }, + serverSelector: { + transform({ values: [value] }) { + if (typeof value !== 'function') { + throw new MongoParseError('serverSelector must be a function'); + } + return value; + } + }, servername: { type: 'string' }, diff --git a/src/index.ts b/src/index.ts index dd20169bb9..f16eed1747 100644 --- a/src/index.ts +++ b/src/index.ts @@ -593,7 +593,11 @@ export type { TagSet, TopologyVersion } from './sdam/server_description'; -export type { DeprioritizedServers, ServerSelector } from './sdam/server_selection'; +export type { + DeprioritizedServers, + ServerSelector, + ServerSelectorFunction +} from './sdam/server_selection'; export type { SrvPoller, SrvPollerEvents, SrvPollerOptions } from './sdam/srv_polling'; export type { ConnectOptions, diff --git a/src/mongo_client.ts b/src/mongo_client.ts index 28891d903d..61ae19bb2c 100644 --- a/src/mongo_client.ts +++ b/src/mongo_client.ts @@ -49,7 +49,11 @@ import { ReadPreference, type ReadPreferenceMode } from './read_preference'; import { type Runtime, type RuntimeAdapters } from './runtime_adapters'; import type { ServerMonitoringMode } from './sdam/monitor'; import type { TagSet } from './sdam/server_description'; -import { DeprioritizedServers, readPreferenceServerSelector } from './sdam/server_selection'; +import { + DeprioritizedServers, + readPreferenceServerSelector, + type ServerSelectorFunction +} from './sdam/server_selection'; import type { SrvPoller } from './sdam/srv_polling'; import { Topology, type TopologyEvents } from './sdam/topology'; import { ClientSession, type ClientSessionOptions, ServerSessionPool } from './sessions'; @@ -259,6 +263,104 @@ export interface MongoClientOptions extends BSONSerializeOptions, SupportedNodeC localThresholdMS?: number; /** Specifies how long (in milliseconds) to block for server selection before throwing an exception. */ serverSelectionTimeoutMS?: number; + /** + * Further filters the servers that are already eligible for an operation. + * + * The driver first applies the operation's read preference, tag sets, max + * staleness, and deprioritized-server rules. It then calls this function + * with copies of the servers that remain, and only after the function + * returns does it apply `localThresholdMS`. Array order is not a preference. + * Inside the latency window the driver still chooses by current operation + * count. Return a one-element array to send the operation to that member. + * Returning one slower member keeps it: that member becomes the low end of + * the window. Returning it together with a faster member can drop it again. + * + * Return only servers from the array you were given. A server from an + * earlier call, or any other address, throws `MongoInvalidArgumentError`. + * An empty array means none of the eligible servers are acceptable. + * Selection then waits until `serverSelectionTimeoutMS` and throws + * `MongoServerSelectionError`. Prefer `return preferred ? [preferred] : servers` + * when the member is a preference. Use `[]` only when the operation must + * not run anywhere else. + * + * The descriptions are copies. Editing `roundTripTime`, `tags`, `hosts`, + * `electionId`, `topologyVersion`, `$clusterTime`, or `error` does not + * change the topology. The driver selects the original server with the same + * address, not the object you return. Do not keep the copies for a later + * call. The next selection builds new ones, and a saved description whose + * address is no longer eligible is rejected. + * + * What a call is allowed to see depends on the operation: + * + * - `insertOne`, `updateOne`, `deleteOne`, and other ordinary writes are + * offered writable servers only. On a replica set that is the primary, so + * this function cannot send those writes to a secondary. + * - An aggregate that writes (`$out` or `$merge`) can be offered secondaries + * when its read preference allows that and the members are new enough. + * Returning one of those secondaries does send the aggregate there. + * - A `find` or a read aggregate is offered the servers that match its read + * preference. `primary` sees only the primary. `nearest` sees every + * suitable member, including members outside `localThresholdMS`. + * + * The function is not called for a Single topology (`directConnection`), a + * LoadBalanced topology, or when no server is eligible. It is not called for + * an operation pinned to a server (`getMore`, `killCursors`). It is not + * called for the connect handshake or for `endSessions`, so a function that + * returns an empty array can still `connect()` and `close()` can still end + * pooled sessions. + * + * Transactions on a replica set use read preference primary, so the function + * only sees the primary. On a sharded cluster the function runs for the + * first command in the transaction. The driver then pins that mongos, and + * later commands in the transaction reuse it without calling the function. + * + * This is not a connection-string option. Pass it to the `MongoClient` + * constructor. It matches PyMongo's `server_selector` keyword. + * + * @example Prefer one member, and keep the rest of the eligible list when it is absent + * ```ts + * // nearest reads that can see db-2 use db-2. Writes still see only the + * // primary, so they return that primary instead of waiting out the + * // selection timeout. + * new MongoClient(uri, { + * readPreference: 'nearest', + * serverSelector(servers) { + * const preferred = servers.find(server => server.address === 'db-2.example.com:27017'); + * return preferred ? [preferred] : servers; + * } + * }); + * ``` + * + * @example Require one member for operations that are allowed to use it + * ```ts + * // [] waits until serverSelectionTimeoutMS. A primary write is only offered + * // the primary, so this times out those writes while db-2 is a secondary. + * // Use it when a read must not fall back to another member. + * new MongoClient(uri, { + * readPreference: 'nearest', + * serverSelector(servers) { + * const member = servers.find(server => server.address === 'db-2.example.com:27017'); + * return member ? [member] : []; + * } + * }); + * ``` + * + * @example Choose one mongos, or one secondary for a writing aggregate + * ```ts + * // Every eligible mongos is passed in. Returning one sends the command + * // there and leaves the others for the next call that does not match. + * // The same shape works for an aggregate with `$out` or `$merge`: the + * // eligible list may contain secondaries, and returning one of them runs + * // that aggregate on it. + * new MongoClient(uri, { + * serverSelector(servers) { + * const chosen = servers.find(server => server.address === 'mongos-a.example.com:27017'); + * return chosen ? [chosen] : servers; + * } + * }); + * ``` + */ + serverSelector?: ServerSelectorFunction; /** heartbeatFrequencyMS controls when the driver checks the state of the MongoDB deployment. Specify the interval (in milliseconds) between checks, counted from the end of the previous check until the beginning of the next one. */ heartbeatFrequencyMS?: number; /** Sets the minimum heartbeat frequency. In the event that the driver has to frequently re-check a server's availability, it will wait at least this long since the previous check to avoid wasted effort. */ @@ -849,7 +951,8 @@ export class MongoClient extends TypedEventEmitter implements { description: topologyDescription }: Topology ) { // If we would attempt to select a server and get nothing back we short circuit - // to avoid the server selection timeout. + // to avoid the server selection timeout. This check ignores `serverSelector`: + // session cleanup must run even when the filter would select no server. const selector = readPreferenceServerSelector(ReadPreference.primaryPreferred); const serverDescriptions = Array.from(topologyDescription.servers.values()); const servers = selector(topologyDescription, serverDescriptions, new DeprioritizedServers()); @@ -1096,6 +1199,8 @@ export interface MongoOptions proxyUsername?: string; proxyPassword?: string; serverMonitoringMode: ServerMonitoringMode; + /** @see MongoClientOptions.serverSelector */ + serverSelector?: ServerSelectorFunction; /** @internal */ connectionType?: typeof Connection; /** @internal */ diff --git a/src/sdam/server_selection.ts b/src/sdam/server_selection.ts index 8f01beaa94..a10664b1f4 100644 --- a/src/sdam/server_selection.ts +++ b/src/sdam/server_selection.ts @@ -1,7 +1,15 @@ -import { MongoInvalidArgumentError, MongoRuntimeError } from '../error'; +import { Binary, Long, ObjectId, Timestamp } from '../bson'; +import { + MongoError, + MongoInvalidArgumentError, + MongoNetworkError, + type MongoNetworkErrorOptions, + MongoRuntimeError, + MongoServerError +} from '../error'; import { ReadPreference } from '../read_preference'; -import { ServerType, TopologyType } from './common'; -import type { ServerDescription, TagSet } from './server_description'; +import { type ClusterTime, ServerType, TopologyType } from './common'; +import type { ServerDescription, TagSet, TopologyVersion } from './server_description'; import type { TopologyDescription } from './topology_description'; // max staleness constants @@ -18,6 +26,31 @@ export type ServerSelector = ( deprioritized: DeprioritizedServers ) => ServerDescription[]; +/** + * Narrows the servers eligible for an operation. + * + * PyMongo exposes the same hook as `MongoClient(server_selector=...)`. + * The Node driver calls this after read preference, tag, and max-staleness + * filtering and before `localThresholdMS`. + * + * The argument is a list of copies, including `electionId`, `topologyVersion`, + * `$clusterTime`, and `error`. Return servers from that list. An empty array + * selects nothing and the operation waits until `serverSelectionTimeoutMS`. + * A server that was not in the list is rejected. + * + * @param servers - Copies of the servers that already satisfy the operation's selection rules. + * @returns The servers to consider. A one-element array targets that member. + * @public + */ +export type ServerSelectorFunction = (servers: ServerDescription[]) => ServerDescription[]; + +/** + * Selectors in this set choose one already-known server (cursor getMore and + * killCursors). The user filter must not move those operations. + * @internal + */ +const selectorsIgnoringUserServerSelector = new WeakSet(); + /** @internal */ export class DeprioritizedServers { private deprioritized: Set = new Set(); @@ -69,7 +102,7 @@ export function writableServerSelector(): ServerSelector { * if it is in a state that it can have commands sent to it. */ export function sameServerSelector(description?: ServerDescription): ServerSelector { - return function sameServerSelector( + const selector: ServerSelector = function sameServerSelector( _topologyDescription: TopologyDescription, servers: ServerDescription[], _deprioritized: DeprioritizedServers @@ -81,6 +114,8 @@ export function sameServerSelector(description?: ServerDescription): ServerSelec return sd.address === description.address && sd.type !== ServerType.Unknown; }); }; + selectorsIgnoringUserServerSelector.add(selector); + return selector; } /** @@ -253,6 +288,265 @@ function latencyWindowReducer( return servers.filter(server => server.roundTripTime <= high && server.roundTripTime >= low); } +const TOPOLOGIES_IGNORING_USER_SERVER_SELECTOR = new Set([ + TopologyType.Single, + TopologyType.LoadBalanced, + TopologyType.Unknown +]); + +/** + * Connect and session cleanup must be able to use any suitable server. A user + * filter that returns no server would otherwise time out `MongoClient.connect` + * and skip `endSessions` on close. + * @internal + */ +export function operationSkipsUserServerSelector(operationName: string | undefined): boolean { + return operationName === 'handshake' || operationName === 'endSessions'; +} + +/** + * `Object.assign` copies the description by reference for every object field. + * Several of those objects are the topology's own state, not a snapshot: + * + * - `TopologyDescription.updateRsFromPrimary` keeps `electionId` and later + * compares it with the next hello. + * - `isStaleServerDescription` reads `topologyVersion` off the stored + * description, so editing `counter` or `processId` can make a newer hello + * look stale. + * - `_advanceClusterTime` stores the hello `$clusterTime` object and later + * writes that same object on the wire. + * - `error` is the hello failure itself. `error.topologyVersion` is often the + * same object as `description.topologyVersion`. + * + * A caller that follows the public contract and edits any of those in place + * would change server selection and cluster-time gossip for every later + * operation. Replace each of them with a value the callback can edit freely. + * Selection still returns the original descriptions, matched by address. + */ +function copyServerDescription(server: ServerDescription): ServerDescription { + const copy: ServerDescription = Object.assign( + Object.create(Object.getPrototypeOf(server)), + server + ); + copy.hosts = server.hosts.slice(); + copy.passives = server.passives.slice(); + copy.arbiters = server.arbiters.slice(); + copy.tags = { ...server.tags }; + copy.electionId = server.electionId == null ? null : copyObjectId(server.electionId); + copy.topologyVersion = + server.topologyVersion == null ? null : copyTopologyVersion(server.topologyVersion); + copy.$clusterTime = copyClusterTime(server.$clusterTime); + copy.error = server.error == null ? null : copyError(server.error); + return copy; +} + +/** + * Builds a new ObjectId from a copy of the 12 id bytes. + * + * Current BSON stores an ObjectId as four packed integers, and + * `new ObjectId(existing)` copies those integers. `.id` also returns a new + * buffer on every read, then the setter writes into the instance's integers. + * Copying the bytes first does not depend on that layout: `copy.id = ...` + * changes the copy and leaves the id SDAM stored unchanged. + */ +function copyObjectId(id: ObjectId): ObjectId { + return new ObjectId(Uint8Array.from(id.id)); +} + +function copyTopologyVersion(version: TopologyVersion): TopologyVersion { + return { + processId: copyObjectId(version.processId), + // `Long.fromValue` allocates a new Long. Assigning `.low` on the shared + // counter would change staleness checks for the real description. + counter: Long.fromValue(version.counter) + }; +} + +function copyClusterTime(clusterTime: ClusterTime | undefined): ClusterTime | undefined { + if (clusterTime == null) { + return undefined; + } + // `t` and `i` are the two unsigned halves of the gossiped cluster time. + // `new Timestamp({ t, i })` stores them on a new Timestamp. The topology + // may later put this object on the wire, so the copy must not be the + // instance SDAM is still holding. `new Timestamp(existing)` also copies + // the bits in the current BSON library; the component form does not + // depend on that overload. + const copy: ClusterTime = { + clusterTime: new Timestamp({ t: clusterTime.clusterTime.t, i: clusterTime.clusterTime.i }) + }; + if (clusterTime.signature != null) { + copy.signature = { + // On Node, `value()` returns the Binary's own buffer, and + // `new Binary(existingBuffer)` keeps that buffer. Copy the bytes + // and preserve the subtype. `keyId` is a Long, so `fromValue` + // allocates one the callback can edit. + hash: new Binary( + Uint8Array.from(clusterTime.signature.hash.value()), + clusterTime.signature.hash.sub_type + ), + keyId: Long.fromValue(clusterTime.signature.keyId) + }; + } + return copy; +} + +/** + * `Object.assign` does not copy an Error usefully: `message` is often + * non-enumerable, and the label set is private. Build a new error and give + * it a new `topologyVersion`. + * + * A failed hello is a `MongoNetworkError` or `MongoNetworkTimeoutError`, and + * a server reply is a `MongoServerError`. Those keep their class so the + * callback can branch on them. Anything else becomes a `MongoError`. The + * server document is shallow-copied apart from `topologyVersion`. `cause` + * stays shared; it is not SDAM state. + */ +function copyError(error: MongoError): MongoError { + if (error instanceof MongoServerError) { + const response = { ...error.errorResponse }; + if (response.topologyVersion != null) { + response.topologyVersion = copyTopologyVersion(response.topologyVersion); + } + // The constructor points the new instance's `topologyVersion` at the + // document copy. `finishErrorCopy` replaces it, because the stored + // description often uses one object for `description.topologyVersion`, + // `error.topologyVersion`, and `errorResponse.topologyVersion`. + return finishErrorCopy(error, new MongoServerError(response)); + } + + if (error instanceof MongoNetworkError) { + const ctor = error.constructor as new ( + message: string, + options?: MongoNetworkErrorOptions + ) => MongoNetworkError; + return finishErrorCopy( + error, + new ctor(error.message, { + beforeHandshake: error.beforeHandshake, + cause: error.cause + }) + ); + } + + return finishErrorCopy( + error, + new MongoError(error.message, error.cause == null ? undefined : { cause: error.cause }) + ); +} + +function finishErrorCopy(source: MongoError, copy: MongoError): MongoError { + copy.code = source.code; + copy.connectionGeneration = source.connectionGeneration; + if (source.topologyVersion != null) { + copy.topologyVersion = copyTopologyVersion(source.topologyVersion); + } + for (const label of source.errorLabels) { + copy.addErrorLabel(label); + } + return copy; +} + +/** + * Keeps only servers from `eligible`, in the order the user returned them. + * A returned address that was not eligible is a programming error: accepting + * it would let the function bypass read preference, staleness, and deprioritization. + */ +function serversReturnedByUser( + eligible: ServerDescription[], + selected: ServerDescription[] +): ServerDescription[] { + const eligibleByAddress = new Map(eligible.map(server => [server.address, server])); + const narrowed: ServerDescription[] = []; + const seen = new Set(); + for (const server of selected) { + const match = eligibleByAddress.get(server.address); + if (match == null) { + throw new MongoInvalidArgumentError( + `serverSelector returned ${server.address}, which is not eligible for this operation` + ); + } + if (seen.has(server.address)) { + continue; + } + seen.add(server.address); + narrowed.push(match); + } + return narrowed; +} + +/** + * Built-in selectors apply the latency window before returning. Copy the + * description with an unlimited window so a user filter can still see, and + * choose, a slower eligible member. The real window is applied to the filter's + * result. + */ +function withoutLatencyWindow(topologyDescription: TopologyDescription): TopologyDescription { + const widened = Object.create(Object.getPrototypeOf(topologyDescription)); + return Object.assign(widened, topologyDescription, { + localThresholdMS: Number.POSITIVE_INFINITY + }); +} + +function isServerDescriptionList(value: unknown): value is ServerDescription[] { + if (!Array.isArray(value)) { + return false; + } + return value.every( + server => + server != null && + typeof server.address === 'string' && + typeof server.roundTripTime === 'number' + ); +} + +/** + * Applies a MongoClient `serverSelector` to an internal selector. + * + * The user function is skipped when no server is eligible, when the topology + * does not use server selection, and when the operation is pinned to a server. + * @internal + */ +export function applyUserServerSelector( + selector: ServerSelector, + userServerSelector: ServerSelectorFunction | undefined +): ServerSelector { + if (userServerSelector == null || selectorsIgnoringUserServerSelector.has(selector)) { + return selector; + } + + return function userServerSelectorWrapper( + topologyDescription: TopologyDescription, + servers: ServerDescription[], + deprioritized: DeprioritizedServers + ): ServerDescription[] { + if (TOPOLOGIES_IGNORING_USER_SERVER_SELECTOR.has(topologyDescription.type)) { + return selector(topologyDescription, servers, deprioritized); + } + + const eligible = selector(withoutLatencyWindow(topologyDescription), servers, deprioritized); + if (eligible.length === 0) { + return []; + } + + const selected = userServerSelector(eligible.map(copyServerDescription)); + if (!isServerDescriptionList(selected)) { + throw new MongoInvalidArgumentError( + 'serverSelector must return an array of server descriptions' + ); + } + // An empty array is the caller's way to say none of the eligible servers + // will do. The wait queue retries until serverSelectionTimeoutMS. A + // non-empty list is matched back to the original descriptions; + // serversReturnedByUser throws if any address was not eligible, so the + // matched list is non-empty whenever we get past that call. + if (selected.length === 0) { + return []; + } + return latencyWindowReducer(topologyDescription, serversReturnedByUser(eligible, selected)); + }; +} + // filters function primaryFilter(server: ServerDescription): boolean { return server.type === ServerType.RSPrimary; diff --git a/src/sdam/topology.ts b/src/sdam/topology.ts index a27ea95e2d..096e18e4a7 100644 --- a/src/sdam/topology.ts +++ b/src/sdam/topology.ts @@ -71,9 +71,12 @@ import type { ServerMonitoringMode } from './monitor'; import { Server, type ServerEvents, type ServerOptions } from './server'; import { compareTopologyVersion, ServerDescription } from './server_description'; import { + applyUserServerSelector, DeprioritizedServers, + operationSkipsUserServerSelector, readPreferenceServerSelector, - type ServerSelector + type ServerSelector, + type ServerSelectorFunction } from './server_selection'; import { ServerSelectionFailedEvent, @@ -162,6 +165,8 @@ export interface TopologyOptions extends BSONSerializeOptions, ServerOptions { serverMonitoringMode: ServerMonitoringMode; /** MongoDB server API version */ serverApi?: ServerApi; + /** @see MongoClientOptions.serverSelector */ + serverSelector?: ServerSelectorFunction; __skipPingOnConnect?: boolean; } @@ -560,6 +565,10 @@ export class Topology extends TypedEventEmitter { serverSelector = selector; } + if (!operationSkipsUserServerSelector(options.operationName)) { + serverSelector = applyUserServerSelector(serverSelector, this.s.options.serverSelector); + } + options = { serverSelectionTimeoutMS: this.s.serverSelectionTimeoutMS, ...options }; if ( this.client.mongoLogger?.willLog(MongoLoggableComponent.SERVER_SELECTION, SeverityLevel.DEBUG) diff --git a/test/integration/server-selection/server_selector.test.ts b/test/integration/server-selection/server_selector.test.ts new file mode 100644 index 0000000000..fe897ab548 --- /dev/null +++ b/test/integration/server-selection/server_selector.test.ts @@ -0,0 +1,97 @@ +import { expect } from 'chai'; +import { setTimeout } from 'timers/promises'; + +import { + type CommandStartedEvent, + HostAddress, + type MongoClient, + ReadPreference, + ServerType +} from '../../mongodb'; + +const METADATA: MongoDBMetadataUI = { requires: { topology: 'replicaset' } }; + +describe('serverSelector option', function () { + let client: MongoClient | undefined; + + afterEach(async function () { + await client?.close(); + client = undefined; + }); + + it('sends nearest reads to the member the selector returns', METADATA, async function () { + client = this.configuration.newClient( + {}, + { + monitorCommands: true, + serverSelector(servers) { + return [highestPort(servers)]; + } + } + ); + const activeClient = client; + await activeClient.connect(); + await waitForMembers(activeClient); + + const finds: CommandStartedEvent[] = []; + activeClient.on('commandStarted', event => { + if (event.commandName === 'find') finds.push(event); + }); + + const collection = activeClient.db('serverSelectorTest').collection(`members_${Date.now()}`); + await collection.insertMany([{ name: 'Ada' }, { name: 'Grace' }]); + for (let attempt = 0; attempt < 5; attempt++) { + await collection.findOne({ name: 'Ada' }, { readPreference: ReadPreference.nearest }); + } + + const topology = activeClient.topology; + if (topology == null) { + throw new Error('client is not connected'); + } + const expected = highestPort( + Array.from(topology.description.servers.values()).filter( + server => server.type === ServerType.RSPrimary || server.type === ServerType.RSSecondary + ) + ); + expect(finds).to.not.be.empty; + for (const event of finds) { + expect(event.address).to.equal(expected.address); + } + await collection.drop().catch(() => null); + }); +}); + +function highestPort(servers: T[]): T { + let selected = servers[0]; + for (const server of servers) { + if (portOf(server.address) > portOf(selected.address)) selected = server; + } + return selected; +} + +function portOf(address: string): number { + return HostAddress.fromString(address).port ?? 0; +} + +async function waitForMembers(client: MongoClient): Promise { + const hello = await client.db('admin').command({ hello: 1 }); + const hosts = new Set( + ((hello.hosts as string[] | undefined) ?? []).map(host => host.toLowerCase()) + ); + if (hosts.size === 0) { + throw new Error('hello did not report replica set hosts'); + } + + const deadline = Date.now() + 10_000; + while (Date.now() < deadline) { + const known = new Set( + Array.from(client.topology?.description.servers.values() ?? []) + .filter(server => server.type !== ServerType.Unknown) + .map(server => server.address.toLowerCase()) + ); + const discovered = Array.from(hosts).every(host => known.has(host)); + if (discovered) return; + await setTimeout(100); + } + throw new Error('timed out waiting for replica set members'); +} diff --git a/test/mongodb_bundled.ts b/test/mongodb_bundled.ts index de2cde27d6..da51ff6583 100644 --- a/test/mongodb_bundled.ts +++ b/test/mongodb_bundled.ts @@ -19,6 +19,7 @@ export const { AggregateOperation, AggregationCursor, applySession, + applyUserServerSelector, Aspect, AUTH_MECHS_AUTH_SRC_EXTERNAL, AuthContext, diff --git a/test/unit/sdam/server_selector.test.ts b/test/unit/sdam/server_selector.test.ts new file mode 100644 index 0000000000..32eda8eee3 --- /dev/null +++ b/test/unit/sdam/server_selector.test.ts @@ -0,0 +1,632 @@ +import { expect } from 'chai'; + +import { + applyUserServerSelector, + Binary, + DeprioritizedServers, + Long, + MIN_SECONDARY_WRITE_WIRE_VERSION, + MongoClient, + MongoInvalidArgumentError, + MongoNetworkError, + MongoParseError, + MongoServerError, + MongoServerSelectionError, + ObjectId, + parseOptions, + ReadPreference, + readPreferenceServerSelector, + sameServerSelector, + secondaryWritableServerSelector, + ServerDescription, + type ServerSelectorFunction, + Timestamp, + TopologyDescription, + TopologyType, + writableServerSelector +} from '../../mongodb'; +import { topologyWithPlaceholderClient } from '../../tools/utils'; + +const LOCAL_THRESHOLD_MS = 15; + +function member( + address: string, + type: 'primary' | 'secondary' | 'mongos' | 'standalone' | 'loadBalancer', + roundTripTime: number +): ServerDescription { + if (type === 'mongos') { + return new ServerDescription(address, { msg: 'isdbgrid', ok: 1 }, { roundTripTime }); + } + if (type === 'loadBalancer') { + return new ServerDescription(address, { ok: 1 }, { loadBalanced: true, roundTripTime }); + } + if (type === 'secondary') { + return new ServerDescription( + address, + { setName: 'rs', secondary: true, ok: 1 }, + { roundTripTime } + ); + } + return new ServerDescription( + address, + { setName: type === 'primary' ? 'rs' : undefined, isWritablePrimary: true, ok: 1 }, + { roundTripTime } + ); +} + +function description( + type: TopologyType, + servers: ServerDescription[], + localThresholdMS = LOCAL_THRESHOLD_MS +): TopologyDescription { + return new TopologyDescription( + type, + new Map(servers.map(server => [server.address, server])), + type === TopologyType.Sharded || type === TopologyType.Single ? null : 'rs', + null, + null, + MIN_SECONDARY_WRITE_WIRE_VERSION, + { localThresholdMS } + ); +} + +function addresses(servers: ServerDescription[]): string[] { + return servers.map(server => server.address); +} + +describe('serverSelector', function () { + describe('option parsing', function () { + const serverSelector: ServerSelectorFunction = servers => servers; + + it('accepts a function passed to MongoClient', function () { + const options = parseOptions('mongodb://localhost:27017', { serverSelector }); + expect(options.serverSelector).to.equal(serverSelector); + + const client = new MongoClient('mongodb://localhost:27017', { serverSelector }); + expect(client.options.serverSelector).to.equal(serverSelector); + }); + + it('leaves the option unset when it is omitted or null', function () { + expect(parseOptions('mongodb://localhost:27017').serverSelector).to.equal(undefined); + expect( + parseOptions('mongodb://localhost:27017', { serverSelector: null as never }).serverSelector + ).to.equal(undefined); + }); + + it('rejects a non-function', function () { + expect(() => + parseOptions('mongodb://localhost:27017', { serverSelector: 'host' as never }) + ).to.throw(MongoParseError, /serverSelector must be a function/); + }); + + it('rejects the option in a connection string', function () { + expect(() => parseOptions('mongodb://localhost:27017/?serverSelector=host')).to.throw( + MongoParseError, + /URI cannot contain `serverSelector`, it can only be passed to the client/ + ); + }); + }); + + describe('selection', function () { + const primary = member('127.0.0.1:27017', 'primary', 10); + const secondary = member('127.0.0.1:27018', 'secondary', 20); + const slowSecondary = member('127.0.0.1:27019', 'secondary', 100); + + it('sees every eligible server before the latency window', function () { + let seen: string[] = []; + const selector = applyUserServerSelector(readPreferenceNearest(), servers => { + seen = addresses(servers); + return servers; + }); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary, slowSecondary]), + [primary, secondary, slowSecondary], + new DeprioritizedServers() + ); + + expect(seen).to.deep.equal([primary.address, secondary.address, slowSecondary.address]); + expect(addresses(selected)).to.deep.equal([primary.address, secondary.address]); + }); + + it('keeps a single slower member outside the latency window', function () { + const selector = applyUserServerSelector(readPreferenceNearest(), servers => + servers.filter(server => server.address === slowSecondary.address) + ); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary, slowSecondary]), + [primary, secondary, slowSecondary], + new DeprioritizedServers() + ); + + expect(addresses(selected)).to.deep.equal([slowSecondary.address]); + }); + + it('passes only the primary to the filter for a primary read', function () { + let seen: string[] = []; + const selector = applyUserServerSelector(readPreference(ReadPreference.primary), servers => { + seen = addresses(servers); + return servers; + }); + + selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, slowSecondary]), + [primary, slowSecondary], + new DeprioritizedServers() + ); + + expect(seen).to.deep.equal([primary.address]); + }); + + it('does not call the filter when no server is eligible', function () { + let calls = 0; + const selector = applyUserServerSelector(readPreference(ReadPreference.primary), servers => { + calls += 1; + return servers; + }); + + const selected = selector( + description(TopologyType.ReplicaSetNoPrimary, [secondary, slowSecondary]), + [secondary, slowSecondary], + new DeprioritizedServers() + ); + + expect(calls).to.equal(0); + expect(selected).to.deep.equal([]); + }); + + it('can narrow secondary writes to one secondary', function () { + const selector = applyUserServerSelector( + secondaryWritableServerSelector(MIN_SECONDARY_WRITE_WIRE_VERSION, ReadPreference.secondary), + servers => servers.filter(server => server.address === slowSecondary.address) + ); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary, slowSecondary]), + [primary, secondary, slowSecondary], + new DeprioritizedServers() + ); + + expect(addresses(selected)).to.deep.equal([slowSecondary.address]); + }); + + it('ignores the filter for a single topology', function () { + let calls = 0; + const standalone = member('127.0.0.1:27021', 'standalone', 10); + const slow = member('127.0.0.1:27022', 'standalone', 100); + const selector = applyUserServerSelector(readPreferenceNearest(), () => { + calls += 1; + return []; + }); + + const selected = selector( + description(TopologyType.Single, [standalone, slow]), + [standalone, slow], + new DeprioritizedServers() + ); + + expect(calls).to.equal(0); + expect(addresses(selected)).to.deep.equal([standalone.address]); + }); + + it('ignores the filter for a load-balanced topology', function () { + let calls = 0; + const balancer = member('127.0.0.1:27020', 'loadBalancer', 10); + const selector = applyUserServerSelector(readPreferenceNearest(), () => { + calls += 1; + return []; + }); + + const selected = selector( + description(TopologyType.LoadBalanced, [balancer]), + [balancer], + new DeprioritizedServers() + ); + + expect(calls).to.equal(0); + expect(selected).to.deep.equal([balancer]); + }); + + it('does not move an operation pinned to a server', function () { + let calls = 0; + const selector = applyUserServerSelector(sameServerSelector(slowSecondary), () => { + calls += 1; + return [primary]; + }); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, slowSecondary]), + [primary, slowSecondary], + new DeprioritizedServers() + ); + + expect(calls).to.equal(0); + expect(addresses(selected)).to.deep.equal([slowSecondary.address]); + }); + + it('rejects a return value that is not a list of server descriptions', function () { + const selector = applyUserServerSelector(readPreferenceNearest(), () => 'nope' as never); + expect(() => + selector( + description(TopologyType.ReplicaSetWithPrimary, [primary]), + [primary], + new DeprioritizedServers() + ) + ).to.throw(MongoInvalidArgumentError, /array of server descriptions/); + }); + + it('gives the filter copies and selects the eligible instances', function () { + let seen: ServerDescription[] = []; + const selector = applyUserServerSelector(readPreferenceNearest(), servers => { + seen = servers; + servers[0].roundTripTime = 999; + servers[0].tags.nodeType = 'analytics'; + servers[0].hosts.push('evil:27017'); + return servers; + }); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary, slowSecondary]), + [primary, secondary, slowSecondary], + new DeprioritizedServers() + ); + + expect(seen[0]).to.not.equal(primary); + expect(primary.roundTripTime).to.equal(10); + expect(primary.tags).to.deep.equal({}); + expect(primary.hosts).to.deep.equal([]); + expect(selected[0]).to.equal(primary); + expect(selected[1]).to.equal(secondary); + }); + + it('keeps nested SDAM fields isolated from the filter', function () { + const electionId = new ObjectId(); + const processId = new ObjectId(); + const counter = Long.fromNumber(4); + const clusterTime = new Timestamp({ t: 2, i: 3 }); + const keyId = Long.fromNumber(1); + const hash = new Binary(Uint8Array.from([1, 2, 3])); + // One topologyVersion object is stored on the description, on the server + // error, and inside that error's response document. Same for electionId + // and $clusterTime: the description keeps the hello's objects. + const topologyVersion = { processId, counter }; + const error = new MongoServerError({ + message: 'ping failed', + ok: 0, + topologyVersion + }); + const described = new ServerDescription( + '127.0.0.1:27017', + { + setName: 'rs', + isWritablePrimary: true, + ok: 1, + electionId, + $clusterTime: { clusterTime, signature: { hash, keyId } } + }, + { roundTripTime: 10, error } + ); + expect(described.electionId).to.equal(electionId); + expect(described.topologyVersion).to.equal(error.topologyVersion); + expect(error.errorResponse.topologyVersion).to.equal(topologyVersion); + expect(described.$clusterTime?.clusterTime).to.equal(clusterTime); + expect(described.$clusterTime?.signature?.hash).to.equal(hash); + + const selector = applyUserServerSelector(readPreference(ReadPreference.primary), servers => { + const seen = servers[0]; + if ( + seen.electionId == null || + seen.topologyVersion == null || + seen.$clusterTime?.signature == null || + !(seen.error instanceof MongoServerError) || + seen.error.topologyVersion == null || + seen.error.errorResponse.topologyVersion == null + ) { + throw new Error('expected nested SDAM fields on the copy'); + } + // Edit the copy in place. Replacing the property would hide a shared + // Long, ObjectId, Timestamp, or Binary. + const idBytes = Uint8Array.from(seen.electionId.id); + idBytes[0] ^= 0xff; + seen.electionId.id = idBytes; + seen.topologyVersion.processId = new ObjectId(); + seen.topologyVersion.counter.low = 1; + seen.$clusterTime.clusterTime.low = 1; + seen.$clusterTime.signature.hash.value()[0] = 9; + seen.$clusterTime.signature.keyId.low = 9; + seen.error.message = 'changed'; + seen.error.topologyVersion.counter.low = 7; + seen.error.errorResponse.message = 'changed response'; + seen.error.errorResponse.topologyVersion.counter.low = 8; + return servers; + }); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [described]), + [described], + new DeprioritizedServers() + ); + + expect(selected[0]).to.equal(described); + expect(described.electionId).to.equal(electionId); + expect(described.topologyVersion?.processId).to.equal(processId); + expect(described.topologyVersion?.counter.toNumber()).to.equal(4); + expect(described.$clusterTime?.clusterTime.t).to.equal(2); + expect(described.$clusterTime?.clusterTime.i).to.equal(3); + expect(Array.from(described.$clusterTime?.signature?.hash.value() ?? [])).to.deep.equal([ + 1, 2, 3 + ]); + expect(described.$clusterTime?.signature?.keyId.toNumber()).to.equal(1); + expect(described.error).to.be.instanceOf(MongoServerError); + if (!(described.error instanceof MongoServerError)) { + throw new Error('expected the original MongoServerError'); + } + expect(described.error.message).to.equal('ping failed'); + expect(described.error.topologyVersion?.counter.toNumber()).to.equal(4); + expect(described.error.errorResponse.message).to.equal('ping failed'); + expect(described.error.errorResponse.topologyVersion?.counter.toNumber()).to.equal(4); + }); + + it('keeps a network error class and isolates its topologyVersion', function () { + const processId = new ObjectId(); + const counter = Long.fromNumber(4); + const topologyVersion = { processId, counter }; + const error = new MongoNetworkError('connection closed', { beforeHandshake: false }); + error.topologyVersion = topologyVersion; + error.addErrorLabel('ResetPool'); + const described = new ServerDescription( + '127.0.0.1:27017', + { setName: 'rs', isWritablePrimary: true, ok: 1 }, + { roundTripTime: 10, error } + ); + + let seenError: MongoNetworkError | undefined; + const selector = applyUserServerSelector(readPreference(ReadPreference.primary), servers => { + const seen = servers[0].error; + if (!(seen instanceof MongoNetworkError) || seen.topologyVersion == null) { + throw new Error('expected a MongoNetworkError copy'); + } + seenError = seen; + seen.message = 'changed'; + seen.topologyVersion.counter.low = 1; + seen.topologyVersion.processId = new ObjectId(); + return servers; + }); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [described]), + [described], + new DeprioritizedServers() + ); + + expect(selected[0]).to.equal(described); + expect(seenError).to.be.instanceOf(MongoNetworkError); + expect(seenError?.beforeHandshake).to.equal(false); + expect(seenError?.hasErrorLabel('ResetPool')).to.equal(true); + expect(error.message).to.equal('connection closed'); + expect(error.topologyVersion?.processId).to.equal(processId); + expect(error.topologyVersion?.counter.toNumber()).to.equal(4); + expect(described.topologyVersion?.counter.toNumber()).to.equal(4); + }); + + it('rejects a server that was not eligible', function () { + const other = member('127.0.0.1:27999', 'secondary', 1); + const selector = applyUserServerSelector(readPreference(ReadPreference.primary), () => [ + other + ]); + + expect(() => + selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, other]), + [primary, other], + new DeprioritizedServers() + ) + ).to.throw(MongoInvalidArgumentError, /127\.0\.0\.1:27999, which is not eligible/); + }); + + it('rejects a secondary returned for a write', function () { + const selector = applyUserServerSelector(writableServerSelector(), () => [secondary]); + + expect(() => + selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary]), + [primary, secondary], + new DeprioritizedServers() + ) + ).to.throw(MongoInvalidArgumentError, /not eligible/); + }); + + it('selects the eligible server when the filter returns another object with its address', function () { + const selector = applyUserServerSelector(readPreferenceNearest(), () => [ + { address: slowSecondary.address, roundTripTime: 0 } as ServerDescription + ]); + + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary, slowSecondary]), + [primary, secondary, slowSecondary], + new DeprioritizedServers() + ); + + expect(selected).to.have.lengthOf(1); + expect(selected[0]).to.equal(slowSecondary); + }); + + it('returns an empty list when the filter returns an empty list', function () { + const selector = applyUserServerSelector(readPreferenceNearest(), () => []); + const selected = selector( + description(TopologyType.ReplicaSetWithPrimary, [primary, secondary]), + [primary, secondary], + new DeprioritizedServers() + ); + expect(selected).to.deep.equal([]); + }); + + it('applies the latency window to mongos after the filter', function () { + const fast = member('127.0.0.1:27017', 'mongos', 10); + const slow = member('127.0.0.1:27018', 'mongos', 100); + let seen: string[] = []; + const selector = applyUserServerSelector(readPreference(ReadPreference.primary), servers => { + seen = addresses(servers); + return [slow]; + }); + + const selected = selector( + description(TopologyType.Sharded, [fast, slow]), + [fast, slow], + new DeprioritizedServers() + ); + + expect(seen).to.deep.equal([fast.address, slow.address]); + expect(addresses(selected)).to.deep.equal([slow.address]); + }); + }); + + describe('Topology#selectServer', function () { + const primary = member('127.0.0.1:27017', 'primary', 10); + const slowSecondary = member('127.0.0.1:27019', 'secondary', 100); + + it('uses the serverSelector parsed onto the client options', async function () { + const parsed = parseOptions('mongodb://127.0.0.1:27017,127.0.0.1:27019', { + serverSelectionTimeoutMS: 500, + localThresholdMS: LOCAL_THRESHOLD_MS, + serverSelector: servers => + servers.filter(server => server.address === slowSecondary.address) + }); + const topology = topologyWithPlaceholderClient( + ['127.0.0.1:27017', '127.0.0.1:27019'], + parsed + ); + install(topology, TopologyType.ReplicaSetWithPrimary, [primary, slowSecondary]); + + const selected = await topology.selectServer(ReadPreference.nearest, { + operationName: 'find', + deprioritizedServers: new DeprioritizedServers() + }); + + expect(selected).to.equal(topology.s.servers.get(slowSecondary.address)); + }); + + it('keeps getMore on its pinned server when serverSelector would reject it', async function () { + const topology = topologyWithPlaceholderClient(['127.0.0.1:27017', '127.0.0.1:27019'], { + serverSelectionTimeoutMS: 500, + serverSelector() { + throw new Error('serverSelector must not run for a pinned operation'); + } + }); + install(topology, TopologyType.ReplicaSetWithPrimary, [primary, slowSecondary]); + + const selected = await topology.selectServer(sameServerSelector(slowSecondary), { + operationName: 'getMore', + deprioritizedServers: new DeprioritizedServers() + }); + + expect(selected).to.equal(topology.s.servers.get(slowSecondary.address)); + }); + + it('fails selection when the filter throws', async function () { + const topology = topologyWithPlaceholderClient('127.0.0.1:27017', { + serverSelectionTimeoutMS: 500, + serverSelector() { + throw new Error('selector failed'); + } + }); + install(topology, TopologyType.ReplicaSetWithPrimary, [primary]); + + const error = await topology + .selectServer(ReadPreference.primary, { + operationName: 'find', + deprioritizedServers: new DeprioritizedServers() + }) + .catch(caught => caught); + + expect(error).to.be.instanceOf(Error); + expect(error.message).to.equal('selector failed'); + }); + + it('does not call the filter for the connect handshake', async function () { + let calls = 0; + const topology = topologyWithPlaceholderClient('127.0.0.1:27017', { + serverSelectionTimeoutMS: 500, + serverSelector() { + calls += 1; + return []; + } + }); + install(topology, TopologyType.ReplicaSetWithPrimary, [primary]); + + const selected = await topology.selectServer(ReadPreference.primary, { + operationName: 'handshake', + deprioritizedServers: new DeprioritizedServers() + }); + + expect(calls).to.equal(0); + expect(selected).to.equal(topology.s.servers.get(primary.address)); + }); + + it('does not call the filter for endSessions', async function () { + let calls = 0; + const topology = topologyWithPlaceholderClient('127.0.0.1:27017', { + serverSelectionTimeoutMS: 500, + serverSelector() { + calls += 1; + throw new Error('serverSelector must not run while ending sessions'); + } + }); + install(topology, TopologyType.ReplicaSetWithPrimary, [primary]); + + const selected = await topology.selectServer(ReadPreference.primaryPreferred, { + operationName: 'endSessions', + deprioritizedServers: new DeprioritizedServers() + }); + + expect(calls).to.equal(0); + expect(selected).to.equal(topology.s.servers.get(primary.address)); + }); + + it('times out when the filter returns no servers', async function () { + const topology = topologyWithPlaceholderClient('127.0.0.1:27017', { + serverSelectionTimeoutMS: 50, + serverSelector() { + return []; + } + }); + install(topology, TopologyType.ReplicaSetWithPrimary, [primary]); + + const error = await topology + .selectServer(ReadPreference.primary, { + operationName: 'find', + deprioritizedServers: new DeprioritizedServers() + }) + .catch(caught => caught); + + expect(error).to.be.instanceOf(MongoServerSelectionError); + }); + }); +}); + +function readPreference(preference: ReadPreference) { + return readPreferenceServerSelector(preference); +} + +function readPreferenceNearest() { + return readPreference(ReadPreference.nearest); +} + +function install( + topology: ReturnType, + type: TopologyType, + servers: ServerDescription[] +) { + topology.s.state = 'connected'; + topology.s.description = description(type, servers); + topology.s.servers.clear(); + for (const server of servers) { + topology.s.servers.set(server.address, { + address: server.address, + requestCheck() { + return Promise.resolve(); + } + } as never); + } +}