opusdev/vector-similarity-api

1import { MongoInvalidArgumentError } from '../error';2import { ReadPreference } from '../read_preference';3import { ServerType, TopologyType } from './common';4import type { ServerDescription, TagSet } from './server_description';5import type { TopologyDescription } from './topology_description';6 7// max staleness constants8const IDLE_WRITE_PERIOD = 10000;9const SMALLEST_MAX_STALENESS_SECONDS = 90;10 11//  Minimum version to try writes on secondaries.12export const MIN_SECONDARY_WRITE_WIRE_VERSION = 13;13 14/** @internal */15export type ServerSelector = (16  topologyDescription: TopologyDescription,17  servers: ServerDescription[],18  deprioritized?: ServerDescription[]19) => ServerDescription[];20 21/**22 * Returns a server selector that selects for writable servers23 */24export function writableServerSelector(): ServerSelector {25  return function writableServer(26    topologyDescription: TopologyDescription,27    servers: ServerDescription[]28  ): ServerDescription[] {29    return latencyWindowReducer(30      topologyDescription,31      servers.filter((s: ServerDescription) => s.isWritable)32    );33  };34}35 36/**37 * The purpose of this selector is to select the same server, only38 * if it is in a state that it can have commands sent to it.39 */40export function sameServerSelector(description?: ServerDescription): ServerSelector {41  return function sameServerSelector(42    topologyDescription: TopologyDescription,43    servers: ServerDescription[]44  ): ServerDescription[] {45    if (!description) return [];46    // Filter the servers to match the provided description only if47    // the type is not unknown.48    return servers.filter(sd => {49      return sd.address === description.address && sd.type !== ServerType.Unknown;50    });51  };52}53 54/**55 * Returns a server selector that uses a read preference to select a56 * server potentially for a write on a secondary.57 */58export function secondaryWritableServerSelector(59  wireVersion?: number,60  readPreference?: ReadPreference61): ServerSelector {62  // If server version < 5.0, read preference always primary.63  // If server version >= 5.0...64  // - If read preference is supplied, use that.65  // - If no read preference is supplied, use primary.66  if (67    !readPreference ||68    !wireVersion ||69    (wireVersion && wireVersion < MIN_SECONDARY_WRITE_WIRE_VERSION)70  ) {71    return readPreferenceServerSelector(ReadPreference.primary);72  }73  return readPreferenceServerSelector(readPreference);74}75 76/**77 * Reduces the passed in array of servers by the rules of the "Max Staleness" specification78 * found here:79 *80 * @see https://github.com/mongodb/specifications/blob/master/source/max-staleness/max-staleness.md81 *82 * @param readPreference - The read preference providing max staleness guidance83 * @param topologyDescription - The topology description84 * @param servers - The list of server descriptions to be reduced85 * @returns The list of servers that satisfy the requirements of max staleness86 */87function maxStalenessReducer(88  readPreference: ReadPreference,89  topologyDescription: TopologyDescription,90  servers: ServerDescription[]91): ServerDescription[] {92  if (readPreference.maxStalenessSeconds == null || readPreference.maxStalenessSeconds < 0) {93    return servers;94  }95 96  const maxStaleness = readPreference.maxStalenessSeconds;97  const maxStalenessVariance =98    (topologyDescription.heartbeatFrequencyMS + IDLE_WRITE_PERIOD) / 1000;99  if (maxStaleness < maxStalenessVariance) {100    throw new MongoInvalidArgumentError(101      `Option "maxStalenessSeconds" must be at least ${maxStalenessVariance} seconds`102    );103  }104 105  if (maxStaleness < SMALLEST_MAX_STALENESS_SECONDS) {106    throw new MongoInvalidArgumentError(107      `Option "maxStalenessSeconds" must be at least ${SMALLEST_MAX_STALENESS_SECONDS} seconds`108    );109  }110 111  if (topologyDescription.type === TopologyType.ReplicaSetWithPrimary) {112    const primary: ServerDescription = Array.from(topologyDescription.servers.values()).filter(113      primaryFilter114    )[0];115 116    return servers.reduce((result: ServerDescription[], server: ServerDescription) => {117      const stalenessMS =118        server.lastUpdateTime -119        server.lastWriteDate -120        (primary.lastUpdateTime - primary.lastWriteDate) +121        topologyDescription.heartbeatFrequencyMS;122 123      const staleness = stalenessMS / 1000;124      const maxStalenessSeconds = readPreference.maxStalenessSeconds ?? 0;125      if (staleness <= maxStalenessSeconds) {126        result.push(server);127      }128 129      return result;130    }, []);131  }132 133  if (topologyDescription.type === TopologyType.ReplicaSetNoPrimary) {134    if (servers.length === 0) {135      return servers;136    }137 138    const sMax = servers.reduce((max: ServerDescription, s: ServerDescription) =>139      s.lastWriteDate > max.lastWriteDate ? s : max140    );141 142    return servers.reduce((result: ServerDescription[], server: ServerDescription) => {143      const stalenessMS =144        sMax.lastWriteDate - server.lastWriteDate + topologyDescription.heartbeatFrequencyMS;145 146      const staleness = stalenessMS / 1000;147      const maxStalenessSeconds = readPreference.maxStalenessSeconds ?? 0;148      if (staleness <= maxStalenessSeconds) {149        result.push(server);150      }151 152      return result;153    }, []);154  }155 156  return servers;157}158 159/**160 * Determines whether a server's tags match a given set of tags161 *162 * @param tagSet - The requested tag set to match163 * @param serverTags - The server's tags164 */165function tagSetMatch(tagSet: TagSet, serverTags: TagSet) {166  const keys = Object.keys(tagSet);167  const serverTagKeys = Object.keys(serverTags);168  for (let i = 0; i < keys.length; ++i) {169    const key = keys[i];170    if (serverTagKeys.indexOf(key) === -1 || serverTags[key] !== tagSet[key]) {171      return false;172    }173  }174 175  return true;176}177 178/**179 * Reduces a set of server descriptions based on tags requested by the read preference180 *181 * @param readPreference - The read preference providing the requested tags182 * @param servers - The list of server descriptions to reduce183 * @returns The list of servers matching the requested tags184 */185function tagSetReducer(186  readPreference: ReadPreference,187  servers: ServerDescription[]188): ServerDescription[] {189  if (190    readPreference.tags == null ||191    (Array.isArray(readPreference.tags) && readPreference.tags.length === 0)192  ) {193    return servers;194  }195 196  for (let i = 0; i < readPreference.tags.length; ++i) {197    const tagSet = readPreference.tags[i];198    const serversMatchingTagset = servers.reduce(199      (matched: ServerDescription[], server: ServerDescription) => {200        if (tagSetMatch(tagSet, server.tags)) matched.push(server);201        return matched;202      },203      []204    );205 206    if (serversMatchingTagset.length) {207      return serversMatchingTagset;208    }209  }210 211  return [];212}213 214/**215 * Reduces a list of servers to ensure they fall within an acceptable latency window. This is216 * further specified in the "Server Selection" specification, found here:217 *218 * @see https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.md219 *220 * @param topologyDescription - The topology description221 * @param servers - The list of servers to reduce222 * @returns The servers which fall within an acceptable latency window223 */224function latencyWindowReducer(225  topologyDescription: TopologyDescription,226  servers: ServerDescription[]227): ServerDescription[] {228  const low = servers.reduce(229    (min: number, server: ServerDescription) => Math.min(server.roundTripTime, min),230    Infinity231  );232 233  const high = low + topologyDescription.localThresholdMS;234  return servers.reduce((result: ServerDescription[], server: ServerDescription) => {235    if (server.roundTripTime <= high && server.roundTripTime >= low) result.push(server);236    return result;237  }, []);238}239 240// filters241function primaryFilter(server: ServerDescription): boolean {242  return server.type === ServerType.RSPrimary;243}244 245function secondaryFilter(server: ServerDescription): boolean {246  return server.type === ServerType.RSSecondary;247}248 249function nearestFilter(server: ServerDescription): boolean {250  return server.type === ServerType.RSSecondary || server.type === ServerType.RSPrimary;251}252 253function knownFilter(server: ServerDescription): boolean {254  return server.type !== ServerType.Unknown;255}256 257function loadBalancerFilter(server: ServerDescription): boolean {258  return server.type === ServerType.LoadBalancer;259}260 261/**262 * Returns a function which selects servers based on a provided read preference263 *264 * @param readPreference - The read preference to select with265 */266export function readPreferenceServerSelector(readPreference: ReadPreference): ServerSelector {267  if (!readPreference.isValid()) {268    throw new MongoInvalidArgumentError('Invalid read preference specified');269  }270 271  return function readPreferenceServers(272    topologyDescription: TopologyDescription,273    servers: ServerDescription[],274    deprioritized: ServerDescription[] = []275  ): ServerDescription[] {276    if (topologyDescription.type === TopologyType.LoadBalanced) {277      return servers.filter(loadBalancerFilter);278    }279 280    if (topologyDescription.type === TopologyType.Unknown) {281      return [];282    }283 284    if (topologyDescription.type === TopologyType.Single) {285      return latencyWindowReducer(topologyDescription, servers.filter(knownFilter));286    }287 288    if (topologyDescription.type === TopologyType.Sharded) {289      const filtered = servers.filter(server => {290        return !deprioritized.includes(server);291      });292      const selectable = filtered.length > 0 ? filtered : deprioritized;293      return latencyWindowReducer(topologyDescription, selectable.filter(knownFilter));294    }295 296    const mode = readPreference.mode;297    if (mode === ReadPreference.PRIMARY) {298      return servers.filter(primaryFilter);299    }300 301    if (mode === ReadPreference.PRIMARY_PREFERRED) {302      const result = servers.filter(primaryFilter);303      if (result.length) {304        return result;305      }306    }307 308    const filter = mode === ReadPreference.NEAREST ? nearestFilter : secondaryFilter;309    const selectedServers = latencyWindowReducer(310      topologyDescription,311      tagSetReducer(312        readPreference,313        maxStalenessReducer(readPreference, topologyDescription, servers.filter(filter))314      )315    );316 317    if (mode === ReadPreference.SECONDARY_PREFERRED && selectedServers.length === 0) {318      return servers.filter(primaryFilter);319    }320 321    return selectedServers;322  };323}324