feat: make measurement cloning work

alexey-yarmosh · alexey-yarmosh · commit 93099931d17f · 2023-11-29T13:08:06.000+01:00
diff --git a/src/measurement/route/get-measurement.ts b/src/measurement/route/get-measurement.ts
@@ -12,7 +12,7 @@ const handle = async (ctx: ParameterizedContext<DefaultState, DefaultContext & R
 		return;
 	}
 
-	const result = await store.getMeasurementResults(id);
+	const result = await store.getMeasurement(id);
 
 	if (!result) {
 		ctx.status = 404;
diff --git a/src/measurement/store.ts b/src/measurement/store.ts
@@ -5,7 +5,7 @@ import type { Probe } from '../probe/types.js';
 import type { RedisClient } from '../lib/redis/client.js';
 import { getRedisClient } from '../lib/redis/client.js';
 import { scopedLogger } from '../lib/logger.js';
-import type { MeasurementRecord, MeasurementResult, MeasurementRequest, MeasurementProgressMessage } from './types.js';
+import type { MeasurementRecord, MeasurementResult, MeasurementRequest, MeasurementProgressMessage, RequestType } from './types.js';
 import { getDefaults } from './schema/utils.js';
 
 const logger = scopedLogger('store');
@@ -38,10 +38,14 @@ const substractObjects = (obj1: Record<string, unknown>, obj2: Record<string, un
 export class MeasurementStore {
 	constructor (private readonly redis: RedisClient) {}
 
-	async getMeasurementResults (id: string): Promise<string> {
+	async getMeasurement (id: string): Promise<string> {
 		return this.redis.sendCommand([ 'JSON.GET', getMeasurementKey(id) ]);
 	}
 
+	async getMeasurementJson (id: string): Promise<MeasurementRecord> {
+		return await this.redis.json.get(getMeasurementKey(id)) as MeasurementRecord;
+	}
+
 	async createMeasurement (request: MeasurementRequest, probes: Probe[]): Promise<string> {
 		const id = cryptoRandomString({ length: 16, type: 'alphanumeric' });
 		const key = getMeasurementKey(id);
@@ -66,7 +70,7 @@ export class MeasurementStore {
 
 		await Promise.all([
 			this.redis.hSet('gp:in-progress', id, startTime.getTime()),
-			this.redis.set(getMeasurementKey(id, 'probes_awaiting'), probes.length, { EX: probesAwaitingTtl }),
+			this.redis.set(getMeasurementKey(id, 'probes_awaiting'), probes.filter(probe => probe.status !== 'offline').length, { EX: probesAwaitingTtl }),
 			this.redis.json.set(getMeasurementKey(id, 'ips'), '$', probes.map(probe => probe.ipAddress)),
 			this.redis.json.set(key, '$', measurementWithoutDefaults),
 			this.redis.expire(getMeasurementKey(id, 'ips'), config.get<number>('measurement.resultTTL')),
@@ -166,7 +170,7 @@ export class MeasurementStore {
 		return substractObjects(measurement, defaults) as Partial<MeasurementRecord>;
 	}
 
-	probesToResults (probes: Probe[], type: string) {
+	probesToResults (probes: Probe[], type: RequestType) {
 		const results = probes.map(probe => ({
 			probe: {
 				continent: probe.location.continent,
@@ -181,13 +185,20 @@ export class MeasurementStore {
 				tags: probe.tags.map(({ value }) => value),
 				resolvers: probe.resolvers,
 			},
-			result: this.getInitialResult(type),
+			result: this.getInitialResult(type, probe.status),
 		} as MeasurementResult));
 
 		return results;
 	}
 
-	getInitialResult (type: string) {
+	getInitialResult (type: RequestType, status: Probe['status']) {
+		if (status === 'offline') {
+			return {
+				status: 'failed',
+				rawOutput: 'This probe is currently offline. Please try again later.',
+			};
+		}
+
 		if (type === 'http') {
 			return {
 				status: 'in-progress',
diff --git a/src/probe/probes-location-filter.ts b/src/probe/probes-location-filter.ts
@@ -0,0 +1,152 @@
+import config from 'config';
+import _ from 'lodash';
+import type { Location } from '../lib/location/types.js';
+import type { Probe, ProbeLocation } from './types.js';
+
+/*
+ * [
+ *    [ public key, internal key]
+ * ]
+ *
+ * */
+const locationKeyMap = [
+	[ 'network', 'normalizedNetwork' ],
+	[ 'city', 'normalizedCity' ],
+];
+
+export class ProbesLocationFilter {
+	static magicFilter (probes: Probe[], magicLocation: string) {
+		let filteredProbes = probes;
+		const keywords = magicLocation.split('+');
+
+		for (const keyword of keywords) {
+			const closestExactMatchPosition = probes.reduce((smallestExactMatchPosition, probe) => {
+				const exactMatchPosition = ProbesLocationFilter.getExactIndexPosition(probe, keyword);
+
+				if (exactMatchPosition === -1) {
+					return smallestExactMatchPosition;
+				}
+
+				return exactMatchPosition < smallestExactMatchPosition ? exactMatchPosition : smallestExactMatchPosition;
+			}, Number.POSITIVE_INFINITY);
+			const noExactMatches = closestExactMatchPosition === Number.POSITIVE_INFINITY;
+
+			if (noExactMatches) {
+				filteredProbes = filteredProbes.filter(probe => ProbesLocationFilter.getIndexPosition(probe, keyword) !== -1);
+			} else {
+				filteredProbes = filteredProbes.filter(probe => ProbesLocationFilter.getExactIndexPosition(probe, keyword) === closestExactMatchPosition);
+			}
+		}
+
+		return filteredProbes;
+	}
+
+	static getExactIndexPosition (probe: Probe, value: string) {
+		return probe.index.findIndex(category => category.some(index => index === value.replaceAll('-', ' ').trim()));
+	}
+
+	static getIndexPosition (probe: Probe, value: string) {
+		return probe.index.findIndex(category => category.some(index => index.includes(value.replaceAll('-', ' ').trim())));
+	}
+
+	static hasTag (probe: Probe, tag: string) {
+		return probe.tags.some(({ value }) => value.toLowerCase() === tag);
+	}
+
+	public filterGloballyDistibuted (probes: Probe[], limit: number): Probe[] {
+		const distribution = this.getDistibutionConfig();
+		return this.filterByLocationAndWeight(probes, distribution, limit);
+	}
+
+	public filterByLocation (probes: Probe[], location: Location): Probe[] {
+		if (location.magic === 'world') {
+			return _.shuffle(this.filterGloballyDistibuted(probes, probes.length));
+		}
+
+		let filteredProbes = probes;
+
+		Object.keys(location).forEach((key) => {
+			if (key === 'tags') {
+				// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+				probes = probes.filter(probe => location.tags!.every(tag => ProbesLocationFilter.hasTag(probe, tag)));
+			} else if (key === 'magic') {
+				// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+				filteredProbes = ProbesLocationFilter.magicFilter(filteredProbes, location.magic!);
+			} else {
+				const probeKey = locationKeyMap.find(m => m.includes(key))?.[1] ?? key;
+				filteredProbes = filteredProbes.filter(probe => location[key as keyof Location] === probe.location[probeKey as keyof ProbeLocation]);
+			}
+		});
+
+		const isMagicSorting = Object.keys(location).includes('magic');
+		// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+		return isMagicSorting ? this.magicSort(filteredProbes, location.magic!) : _.shuffle(filteredProbes);
+	}
+
+	public filterByLocationAndWeight (probes: Probe[], distribution: Map<Location, number>, limit: number): Probe[] {
+		const groupedByLocation = new Map<Location, Probe[]>();
+
+		for (const [ location ] of distribution) {
+			const foundProbes = this.filterByLocation(probes, location);
+
+			if (foundProbes.length > 0) {
+				groupedByLocation.set(location, foundProbes);
+			}
+		}
+
+		const pickedProbes = new Set<Probe>();
+
+		while (groupedByLocation.size > 0 && pickedProbes.size < limit) {
+			const selectedCount = pickedProbes.size;
+
+			for (const [ location, locationProbes ] of groupedByLocation) {
+				if (pickedProbes.size === limit) {
+					break;
+				}
+
+				const locationWeight = distribution.get(location);
+
+				if (!locationWeight) {
+					continue;
+				}
+
+				let count = Math.ceil((limit - selectedCount) * locationWeight / 100);
+				const remainingSpace = limit - pickedProbes.size;
+				count = count > remainingSpace ? remainingSpace : count;
+
+				for (const s of locationProbes.splice(0, count)) {
+					pickedProbes.add(s);
+				}
+
+				if (locationProbes.length === 0) {
+					groupedByLocation.delete(location);
+				}
+			}
+		}
+
+		return [ ...pickedProbes ];
+	}
+
+	private magicSort (probes: Probe[], magicString: string): Probe[] {
+		const getClosestIndexPosition = (probe: Probe) => {
+			const keywords = magicString.split('+');
+			const closestIndexPosition = keywords.reduce((smallesIndex, keyword) => {
+				const indexPosition = ProbesLocationFilter.getIndexPosition(probe, keyword);
+				return indexPosition < smallesIndex ? indexPosition : smallesIndex;
+			}, Number.POSITIVE_INFINITY);
+			return closestIndexPosition;
+		};
+
+		const probesGroupedByIndexPosition = _.groupBy(probes, getClosestIndexPosition);
+		const groupsSortedByIndexPosition = Object.values(probesGroupedByIndexPosition); // Object.values sorts values by key
+		const groupsWithShuffledItems = groupsSortedByIndexPosition.map(group => _.shuffle(group));
+		const resultProbes = groupsWithShuffledItems.flat();
+
+		return resultProbes;
+	}
+
+	private getDistibutionConfig () {
+		return new Map<Location, number>(_.shuffle(Object.entries(config.get<Record<string, number>>('measurement.globalDistribution')))
+			.map(([ value, weight ]) => [{ continent: value }, weight ]));
+	}
+}
diff --git a/src/probe/router.ts b/src/probe/router.ts
@@ -1,14 +1,13 @@
 import _ from 'lodash';
-import type { RemoteProbeSocket } from '../lib/ws/server.js';
 import { fetchSockets } from '../lib/ws/fetch-sockets.js';
-import type { LocationWithLimit } from '../measurement/types.js';
+import type { LocationWithLimit, MeasurementRecord } from '../measurement/types.js';
 import type { Location } from '../lib/location/types.js';
 import type { Probe } from './types.js';
-import { SocketsLocationFilter } from './sockets-location-filter.js';
+import { ProbesLocationFilter } from './probes-location-filter.js';
 import { getMeasurementStore } from '../measurement/store.js';
 
 export class ProbeRouter {
-	private readonly socketsFilter = new SocketsLocationFilter();
+	private readonly probesFilter = new ProbesLocationFilter();
 
 	private readonly store = getMeasurementStore();
 
@@ -18,51 +17,51 @@ export class ProbeRouter {
 		locations: LocationWithLimit[] | string = [],
 		globalLimit = 1,
 	): Promise<Probe[]> {
-		const sockets = await this.fetchSockets();
-		let filtered: RemoteProbeSocket[] = [];
+		const probes = await this.fetchProbes();
+		let filtered: Probe[] = [];
 
 		if (typeof locations === 'string') {
-			filtered = await this.findWithMeasurementId(sockets, locations);
+			filtered = await this.findWithMeasurementId(probes, locations);
 		} else if (locations.some(l => l.limit)) {
-			filtered = this.findWithLocationLimit(sockets, locations);
+			filtered = this.findWithLocationLimit(probes, locations);
 		} else if (locations.length > 0) {
-			filtered = this.findWithGlobalLimit(sockets, locations, globalLimit);
+			filtered = this.findWithGlobalLimit(probes, locations, globalLimit);
 		} else {
-			filtered = this.findGloballyDistributed(sockets, globalLimit);
+			filtered = this.findGloballyDistributed(probes, globalLimit);
 		}
 
-		return filtered.map(s => s.data.probe);
+		return filtered;
 	}
 
-	private async fetchSockets (): Promise<RemoteProbeSocket[]> {
+	private async fetchProbes (): Promise<Probe[]> {
 		const sockets = await this.fetchWsSockets();
-		return sockets.filter(s => s.data.probe.status === 'ready');
+		return sockets.filter(s => s.data.probe.status === 'ready').map(s => s.data.probe);
 	}
 
-	private findGloballyDistributed (sockets: RemoteProbeSocket[], limit: number): RemoteProbeSocket[] {
-		return this.socketsFilter.filterGloballyDistibuted(sockets, limit);
+	private findGloballyDistributed (probes: Probe[], limit: number): Probe[] {
+		return this.probesFilter.filterGloballyDistibuted(probes, limit);
 	}
 
-	private findWithGlobalLimit (sockets: RemoteProbeSocket[], locations: Location[], limit: number): RemoteProbeSocket[] {
+	private findWithGlobalLimit (probes: Probe[], locations: Location[], limit: number): Probe[] {
 		const weight = Math.floor(100 / locations.length);
 		const distribution = new Map(locations.map(l => [ l, weight ]));
 
-		return this.socketsFilter.filterByLocationAndWeight(sockets, distribution, limit);
+		return this.probesFilter.filterByLocationAndWeight(probes, distribution, limit);
 	}
 
-	private findWithLocationLimit (sockets: RemoteProbeSocket[], locations: LocationWithLimit[]): RemoteProbeSocket[] {
-		const grouped = new Map<LocationWithLimit, RemoteProbeSocket[]>();
+	private findWithLocationLimit (probes: Probe[], locations: LocationWithLimit[]): Probe[] {
+		const grouped = new Map<LocationWithLimit, Probe[]>();
 
 		for (const location of locations) {
 			const { limit, ...l } = location;
-			const found = this.socketsFilter.filterByLocation(sockets, l);
+			const found = this.probesFilter.filterByLocation(probes, l);
 
 			if (found.length > 0) {
 				grouped.set(location, found);
 			}
 		}
 
-		const picked = new Set<RemoteProbeSocket>();
+		const picked = new Set<Probe>();
 
 		for (const [ loc, soc ] of grouped) {
 			for (const s of _.take(soc, loc.limit)) {
@@ -73,9 +72,68 @@ export class ProbeRouter {
 		return [ ...picked ];
 	}
 
-	private async findWithMeasurementId (sockets: RemoteProbeSocket[], measurementId: string): Promise<RemoteProbeSocket[]> {
-		const ips = await this.store.getIpsByMeasurementId(measurementId);
-		return sockets.filter(socket => ips.includes(socket.data.probe.ipAddress));
+	private async findWithMeasurementId (probes: Probe[], measurementId: string): Promise<Probe[]> {
+		let prevMeasurement: MeasurementRecord | undefined;
+		const prevIps = await this.store.getIpsByMeasurementId(measurementId);
+		const ipToProbe = new Map(probes.map(probe => [ probe.ipAddress, probe ]));
+		const result: Probe[] = [];
+
+		for (let i = 0; i < prevIps.length; i++) {
+			const ip = prevIps[i]!;
+			const probe = ipToProbe.get(ip);
+
+			if (probe) {
+				result.push(probe);
+			} else {
+				if (!prevMeasurement) {
+					prevMeasurement = await this.store.getMeasurementJson(measurementId);
+
+					if (!prevMeasurement) { return []; }
+				}
+
+				const prevTest = prevMeasurement.results[i];
+
+				if (!prevTest) { return []; }
+
+				const offlineProbe: Probe = {
+					status: 'offline',
+					client: '',
+					version: '',
+					nodeVersion: '',
+					uuid: '',
+					isHardware: false,
+					hardwareDevice: null,
+					ipAddress: ip,
+					host: '',
+					location: {
+						continent: prevTest.probe.continent,
+						region: prevTest.probe.region,
+						country: prevTest.probe.country,
+						city: prevTest.probe.city,
+						normalizedCity: prevTest.probe.city.toLowerCase(),
+						asn: prevTest.probe.asn,
+						latitude: prevTest.probe.latitude,
+						longitude: prevTest.probe.longitude,
+						state: prevTest.probe.state ?? undefined,
+						network: prevTest.probe.network,
+						normalizedNetwork: prevTest.probe.network.toLowerCase(),
+					},
+					index: [],
+					resolvers: prevTest.probe.resolvers,
+					tags: prevTest.probe.tags.map(tag => ({ value: tag, type: 'system' })),
+					stats: {
+						cpu: {
+							count: 0,
+							load: [],
+						},
+						jobs: { count: 0 },
+					},
+				};
+				result.push(offlineProbe);
+			}
+		}
+
+		return result;
 	}
 }
 
diff --git a/src/probe/sockets-location-filter.ts b/src/probe/sockets-location-filter.ts
diff --git a/src/probe/types.ts b/src/probe/types.ts

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ const handle = async (ctx: ParameterizedContext<DefaultState, DefaultContext & R`
`12`	`12`	`return;`
`13`	`13`	`}`
`14`	`14`
`15`		`- const result = await store.getMeasurementResults(id);`
	`15`	`+ const result = await store.getMeasurement(id);`
`16`	`16`
`17`	`17`	`if (!result) {`
`18`	`18`	`ctx.status = 404;`