opusdev/vector-similarity-api

1import { clearTimeout, setTimeout } from 'timers';2 3import { type Document } from './bson';4import { MongoInvalidArgumentError, MongoOperationTimeoutError, MongoRuntimeError } from './error';5import { type ClientSession } from './sessions';6import { csotMin, noop, squashError } from './utils';7 8/** @internal */9export class TimeoutError extends Error {10  duration: number;11  override get name(): 'TimeoutError' {12    return 'TimeoutError';13  }14 15  constructor(message: string, options: { cause?: Error; duration: number }) {16    super(message, options);17    this.duration = options.duration;18  }19 20  static is(error: unknown): error is TimeoutError {21    return (22      error != null && typeof error === 'object' && 'name' in error && error.name === 'TimeoutError'23    );24  }25}26 27type Executor = ConstructorParameters<typeof Promise<never>>[0];28type Reject = Parameters<ConstructorParameters<typeof Promise<never>>[0]>[1];29/**30 * @internal31 * This class is an abstraction over timeouts32 * The Timeout class can only be in the pending or rejected states. It is guaranteed not to resolve33 * if interacted with exclusively through its public API34 * */35export class Timeout extends Promise<never> {36  private id?: NodeJS.Timeout;37 38  public readonly start: number;39  public ended: number | null = null;40  public duration: number;41  private timedOut = false;42  public cleared = false;43 44  get remainingTime(): number {45    if (this.timedOut) return 0;46    if (this.duration === 0) return Infinity;47    return this.start + this.duration - Math.trunc(performance.now());48  }49 50  get timeElapsed(): number {51    return Math.trunc(performance.now()) - this.start;52  }53 54  /** Create a new timeout that expires in `duration` ms */55  private constructor(56    executor: Executor = () => null,57    options?: { duration: number; unref?: true; rejection?: Error }58  ) {59    const duration = options?.duration ?? 0;60    const unref = !!options?.unref;61    const rejection = options?.rejection;62 63    if (duration < 0) {64      throw new MongoInvalidArgumentError('Cannot create a Timeout with a negative duration');65    }66 67    let reject!: Reject;68    super((_, promiseReject) => {69      reject = promiseReject;70 71      executor(noop, promiseReject);72    });73 74    this.duration = duration;75    this.start = Math.trunc(performance.now());76 77    if (rejection == null && this.duration > 0) {78      this.id = setTimeout(() => {79        this.ended = Math.trunc(performance.now());80        this.timedOut = true;81        reject(new TimeoutError(`Expired after ${duration}ms`, { duration }));82      }, this.duration);83      if (typeof this.id.unref === 'function' && unref) {84        // Ensure we do not keep the Node.js event loop running85        this.id.unref();86      }87    } else if (rejection != null) {88      this.ended = Math.trunc(performance.now());89      this.timedOut = true;90      reject(rejection);91    }92  }93 94  /**95   * Clears the underlying timeout. This method is idempotent96   */97  clear(): void {98    clearTimeout(this.id);99    this.id = undefined;100    this.timedOut = false;101    this.cleared = true;102  }103 104  throwIfExpired(): void {105    if (this.timedOut) {106      // This method is invoked when someone wants to throw immediately instead of await the result of this promise107      // Since they won't be handling the rejection from the promise (because we're about to throw here)108      // attach handling to prevent this from bubbling up to Node.js109      this.then(undefined, squashError);110      throw new TimeoutError('Timed out', { duration: this.duration });111    }112  }113 114  public static expires(duration: number, unref?: true): Timeout {115    return new Timeout(undefined, { duration, unref });116  }117 118  static override reject(rejection?: Error): Timeout {119    return new Timeout(undefined, { duration: 0, unref: true, rejection });120  }121}122 123/** @internal */124export type TimeoutContextOptions = (LegacyTimeoutContextOptions | CSOTTimeoutContextOptions) & {125  session?: ClientSession;126};127 128/** @internal */129export type LegacyTimeoutContextOptions = {130  serverSelectionTimeoutMS: number;131  waitQueueTimeoutMS: number;132  socketTimeoutMS?: number;133};134 135/** @internal */136export type CSOTTimeoutContextOptions = {137  timeoutMS: number;138  serverSelectionTimeoutMS: number;139  socketTimeoutMS?: number;140};141 142function isLegacyTimeoutContextOptions(v: unknown): v is LegacyTimeoutContextOptions {143  return (144    v != null &&145    typeof v === 'object' &&146    'serverSelectionTimeoutMS' in v &&147    typeof v.serverSelectionTimeoutMS === 'number' &&148    'waitQueueTimeoutMS' in v &&149    typeof v.waitQueueTimeoutMS === 'number'150  );151}152 153function isCSOTTimeoutContextOptions(v: unknown): v is CSOTTimeoutContextOptions {154  return (155    v != null &&156    typeof v === 'object' &&157    'serverSelectionTimeoutMS' in v &&158    typeof v.serverSelectionTimeoutMS === 'number' &&159    'timeoutMS' in v &&160    typeof v.timeoutMS === 'number'161  );162}163 164/** @internal */165export abstract class TimeoutContext {166  static create(options: TimeoutContextOptions): TimeoutContext {167    if (options.session?.timeoutContext != null) return options.session?.timeoutContext;168    if (isCSOTTimeoutContextOptions(options)) return new CSOTTimeoutContext(options);169    else if (isLegacyTimeoutContextOptions(options)) return new LegacyTimeoutContext(options);170    else throw new MongoRuntimeError('Unrecognized options');171  }172 173  abstract get maxTimeMS(): number | null;174 175  abstract get serverSelectionTimeout(): Timeout | null;176 177  abstract get connectionCheckoutTimeout(): Timeout | null;178 179  abstract get clearServerSelectionTimeout(): boolean;180 181  abstract get timeoutForSocketWrite(): Timeout | null;182 183  abstract get timeoutForSocketRead(): Timeout | null;184 185  abstract csotEnabled(): this is CSOTTimeoutContext;186 187  abstract refresh(): void;188 189  abstract clear(): void;190 191  /** Returns a new instance of the TimeoutContext, with all timeouts refreshed and restarted. */192  abstract refreshed(): TimeoutContext;193 194  abstract addMaxTimeMSToCommand(command: Document, options: { omitMaxTimeMS?: boolean }): void;195 196  abstract getSocketTimeoutMS(): number | undefined;197}198 199/** @internal */200export class CSOTTimeoutContext extends TimeoutContext {201  timeoutMS: number;202  serverSelectionTimeoutMS: number;203  socketTimeoutMS?: number;204 205  clearServerSelectionTimeout: boolean;206 207  private _serverSelectionTimeout?: Timeout | null;208  private _connectionCheckoutTimeout?: Timeout | null;209  public minRoundTripTime = 0;210  public start: number;211 212  constructor(options: CSOTTimeoutContextOptions) {213    super();214    this.start = Math.trunc(performance.now());215 216    this.timeoutMS = options.timeoutMS;217 218    this.serverSelectionTimeoutMS = options.serverSelectionTimeoutMS;219 220    this.socketTimeoutMS = options.socketTimeoutMS;221 222    this.clearServerSelectionTimeout = false;223  }224 225  get maxTimeMS(): number {226    return this.remainingTimeMS - this.minRoundTripTime;227  }228 229  get remainingTimeMS() {230    const timePassed = Math.trunc(performance.now()) - this.start;231    return this.timeoutMS <= 0 ? Infinity : this.timeoutMS - timePassed;232  }233 234  csotEnabled(): this is CSOTTimeoutContext {235    return true;236  }237 238  get serverSelectionTimeout(): Timeout | null {239    // check for undefined240    if (typeof this._serverSelectionTimeout !== 'object' || this._serverSelectionTimeout?.cleared) {241      const { remainingTimeMS, serverSelectionTimeoutMS } = this;242      if (remainingTimeMS <= 0)243        return Timeout.reject(244          new MongoOperationTimeoutError(`Timed out in server selection after ${this.timeoutMS}ms`)245        );246      const usingServerSelectionTimeoutMS =247        serverSelectionTimeoutMS !== 0 &&248        csotMin(remainingTimeMS, serverSelectionTimeoutMS) === serverSelectionTimeoutMS;249      if (usingServerSelectionTimeoutMS) {250        this._serverSelectionTimeout = Timeout.expires(serverSelectionTimeoutMS);251      } else {252        if (remainingTimeMS > 0 && Number.isFinite(remainingTimeMS)) {253          this._serverSelectionTimeout = Timeout.expires(remainingTimeMS);254        } else {255          this._serverSelectionTimeout = null;256        }257      }258    }259 260    return this._serverSelectionTimeout;261  }262 263  get connectionCheckoutTimeout(): Timeout | null {264    if (265      typeof this._connectionCheckoutTimeout !== 'object' ||266      this._connectionCheckoutTimeout?.cleared267    ) {268      if (typeof this._serverSelectionTimeout === 'object') {269        // null or Timeout270        this._connectionCheckoutTimeout = this._serverSelectionTimeout;271      } else {272        throw new MongoRuntimeError(273          'Unreachable. If you are seeing this error, please file a ticket on the NODE driver project on Jira'274        );275      }276    }277    return this._connectionCheckoutTimeout;278  }279 280  get timeoutForSocketWrite(): Timeout | null {281    const { remainingTimeMS } = this;282    if (!Number.isFinite(remainingTimeMS)) return null;283    if (remainingTimeMS > 0) return Timeout.expires(remainingTimeMS);284    return Timeout.reject(new MongoOperationTimeoutError('Timed out before socket write'));285  }286 287  get timeoutForSocketRead(): Timeout | null {288    const { remainingTimeMS } = this;289    if (!Number.isFinite(remainingTimeMS)) return null;290    if (remainingTimeMS > 0) return Timeout.expires(remainingTimeMS);291    return Timeout.reject(new MongoOperationTimeoutError('Timed out before socket read'));292  }293 294  refresh(): void {295    this.start = Math.trunc(performance.now());296    this.minRoundTripTime = 0;297    this._serverSelectionTimeout?.clear();298    this._connectionCheckoutTimeout?.clear();299  }300 301  clear(): void {302    this._serverSelectionTimeout?.clear();303    this._connectionCheckoutTimeout?.clear();304  }305 306  /**307   * @internal308   * Throws a MongoOperationTimeoutError if the context has expired.309   * If the context has not expired, returns the `remainingTimeMS`310   **/311  getRemainingTimeMSOrThrow(message?: string): number {312    const { remainingTimeMS } = this;313    if (remainingTimeMS <= 0)314      throw new MongoOperationTimeoutError(message ?? `Expired after ${this.timeoutMS}ms`);315    return remainingTimeMS;316  }317 318  /**319   * @internal320   * This method is intended to be used in situations where concurrent operation are on the same deadline, but cannot share a single `TimeoutContext` instance.321   * Returns a new instance of `CSOTTimeoutContext` constructed with identical options, but setting the `start` property to `this.start`.322   */323  clone(): CSOTTimeoutContext {324    const timeoutContext = new CSOTTimeoutContext({325      timeoutMS: this.timeoutMS,326      serverSelectionTimeoutMS: this.serverSelectionTimeoutMS327    });328    timeoutContext.start = this.start;329    return timeoutContext;330  }331 332  override refreshed(): CSOTTimeoutContext {333    return new CSOTTimeoutContext(this);334  }335 336  override addMaxTimeMSToCommand(command: Document, options: { omitMaxTimeMS?: boolean }): void {337    if (options.omitMaxTimeMS) return;338    const maxTimeMS = this.remainingTimeMS - this.minRoundTripTime;339    if (maxTimeMS > 0 && Number.isFinite(maxTimeMS)) command.maxTimeMS = maxTimeMS;340  }341 342  override getSocketTimeoutMS(): number | undefined {343    return 0;344  }345}346 347/** @internal */348export class LegacyTimeoutContext extends TimeoutContext {349  options: LegacyTimeoutContextOptions;350  clearServerSelectionTimeout: boolean;351 352  constructor(options: LegacyTimeoutContextOptions) {353    super();354    this.options = options;355    this.clearServerSelectionTimeout = true;356  }357 358  csotEnabled(): this is CSOTTimeoutContext {359    return false;360  }361 362  get serverSelectionTimeout(): Timeout | null {363    if (this.options.serverSelectionTimeoutMS != null && this.options.serverSelectionTimeoutMS > 0)364      return Timeout.expires(this.options.serverSelectionTimeoutMS);365    return null;366  }367 368  get connectionCheckoutTimeout(): Timeout | null {369    if (this.options.waitQueueTimeoutMS != null && this.options.waitQueueTimeoutMS > 0)370      return Timeout.expires(this.options.waitQueueTimeoutMS);371    return null;372  }373 374  get timeoutForSocketWrite(): Timeout | null {375    return null;376  }377 378  get timeoutForSocketRead(): Timeout | null {379    return null;380  }381 382  refresh(): void {383    return;384  }385 386  clear(): void {387    return;388  }389 390  get maxTimeMS() {391    return null;392  }393 394  override refreshed(): LegacyTimeoutContext {395    return new LegacyTimeoutContext(this.options);396  }397 398  override addMaxTimeMSToCommand(_command: Document, _options: { omitMaxTimeMS?: boolean }): void {399    // No max timeMS is added to commands in legacy timeout mode.400  }401 402  override getSocketTimeoutMS(): number | undefined {403    return this.options.socketTimeoutMS;404  }405}406