opusdev/vector-similarity-api
1
1import { clearTimeout, setTimeout } from 'timers';2 3import { type Document, Long } from '../bson';4import { connect, makeConnection, makeSocket, performInitialHandshake } from '../cmap/connect';5import type { Connection, ConnectionOptions } from '../cmap/connection';6import { getFAASEnv } from '../cmap/handshake/client_metadata';7import { LEGACY_HELLO_COMMAND } from '../constants';8import { MongoError, MongoErrorLabel, MongoNetworkTimeoutError } from '../error';9import { MongoLoggableComponent } from '../mongo_logger';10import { CancellationToken, TypedEventEmitter } from '../mongo_types';11import {12 calculateDurationInMs,13 type Callback,14 type EventEmitterWithState,15 makeStateMachine,16 noop,17 now,18 ns19} from '../utils';20import { ServerType, STATE_CLOSED, STATE_CLOSING } from './common';21import {22 ServerHeartbeatFailedEvent,23 ServerHeartbeatStartedEvent,24 ServerHeartbeatSucceededEvent25} from './events';26import { Server } from './server';27import type { TopologyVersion } from './server_description';28 29const STATE_IDLE = 'idle';30const STATE_MONITORING = 'monitoring';31const stateTransition = makeStateMachine({32 [STATE_CLOSING]: [STATE_CLOSING, STATE_IDLE, STATE_CLOSED],33 [STATE_CLOSED]: [STATE_CLOSED, STATE_MONITORING],34 [STATE_IDLE]: [STATE_IDLE, STATE_MONITORING, STATE_CLOSING],35 [STATE_MONITORING]: [STATE_MONITORING, STATE_IDLE, STATE_CLOSING]36});37 38const INVALID_REQUEST_CHECK_STATES = new Set([STATE_CLOSING, STATE_CLOSED, STATE_MONITORING]);39function isInCloseState(monitor: Monitor) {40 return monitor.s.state === STATE_CLOSED || monitor.s.state === STATE_CLOSING;41}42 43/** @public */44export const ServerMonitoringMode = Object.freeze({45 auto: 'auto',46 poll: 'poll',47 stream: 'stream'48} as const);49 50/** @public */51export type ServerMonitoringMode = (typeof ServerMonitoringMode)[keyof typeof ServerMonitoringMode];52 53/** @internal */54export interface MonitorPrivate {55 state: string;56}57 58/** @public */59export interface MonitorOptions60 extends Omit<ConnectionOptions, 'id' | 'generation' | 'hostAddress'> {61 connectTimeoutMS: number;62 heartbeatFrequencyMS: number;63 minHeartbeatFrequencyMS: number;64 serverMonitoringMode: ServerMonitoringMode;65}66 67/** @public */68export type MonitorEvents = {69 serverHeartbeatStarted(event: ServerHeartbeatStartedEvent): void;70 serverHeartbeatSucceeded(event: ServerHeartbeatSucceededEvent): void;71 serverHeartbeatFailed(event: ServerHeartbeatFailedEvent): void;72 resetServer(error?: MongoError): void;73 resetConnectionPool(): void;74 close(): void;75} & EventEmitterWithState;76 77/** @internal */78export class Monitor extends TypedEventEmitter<MonitorEvents> {79 /** @internal */80 s: MonitorPrivate;81 address: string;82 options: Readonly<83 Pick<84 MonitorOptions,85 | 'connectTimeoutMS'86 | 'heartbeatFrequencyMS'87 | 'minHeartbeatFrequencyMS'88 | 'serverMonitoringMode'89 >90 >;91 connectOptions: ConnectionOptions;92 isRunningInFaasEnv: boolean;93 server: Server;94 connection: Connection | null;95 cancellationToken: CancellationToken;96 /** @internal */97 monitorId?: MonitorInterval;98 rttPinger?: RTTPinger;99 /** @internal */100 override component = MongoLoggableComponent.TOPOLOGY;101 /** @internal */102 private rttSampler: RTTSampler;103 104 constructor(server: Server, options: MonitorOptions) {105 super();106 this.on('error', noop);107 108 this.server = server;109 this.connection = null;110 this.cancellationToken = new CancellationToken();111 this.cancellationToken.setMaxListeners(Infinity);112 this.monitorId = undefined;113 this.s = {114 state: STATE_CLOSED115 };116 this.address = server.description.address;117 this.options = Object.freeze({118 connectTimeoutMS: options.connectTimeoutMS ?? 10000,119 heartbeatFrequencyMS: options.heartbeatFrequencyMS ?? 10000,120 minHeartbeatFrequencyMS: options.minHeartbeatFrequencyMS ?? 500,121 serverMonitoringMode: options.serverMonitoringMode122 });123 this.isRunningInFaasEnv = getFAASEnv() != null;124 this.mongoLogger = this.server.topology.client?.mongoLogger;125 this.rttSampler = new RTTSampler(10);126 127 const cancellationToken = this.cancellationToken;128 // TODO: refactor this to pull it directly from the pool, requires new ConnectionPool integration129 const connectOptions = {130 id: '<monitor>' as const,131 generation: server.pool.generation,132 cancellationToken,133 hostAddress: server.description.hostAddress,134 ...options,135 // force BSON serialization options136 raw: false,137 useBigInt64: false,138 promoteLongs: true,139 promoteValues: true,140 promoteBuffers: true141 };142 143 // ensure no authentication is used for monitoring144 delete connectOptions.credentials;145 if (connectOptions.autoEncrypter) {146 delete connectOptions.autoEncrypter;147 }148 149 this.connectOptions = Object.freeze(connectOptions);150 }151 152 connect(): void {153 if (this.s.state !== STATE_CLOSED) {154 return;155 }156 157 // start158 const heartbeatFrequencyMS = this.options.heartbeatFrequencyMS;159 const minHeartbeatFrequencyMS = this.options.minHeartbeatFrequencyMS;160 this.monitorId = new MonitorInterval(monitorServer(this), {161 heartbeatFrequencyMS: heartbeatFrequencyMS,162 minHeartbeatFrequencyMS: minHeartbeatFrequencyMS,163 immediate: true164 });165 }166 167 requestCheck(): void {168 if (INVALID_REQUEST_CHECK_STATES.has(this.s.state)) {169 return;170 }171 172 this.monitorId?.wake();173 }174 175 reset(): void {176 const topologyVersion = this.server.description.topologyVersion;177 if (isInCloseState(this) || topologyVersion == null) {178 return;179 }180 181 stateTransition(this, STATE_CLOSING);182 resetMonitorState(this);183 184 // restart monitor185 stateTransition(this, STATE_IDLE);186 187 // restart monitoring188 const heartbeatFrequencyMS = this.options.heartbeatFrequencyMS;189 const minHeartbeatFrequencyMS = this.options.minHeartbeatFrequencyMS;190 this.monitorId = new MonitorInterval(monitorServer(this), {191 heartbeatFrequencyMS: heartbeatFrequencyMS,192 minHeartbeatFrequencyMS: minHeartbeatFrequencyMS193 });194 }195 196 close(): void {197 if (isInCloseState(this)) {198 return;199 }200 201 stateTransition(this, STATE_CLOSING);202 resetMonitorState(this);203 204 // close monitor205 this.emit('close');206 stateTransition(this, STATE_CLOSED);207 }208 209 get roundTripTime(): number {210 return this.rttSampler.average();211 }212 213 get minRoundTripTime(): number {214 return this.rttSampler.min();215 }216 217 get latestRtt(): number | null {218 return this.rttSampler.last;219 }220 221 addRttSample(rtt: number) {222 this.rttSampler.addSample(rtt);223 }224 225 clearRttSamples() {226 this.rttSampler.clear();227 }228}229 230function resetMonitorState(monitor: Monitor) {231 monitor.monitorId?.stop();232 monitor.monitorId = undefined;233 234 monitor.rttPinger?.close();235 monitor.rttPinger = undefined;236 237 monitor.cancellationToken.emit('cancel');238 239 monitor.connection?.destroy();240 monitor.connection = null;241 242 monitor.clearRttSamples();243}244 245function useStreamingProtocol(monitor: Monitor, topologyVersion: TopologyVersion | null): boolean {246 // If we have no topology version we always poll no matter247 // what the user provided, since the server does not support248 // the streaming protocol.249 if (topologyVersion == null) return false;250 251 const serverMonitoringMode = monitor.options.serverMonitoringMode;252 if (serverMonitoringMode === ServerMonitoringMode.poll) return false;253 if (serverMonitoringMode === ServerMonitoringMode.stream) return true;254 255 // If we are in auto mode, we need to figure out if we're in a FaaS256 // environment or not and choose the appropriate mode.257 if (monitor.isRunningInFaasEnv) return false;258 return true;259}260 261function checkServer(monitor: Monitor, callback: Callback<Document | null>) {262 let start: number;263 let awaited: boolean;264 const topologyVersion = monitor.server.description.topologyVersion;265 const isAwaitable = useStreamingProtocol(monitor, topologyVersion);266 monitor.emitAndLogHeartbeat(267 Server.SERVER_HEARTBEAT_STARTED,268 monitor.server.topology.s.id,269 undefined,270 new ServerHeartbeatStartedEvent(monitor.address, isAwaitable)271 );272 273 function onHeartbeatFailed(err: Error) {274 monitor.connection?.destroy();275 monitor.connection = null;276 monitor.emitAndLogHeartbeat(277 Server.SERVER_HEARTBEAT_FAILED,278 monitor.server.topology.s.id,279 undefined,280 new ServerHeartbeatFailedEvent(monitor.address, calculateDurationInMs(start), err, awaited)281 );282 283 const error = !(err instanceof MongoError)284 ? new MongoError(MongoError.buildErrorMessage(err), { cause: err })285 : err;286 error.addErrorLabel(MongoErrorLabel.ResetPool);287 if (error instanceof MongoNetworkTimeoutError) {288 error.addErrorLabel(MongoErrorLabel.InterruptInUseConnections);289 }290 291 monitor.emit('resetServer', error);292 callback(err);293 }294 295 function onHeartbeatSucceeded(hello: Document) {296 if (!('isWritablePrimary' in hello)) {297 // Provide hello-style response document.298 hello.isWritablePrimary = hello[LEGACY_HELLO_COMMAND];299 }300 301 // NOTE: here we use the latestRtt as this measurement corresponds with the value302 // obtained for this successful heartbeat, if there is no latestRtt, then we calculate the303 // duration304 const duration =305 isAwaitable && monitor.rttPinger306 ? (monitor.rttPinger.latestRtt ?? calculateDurationInMs(start))307 : calculateDurationInMs(start);308 309 monitor.addRttSample(duration);310 311 monitor.emitAndLogHeartbeat(312 Server.SERVER_HEARTBEAT_SUCCEEDED,313 monitor.server.topology.s.id,314 hello.connectionId,315 new ServerHeartbeatSucceededEvent(monitor.address, duration, hello, isAwaitable)316 );317 318 if (isAwaitable) {319 // If we are using the streaming protocol then we immediately issue another 'started'320 // event, otherwise the "check" is complete and return to the main monitor loop321 monitor.emitAndLogHeartbeat(322 Server.SERVER_HEARTBEAT_STARTED,323 monitor.server.topology.s.id,324 undefined,325 new ServerHeartbeatStartedEvent(monitor.address, true)326 );327 // We have not actually sent an outgoing handshake, but when we get the next response we328 // want the duration to reflect the time since we last heard from the server329 start = now();330 } else {331 monitor.rttPinger?.close();332 monitor.rttPinger = undefined;333 334 callback(undefined, hello);335 }336 }337 338 const { connection } = monitor;339 if (connection && !connection.closed) {340 const { serverApi, helloOk } = connection;341 const connectTimeoutMS = monitor.options.connectTimeoutMS;342 const maxAwaitTimeMS = monitor.options.heartbeatFrequencyMS;343 344 const cmd = {345 [serverApi?.version || helloOk ? 'hello' : LEGACY_HELLO_COMMAND]: 1,346 ...(isAwaitable && topologyVersion347 ? { maxAwaitTimeMS, topologyVersion: makeTopologyVersion(topologyVersion) }348 : {})349 };350 351 const options = isAwaitable352 ? {353 socketTimeoutMS: connectTimeoutMS ? connectTimeoutMS + maxAwaitTimeMS : 0,354 exhaustAllowed: true355 }356 : { socketTimeoutMS: connectTimeoutMS };357 358 if (isAwaitable && monitor.rttPinger == null) {359 monitor.rttPinger = new RTTPinger(monitor);360 }361 362 // Record new start time before sending handshake363 start = now();364 365 if (isAwaitable) {366 awaited = true;367 return connection.exhaustCommand(ns('admin.$cmd'), cmd, options, (error, hello) => {368 if (error) return onHeartbeatFailed(error);369 return onHeartbeatSucceeded(hello);370 });371 }372 373 awaited = false;374 connection375 .command(ns('admin.$cmd'), cmd, options)376 .then(onHeartbeatSucceeded, onHeartbeatFailed);377 378 return;379 }380 381 // connecting does an implicit `hello`382 (async () => {383 const socket = await makeSocket(monitor.connectOptions);384 const connection = makeConnection(monitor.connectOptions, socket);385 // The start time is after socket creation but before the handshake386 start = now();387 try {388 await performInitialHandshake(connection, monitor.connectOptions);389 return connection;390 } catch (error) {391 connection.destroy();392 throw error;393 }394 })().then(395 connection => {396 if (isInCloseState(monitor)) {397 connection.destroy();398 return;399 }400 const duration = calculateDurationInMs(start);401 monitor.addRttSample(duration);402 403 monitor.connection = connection;404 monitor.emitAndLogHeartbeat(405 Server.SERVER_HEARTBEAT_SUCCEEDED,406 monitor.server.topology.s.id,407 connection.hello?.connectionId,408 new ServerHeartbeatSucceededEvent(409 monitor.address,410 duration,411 connection.hello,412 useStreamingProtocol(monitor, connection.hello?.topologyVersion)413 )414 );415 416 callback(undefined, connection.hello);417 },418 error => {419 monitor.connection = null;420 awaited = false;421 onHeartbeatFailed(error);422 }423 );424}425 426function monitorServer(monitor: Monitor) {427 return (callback: Callback) => {428 if (monitor.s.state === STATE_MONITORING) {429 process.nextTick(callback);430 return;431 }432 stateTransition(monitor, STATE_MONITORING);433 function done() {434 if (!isInCloseState(monitor)) {435 stateTransition(monitor, STATE_IDLE);436 }437 438 callback();439 }440 441 checkServer(monitor, (err, hello) => {442 if (err) {443 // otherwise an error occurred on initial discovery, also bail444 if (monitor.server.description.type === ServerType.Unknown) {445 return done();446 }447 }448 449 // if the check indicates streaming is supported, immediately reschedule monitoring450 if (useStreamingProtocol(monitor, hello?.topologyVersion)) {451 setTimeout(() => {452 if (!isInCloseState(monitor)) {453 monitor.monitorId?.wake();454 }455 }, 0);456 }457 458 done();459 });460 };461}462 463function makeTopologyVersion(tv: TopologyVersion) {464 return {465 processId: tv.processId,466 // tests mock counter as just number, but in a real situation counter should always be a Long467 // TODO(NODE-2674): Preserve int64 sent from MongoDB468 counter: Long.isLong(tv.counter) ? tv.counter : Long.fromNumber(tv.counter)469 };470}471 472/** @internal */473export interface RTTPingerOptions extends ConnectionOptions {474 heartbeatFrequencyMS: number;475}476 477/** @internal */478export class RTTPinger {479 connection?: Connection;480 /** @internal */481 cancellationToken: CancellationToken;482 /** @internal */483 monitorId: NodeJS.Timeout;484 /** @internal */485 monitor: Monitor;486 closed: boolean;487 /** @internal */488 latestRtt?: number;489 490 constructor(monitor: Monitor) {491 this.connection = undefined;492 this.cancellationToken = monitor.cancellationToken;493 this.closed = false;494 this.monitor = monitor;495 this.latestRtt = monitor.latestRtt ?? undefined;496 497 const heartbeatFrequencyMS = monitor.options.heartbeatFrequencyMS;498 this.monitorId = setTimeout(() => this.measureRoundTripTime(), heartbeatFrequencyMS);499 }500 501 get roundTripTime(): number {502 return this.monitor.roundTripTime;503 }504 505 get minRoundTripTime(): number {506 return this.monitor.minRoundTripTime;507 }508 509 close(): void {510 this.closed = true;511 clearTimeout(this.monitorId);512 513 this.connection?.destroy();514 this.connection = undefined;515 }516 517 private measureAndReschedule(start: number, conn?: Connection) {518 if (this.closed) {519 conn?.destroy();520 return;521 }522 523 if (this.connection == null) {524 this.connection = conn;525 }526 527 this.latestRtt = calculateDurationInMs(start);528 this.monitorId = setTimeout(529 () => this.measureRoundTripTime(),530 this.monitor.options.heartbeatFrequencyMS531 );532 }533 534 private measureRoundTripTime() {535 const start = now();536 537 if (this.closed) {538 return;539 }540 541 const connection = this.connection;542 if (connection == null) {543 connect(this.monitor.connectOptions).then(544 connection => {545 this.measureAndReschedule(start, connection);546 },547 () => {548 this.connection = undefined;549 }550 );551 return;552 }553 554 const commandName =555 connection.serverApi?.version || connection.helloOk ? 'hello' : LEGACY_HELLO_COMMAND;556 557 connection.command(ns('admin.$cmd'), { [commandName]: 1 }, undefined).then(558 () => this.measureAndReschedule(start),559 () => {560 this.connection?.destroy();561 this.connection = undefined;562 return;563 }564 );565 }566}567 568/**569 * @internal570 */571export interface MonitorIntervalOptions {572 /** The interval to execute a method on */573 heartbeatFrequencyMS: number;574 /** A minimum interval that must elapse before the method is called */575 minHeartbeatFrequencyMS: number;576 /** Whether the method should be called immediately when the interval is started */577 immediate: boolean;578}579 580/**581 * @internal582 */583export class MonitorInterval {584 fn: (callback: Callback) => void;585 timerId: NodeJS.Timeout | undefined;586 lastExecutionEnded: number;587 isExpeditedCallToFnScheduled = false;588 stopped = false;589 isExecutionInProgress = false;590 hasExecutedOnce = false;591 592 heartbeatFrequencyMS: number;593 minHeartbeatFrequencyMS: number;594 595 constructor(fn: (callback: Callback) => void, options: Partial<MonitorIntervalOptions> = {}) {596 this.fn = fn;597 this.lastExecutionEnded = -Infinity;598 599 this.heartbeatFrequencyMS = options.heartbeatFrequencyMS ?? 1000;600 this.minHeartbeatFrequencyMS = options.minHeartbeatFrequencyMS ?? 500;601 602 if (options.immediate) {603 this._executeAndReschedule();604 } else {605 this._reschedule(undefined);606 }607 }608 609 wake() {610 const currentTime = now();611 const timeSinceLastCall = currentTime - this.lastExecutionEnded;612 613 // TODO(NODE-4674): Add error handling and logging to the monitor614 if (timeSinceLastCall < 0) {615 return this._executeAndReschedule();616 }617 618 if (this.isExecutionInProgress) {619 return;620 }621 622 // debounce multiple calls to wake within the `minInterval`623 if (this.isExpeditedCallToFnScheduled) {624 return;625 }626 627 // reschedule a call as soon as possible, ensuring the call never happens628 // faster than the `minInterval`629 if (timeSinceLastCall < this.minHeartbeatFrequencyMS) {630 this.isExpeditedCallToFnScheduled = true;631 this._reschedule(this.minHeartbeatFrequencyMS - timeSinceLastCall);632 return;633 }634 635 this._executeAndReschedule();636 }637 638 stop() {639 this.stopped = true;640 if (this.timerId) {641 clearTimeout(this.timerId);642 this.timerId = undefined;643 }644 645 this.lastExecutionEnded = -Infinity;646 this.isExpeditedCallToFnScheduled = false;647 }648 649 toString() {650 return JSON.stringify(this);651 }652 653 toJSON() {654 const currentTime = now();655 const timeSinceLastCall = currentTime - this.lastExecutionEnded;656 return {657 timerId: this.timerId != null ? 'set' : 'cleared',658 lastCallTime: this.lastExecutionEnded,659 isExpeditedCheckScheduled: this.isExpeditedCallToFnScheduled,660 stopped: this.stopped,661 heartbeatFrequencyMS: this.heartbeatFrequencyMS,662 minHeartbeatFrequencyMS: this.minHeartbeatFrequencyMS,663 currentTime,664 timeSinceLastCall665 };666 }667 668 private _reschedule(ms?: number) {669 if (this.stopped) return;670 if (this.timerId) {671 clearTimeout(this.timerId);672 }673 674 this.timerId = setTimeout(this._executeAndReschedule, ms || this.heartbeatFrequencyMS);675 }676 677 private _executeAndReschedule = () => {678 if (this.stopped) return;679 if (this.timerId) {680 clearTimeout(this.timerId);681 }682 683 this.isExpeditedCallToFnScheduled = false;684 this.isExecutionInProgress = true;685 686 this.fn(() => {687 this.lastExecutionEnded = now();688 this.isExecutionInProgress = false;689 this._reschedule(this.heartbeatFrequencyMS);690 });691 };692}693 694/** @internal695 * This class implements the RTT sampling logic specified for [CSOT](https://github.com/mongodb/specifications/blob/bbb335e60cd7ea1e0f7cd9a9443cb95fc9d3b64d/source/client-side-operations-timeout/client-side-operations-timeout.md#drivers-use-minimum-rtt-to-short-circuit-operations)696 *697 * This is implemented as a [circular buffer](https://en.wikipedia.org/wiki/Circular_buffer) keeping698 * the most recent `windowSize` samples699 * */700export class RTTSampler {701 /** Index of the next slot to be overwritten */702 private writeIndex: number;703 private length: number;704 private rttSamples: Float64Array;705 706 constructor(windowSize = 10) {707 this.rttSamples = new Float64Array(windowSize);708 this.length = 0;709 this.writeIndex = 0;710 }711 712 /**713 * Adds an rtt sample to the end of the circular buffer714 * When `windowSize` samples have been collected, `addSample` overwrites the least recently added715 * sample716 */717 addSample(sample: number) {718 this.rttSamples[this.writeIndex++] = sample;719 if (this.length < this.rttSamples.length) {720 this.length++;721 }722 723 this.writeIndex %= this.rttSamples.length;724 }725 726 /**727 * When \< 2 samples have been collected, returns 0728 * Otherwise computes the minimum value samples contained in the buffer729 */730 min(): number {731 if (this.length < 2) return 0;732 let min = this.rttSamples[0];733 for (let i = 1; i < this.length; i++) {734 if (this.rttSamples[i] < min) min = this.rttSamples[i];735 }736 737 return min;738 }739 740 /**741 * Returns mean of samples contained in the buffer742 */743 average(): number {744 if (this.length === 0) return 0;745 let sum = 0;746 for (let i = 0; i < this.length; i++) {747 sum += this.rttSamples[i];748 }749 750 return sum / this.length;751 }752 753 /**754 * Returns most recently inserted element in the buffer755 * Returns null if the buffer is empty756 * */757 get last(): number | null {758 if (this.length === 0) return null;759 return this.rttSamples[this.writeIndex === 0 ? this.length - 1 : this.writeIndex - 1];760 }761 762 /**763 * Clear the buffer764 * NOTE: this does not overwrite the data held in the internal array, just the pointers into765 * this array766 */767 clear() {768 this.length = 0;769 this.writeIndex = 0;770 }771}772