opusdev/vector-similarity-api
1
1import type { Document } from '../bson';2import { type AutoEncrypter } from '../client-side-encryption/auto_encrypter';3import { type CommandOptions, Connection } from '../cmap/connection';4import {5 ConnectionPool,6 type ConnectionPoolEvents,7 type ConnectionPoolOptions8} from '../cmap/connection_pool';9import { PoolClearedError } from '../cmap/errors';10import {11 APM_EVENTS,12 CLOSED,13 CMAP_EVENTS,14 CONNECT,15 DESCRIPTION_RECEIVED,16 ENDED,17 HEARTBEAT_EVENTS,18 SERVER_HEARTBEAT_FAILED,19 SERVER_HEARTBEAT_STARTED,20 SERVER_HEARTBEAT_SUCCEEDED21} from '../constants';22import {23 type AnyError,24 isNodeShuttingDownError,25 isSDAMUnrecoverableError,26 MONGODB_ERROR_CODES,27 MongoError,28 MongoErrorLabel,29 MongoNetworkError,30 MongoNetworkTimeoutError,31 MongoRuntimeError,32 MongoServerClosedError,33 type MongoServerError,34 needsRetryableWriteLabel35} from '../error';36import type { ServerApi } from '../mongo_client';37import { type Abortable, TypedEventEmitter } from '../mongo_types';38import { AggregateOperation } from '../operations/aggregate';39import type { GetMoreOptions } from '../operations/get_more';40import { type AbstractOperation } from '../operations/operation';41import type { ClientSession } from '../sessions';42import { type TimeoutContext } from '../timeout';43import { isTransactionCommand } from '../transactions';44import {45 abortable,46 type EventEmitterWithState,47 makeStateMachine,48 maxWireVersion,49 noop,50 squashError,51 supportsRetryableWrites52} from '../utils';53import { throwIfWriteConcernError } from '../write_concern';54import {55 type ClusterTime,56 STATE_CLOSED,57 STATE_CLOSING,58 STATE_CONNECTED,59 STATE_CONNECTING,60 TopologyType61} from './common';62import type {63 ServerHeartbeatFailedEvent,64 ServerHeartbeatStartedEvent,65 ServerHeartbeatSucceededEvent66} from './events';67import { Monitor, type MonitorOptions } from './monitor';68import { compareTopologyVersion, ServerDescription } from './server_description';69import { MIN_SECONDARY_WRITE_WIRE_VERSION } from './server_selection';70import type { Topology } from './topology';71 72const stateTransition = makeStateMachine({73 [STATE_CLOSED]: [STATE_CLOSED, STATE_CONNECTING],74 [STATE_CONNECTING]: [STATE_CONNECTING, STATE_CLOSING, STATE_CONNECTED, STATE_CLOSED],75 [STATE_CONNECTED]: [STATE_CONNECTED, STATE_CLOSING, STATE_CLOSED],76 [STATE_CLOSING]: [STATE_CLOSING, STATE_CLOSED]77});78 79/** @internal */80export type ServerOptions = Omit<ConnectionPoolOptions, 'id' | 'generation' | 'hostAddress'> &81 MonitorOptions;82 83/** @internal */84export interface ServerPrivate {85 /** The server description for this server */86 description: ServerDescription;87 /** A copy of the options used to construct this instance */88 options: ServerOptions;89 /** The current state of the Server */90 state: string;91 /** MongoDB server API version */92 serverApi?: ServerApi;93 /** A count of the operations currently running against the server. */94 operationCount: number;95}96 97/** @public */98export type ServerEvents = {99 serverHeartbeatStarted(event: ServerHeartbeatStartedEvent): void;100 serverHeartbeatSucceeded(event: ServerHeartbeatSucceededEvent): void;101 serverHeartbeatFailed(event: ServerHeartbeatFailedEvent): void;102 /** Top level MongoClient doesn't emit this so it is marked: @internal */103 connect(server: Server): void;104 descriptionReceived(description: ServerDescription): void;105 closed(): void;106 ended(): void;107} & ConnectionPoolEvents &108 EventEmitterWithState;109 110/** @internal */111export type ServerCommandOptions = Omit<CommandOptions, 'timeoutContext' | 'socketTimeoutMS'> & {112 timeoutContext: TimeoutContext;113 returnFieldSelector?: Document | null;114} & Abortable;115 116/** @internal */117export class Server extends TypedEventEmitter<ServerEvents> {118 /** @internal */119 s: ServerPrivate;120 /** @internal */121 topology: Topology;122 /** @internal */123 pool: ConnectionPool;124 serverApi?: ServerApi;125 hello?: Document;126 monitor: Monitor | null;127 128 /** @event */129 static readonly SERVER_HEARTBEAT_STARTED = SERVER_HEARTBEAT_STARTED;130 /** @event */131 static readonly SERVER_HEARTBEAT_SUCCEEDED = SERVER_HEARTBEAT_SUCCEEDED;132 /** @event */133 static readonly SERVER_HEARTBEAT_FAILED = SERVER_HEARTBEAT_FAILED;134 /** @event */135 static readonly CONNECT = CONNECT;136 /** @event */137 static readonly DESCRIPTION_RECEIVED = DESCRIPTION_RECEIVED;138 /** @event */139 static readonly CLOSED = CLOSED;140 /** @event */141 static readonly ENDED = ENDED;142 143 /**144 * Create a server145 */146 constructor(topology: Topology, description: ServerDescription, options: ServerOptions) {147 super();148 this.on('error', noop);149 150 this.serverApi = options.serverApi;151 152 const poolOptions = { hostAddress: description.hostAddress, ...options };153 154 this.topology = topology;155 this.pool = new ConnectionPool(this, poolOptions);156 157 this.s = {158 description,159 options,160 state: STATE_CLOSED,161 operationCount: 0162 };163 164 for (const event of [...CMAP_EVENTS, ...APM_EVENTS]) {165 this.pool.on(event, (e: any) => this.emit(event, e));166 }167 168 this.pool.on(Connection.CLUSTER_TIME_RECEIVED, (clusterTime: ClusterTime) => {169 this.clusterTime = clusterTime;170 });171 172 if (this.loadBalanced) {173 this.monitor = null;174 // monitoring is disabled in load balancing mode175 return;176 }177 178 // create the monitor179 this.monitor = new Monitor(this, this.s.options);180 181 for (const event of HEARTBEAT_EVENTS) {182 this.monitor.on(event, (e: any) => this.emit(event, e));183 }184 185 this.monitor.on('resetServer', (error: MongoServerError) => markServerUnknown(this, error));186 this.monitor.on(Server.SERVER_HEARTBEAT_SUCCEEDED, (event: ServerHeartbeatSucceededEvent) => {187 this.emit(188 Server.DESCRIPTION_RECEIVED,189 new ServerDescription(this.description.hostAddress, event.reply, {190 roundTripTime: this.monitor?.roundTripTime,191 minRoundTripTime: this.monitor?.minRoundTripTime192 })193 );194 195 if (this.s.state === STATE_CONNECTING) {196 stateTransition(this, STATE_CONNECTED);197 this.emit(Server.CONNECT, this);198 }199 });200 }201 202 get clusterTime(): ClusterTime | undefined {203 return this.topology.clusterTime;204 }205 206 set clusterTime(clusterTime: ClusterTime | undefined) {207 this.topology.clusterTime = clusterTime;208 }209 210 get description(): ServerDescription {211 return this.s.description;212 }213 214 get name(): string {215 return this.s.description.address;216 }217 218 get autoEncrypter(): AutoEncrypter | undefined {219 if (this.s.options && this.s.options.autoEncrypter) {220 return this.s.options.autoEncrypter;221 }222 return;223 }224 225 get loadBalanced(): boolean {226 return this.topology.description.type === TopologyType.LoadBalanced;227 }228 229 /**230 * Initiate server connect231 */232 connect(): void {233 if (this.s.state !== STATE_CLOSED) {234 return;235 }236 237 stateTransition(this, STATE_CONNECTING);238 239 // If in load balancer mode we automatically set the server to240 // a load balancer. It never transitions out of this state and241 // has no monitor.242 if (!this.loadBalanced) {243 this.monitor?.connect();244 } else {245 stateTransition(this, STATE_CONNECTED);246 this.emit(Server.CONNECT, this);247 }248 }249 250 closeCheckedOutConnections() {251 return this.pool.closeCheckedOutConnections();252 }253 254 /** Destroy the server connection */255 close(): void {256 if (this.s.state === STATE_CLOSED) {257 return;258 }259 260 stateTransition(this, STATE_CLOSING);261 262 if (!this.loadBalanced) {263 this.monitor?.close();264 }265 266 this.pool.close();267 stateTransition(this, STATE_CLOSED);268 this.emit('closed');269 }270 271 /**272 * Immediately schedule monitoring of this server. If there already an attempt being made273 * this will be a no-op.274 */275 requestCheck(): void {276 if (!this.loadBalanced) {277 this.monitor?.requestCheck();278 }279 }280 281 public async command<TResult>(282 operation: AbstractOperation<TResult>,283 timeoutContext: TimeoutContext284 ): Promise<InstanceType<typeof operation.SERVER_COMMAND_RESPONSE_TYPE>> {285 if (this.s.state === STATE_CLOSING || this.s.state === STATE_CLOSED) {286 throw new MongoServerClosedError();287 }288 const session = operation.session;289 290 let conn = session?.pinnedConnection;291 292 this.incrementOperationCount();293 if (conn == null) {294 try {295 conn = await this.pool.checkOut({ timeoutContext, signal: operation.options.signal });296 } catch (checkoutError) {297 this.decrementOperationCount();298 if (!(checkoutError instanceof PoolClearedError)) this.handleError(checkoutError);299 throw checkoutError;300 }301 }302 303 let reauthPromise: Promise<void> | null = null;304 const cleanup = () => {305 this.decrementOperationCount();306 if (session?.pinnedConnection !== conn) {307 if (reauthPromise != null) {308 // The reauth promise only exists if it hasn't thrown.309 const checkBackIn = () => {310 this.pool.checkIn(conn);311 };312 void reauthPromise.then(checkBackIn, checkBackIn);313 } else {314 this.pool.checkIn(conn);315 }316 }317 };318 319 let cmd;320 try {321 cmd = operation.buildCommand(conn, session);322 } catch (e) {323 cleanup();324 throw e;325 }326 327 const options = operation.buildOptions(timeoutContext);328 const ns = operation.ns;329 330 if (this.loadBalanced && isPinnableCommand(cmd, session) && !session?.pinnedConnection) {331 session?.pin(conn);332 }333 334 options.directConnection = this.topology.s.options.directConnection;335 336 const omitReadPreference =337 operation instanceof AggregateOperation &&338 operation.hasWriteStage &&339 maxWireVersion(conn) < MIN_SECONDARY_WRITE_WIRE_VERSION;340 if (omitReadPreference) {341 delete options.readPreference;342 }343 344 if (this.description.iscryptd) {345 options.omitMaxTimeMS = true;346 }347 348 try {349 try {350 const res = await conn.command(ns, cmd, options, operation.SERVER_COMMAND_RESPONSE_TYPE);351 throwIfWriteConcernError(res);352 return res;353 } catch (commandError) {354 throw this.decorateCommandError(conn, cmd, options, commandError);355 }356 } catch (operationError) {357 if (358 operationError instanceof MongoError &&359 operationError.code === MONGODB_ERROR_CODES.Reauthenticate360 ) {361 reauthPromise = this.pool.reauthenticate(conn);362 reauthPromise.then(undefined, error => {363 reauthPromise = null;364 squashError(error);365 });366 367 await abortable(reauthPromise, options);368 reauthPromise = null; // only reachable if reauth succeeds369 370 try {371 const res = await conn.command(ns, cmd, options, operation.SERVER_COMMAND_RESPONSE_TYPE);372 throwIfWriteConcernError(res);373 return res;374 } catch (commandError) {375 throw this.decorateCommandError(conn, cmd, options, commandError);376 }377 } else {378 throw operationError;379 }380 } finally {381 cleanup();382 }383 }384 385 /**386 * Handle SDAM error387 * @internal388 */389 handleError(error: AnyError, connection?: Connection) {390 if (!(error instanceof MongoError)) {391 return;392 }393 394 const isStaleError =395 error.connectionGeneration && error.connectionGeneration < this.pool.generation;396 if (isStaleError) {397 return;398 }399 400 const isNetworkNonTimeoutError =401 error instanceof MongoNetworkError && !(error instanceof MongoNetworkTimeoutError);402 const isNetworkTimeoutBeforeHandshakeError =403 error instanceof MongoNetworkError && error.beforeHandshake;404 const isAuthHandshakeError = error.hasErrorLabel(MongoErrorLabel.HandshakeError);405 if (isNetworkNonTimeoutError || isNetworkTimeoutBeforeHandshakeError || isAuthHandshakeError) {406 // In load balanced mode we never mark the server as unknown and always407 // clear for the specific service id.408 if (!this.loadBalanced) {409 error.addErrorLabel(MongoErrorLabel.ResetPool);410 markServerUnknown(this, error);411 } else if (connection) {412 this.pool.clear({ serviceId: connection.serviceId });413 }414 } else {415 if (isSDAMUnrecoverableError(error)) {416 if (shouldHandleStateChangeError(this, error)) {417 const shouldClearPool = isNodeShuttingDownError(error);418 if (this.loadBalanced && connection && shouldClearPool) {419 this.pool.clear({ serviceId: connection.serviceId });420 }421 422 if (!this.loadBalanced) {423 if (shouldClearPool) {424 error.addErrorLabel(MongoErrorLabel.ResetPool);425 }426 markServerUnknown(this, error);427 process.nextTick(() => this.requestCheck());428 }429 }430 }431 }432 }433 434 /**435 * Ensure that error is properly decorated and internal state is updated before throwing436 * @internal437 */438 private decorateCommandError(439 connection: Connection,440 cmd: Document,441 options: CommandOptions | GetMoreOptions | undefined,442 error: unknown443 ): Error {444 if (typeof error !== 'object' || error == null || !('name' in error)) {445 throw new MongoRuntimeError('An unexpected error type: ' + typeof error);446 }447 448 if (error.name === 'AbortError' && 'cause' in error && error.cause instanceof MongoError) {449 error = error.cause;450 }451 452 if (!(error instanceof MongoError)) {453 // Node.js or some other error we have not special handling for454 return error as Error;455 }456 457 if (connectionIsStale(this.pool, connection)) {458 return error;459 }460 461 const session = options?.session;462 if (error instanceof MongoNetworkError) {463 if (session && !session.hasEnded && session.serverSession) {464 session.serverSession.isDirty = true;465 }466 467 // inActiveTransaction check handles commit and abort.468 if (469 inActiveTransaction(session, cmd) &&470 !error.hasErrorLabel(MongoErrorLabel.TransientTransactionError)471 ) {472 error.addErrorLabel(MongoErrorLabel.TransientTransactionError);473 }474 475 if (476 (isRetryableWritesEnabled(this.topology) || isTransactionCommand(cmd)) &&477 supportsRetryableWrites(this) &&478 !inActiveTransaction(session, cmd)479 ) {480 error.addErrorLabel(MongoErrorLabel.RetryableWriteError);481 }482 } else {483 if (484 (isRetryableWritesEnabled(this.topology) || isTransactionCommand(cmd)) &&485 needsRetryableWriteLabel(error, maxWireVersion(this), this.description.type) &&486 !inActiveTransaction(session, cmd)487 ) {488 error.addErrorLabel(MongoErrorLabel.RetryableWriteError);489 }490 }491 492 if (493 session &&494 session.isPinned &&495 error.hasErrorLabel(MongoErrorLabel.TransientTransactionError)496 ) {497 session.unpin({ force: true });498 }499 500 this.handleError(error, connection);501 502 return error;503 }504 505 /**506 * Decrement the operation count, returning the new count.507 */508 private decrementOperationCount(): number {509 return (this.s.operationCount -= 1);510 }511 512 /**513 * Increment the operation count, returning the new count.514 */515 private incrementOperationCount(): number {516 return (this.s.operationCount += 1);517 }518}519 520function markServerUnknown(server: Server, error?: MongoError) {521 // Load balancer servers can never be marked unknown.522 if (server.loadBalanced) {523 return;524 }525 526 if (error instanceof MongoNetworkError && !(error instanceof MongoNetworkTimeoutError)) {527 server.monitor?.reset();528 }529 530 server.emit(531 Server.DESCRIPTION_RECEIVED,532 new ServerDescription(server.description.hostAddress, undefined, { error })533 );534}535 536function isPinnableCommand(cmd: Document, session?: ClientSession): boolean {537 if (session) {538 return (539 session.inTransaction() ||540 (session.transaction.isCommitted && 'commitTransaction' in cmd) ||541 'aggregate' in cmd ||542 'find' in cmd ||543 'getMore' in cmd ||544 'listCollections' in cmd ||545 'listIndexes' in cmd ||546 'bulkWrite' in cmd547 );548 }549 550 return false;551}552 553function connectionIsStale(pool: ConnectionPool, connection: Connection) {554 if (connection.serviceId) {555 return (556 connection.generation !== pool.serviceGenerations.get(connection.serviceId.toHexString())557 );558 }559 560 return connection.generation !== pool.generation;561}562 563function shouldHandleStateChangeError(server: Server, err: MongoError) {564 const etv = err.topologyVersion;565 const stv = server.description.topologyVersion;566 return compareTopologyVersion(stv, etv) < 0;567}568 569function inActiveTransaction(session: ClientSession | undefined, cmd: Document) {570 return session && session.inTransaction() && !isTransactionCommand(cmd);571}572 573/** this checks the retryWrites option passed down from the client options, it574 * does not check if the server supports retryable writes */575function isRetryableWritesEnabled(topology: Topology) {576 return topology.s.options.retryWrites !== false;577}578