opusdev/vector-similarity-api

1import * as dns from 'dns';2import { clearTimeout, setTimeout } from 'timers';3 4import { MongoRuntimeError } from '../error';5import { TypedEventEmitter } from '../mongo_types';6import { checkParentDomainMatch, HostAddress, noop, squashError } from '../utils';7 8/**9 * @internal10 * @category Event11 */12export class SrvPollingEvent {13  srvRecords: dns.SrvRecord[];14  constructor(srvRecords: dns.SrvRecord[]) {15    this.srvRecords = srvRecords;16  }17 18  hostnames(): Set<string> {19    return new Set(this.srvRecords.map(r => HostAddress.fromSrvRecord(r).toString()));20  }21}22 23/** @internal */24export interface SrvPollerOptions {25  srvServiceName: string;26  srvMaxHosts: number;27  srvHost: string;28  heartbeatFrequencyMS: number;29}30 31/** @internal */32export type SrvPollerEvents = {33  srvRecordDiscovery(event: SrvPollingEvent): void;34};35 36/** @internal */37export class SrvPoller extends TypedEventEmitter<SrvPollerEvents> {38  srvHost: string;39  rescanSrvIntervalMS: number;40  heartbeatFrequencyMS: number;41  haMode: boolean;42  generation: number;43  srvMaxHosts: number;44  srvServiceName: string;45  _timeout?: NodeJS.Timeout;46 47  /** @event */48  static readonly SRV_RECORD_DISCOVERY = 'srvRecordDiscovery' as const;49 50  constructor(options: SrvPollerOptions) {51    super();52    this.on('error', noop);53 54    if (!options || !options.srvHost) {55      throw new MongoRuntimeError('Options for SrvPoller must exist and include srvHost');56    }57 58    this.srvHost = options.srvHost;59    this.srvMaxHosts = options.srvMaxHosts ?? 0;60    this.srvServiceName = options.srvServiceName ?? 'mongodb';61    this.rescanSrvIntervalMS = 60000;62    this.heartbeatFrequencyMS = options.heartbeatFrequencyMS ?? 10000;63 64    this.haMode = false;65    this.generation = 0;66 67    this._timeout = undefined;68  }69 70  get srvAddress(): string {71    return `_${this.srvServiceName}._tcp.${this.srvHost}`;72  }73 74  get intervalMS(): number {75    return this.haMode ? this.heartbeatFrequencyMS : this.rescanSrvIntervalMS;76  }77 78  start(): void {79    if (!this._timeout) {80      this.schedule();81    }82  }83 84  stop(): void {85    if (this._timeout) {86      clearTimeout(this._timeout);87      this.generation += 1;88      this._timeout = undefined;89    }90  }91 92  // TODO(NODE-4994): implement new logging logic for SrvPoller failures93  schedule(): void {94    if (this._timeout) {95      clearTimeout(this._timeout);96    }97 98    this._timeout = setTimeout(() => {99      this._poll().then(undefined, squashError);100    }, this.intervalMS);101  }102 103  success(srvRecords: dns.SrvRecord[]): void {104    this.haMode = false;105    this.schedule();106    this.emit(SrvPoller.SRV_RECORD_DISCOVERY, new SrvPollingEvent(srvRecords));107  }108 109  failure(): void {110    this.haMode = true;111    this.schedule();112  }113 114  async _poll(): Promise<void> {115    const generation = this.generation;116    let srvRecords;117 118    try {119      srvRecords = await dns.promises.resolveSrv(this.srvAddress);120    } catch {121      this.failure();122      return;123    }124 125    if (generation !== this.generation) {126      return;127    }128 129    const finalAddresses: dns.SrvRecord[] = [];130    for (const record of srvRecords) {131      try {132        checkParentDomainMatch(record.name, this.srvHost);133        finalAddresses.push(record);134      } catch (error) {135        squashError(error);136      }137    }138 139    if (!finalAddresses.length) {140      this.failure();141      return;142    }143 144    this.success(finalAddresses);145  }146}147