Add methods for easier checking of lockfile existence

Signed-off-by: Christina Ying Wang <christina@balena.io>

src/lib/update-lock.ts

@@ -12,7 +12,7 @@ import {
 import { pathOnRoot, pathExistsOnState } from './host-utils';
 import * as config from '../config';
 import * as lockfile from './lockfile';
-import { NumericIdentifier } from '../types';
+import { NumericIdentifier, StringIdentifier, DockerName } from '../types';
 
 const decodedUid = NumericIdentifier.decode(process.env.LOCKFILE_UID);
 export const LOCKFILE_UID = isRight(decodedUid) ? decodedUid.right : 65534;
@@ -65,7 +65,7 @@ export const readLock: LockFn = Bluebird.promisify(locker.async.readLock, {
 	context: locker,
 });
 
-// Unlock all lockfiles, optionally of an appId | appUuid, then release resources.
+// Unlock all lockfiles of an appId | appUuid, then release resources.
 async function dispose(
 	appIdentifier: string | number,
 	release: () => void,
@@ -82,6 +82,83 @@ async function dispose(
 	}
 }
 
+/**
+ * Given a lockfile path `p`, return a tuple [appId, serviceName] of that path.
+ * Paths are assumed to end in the format /:appId/:serviceName/(resin-)updates.lock.
+ */
+function getIdentifiersFromPath(p: string) {
+	const parts = p.split('/');
+	if (parts.pop()?.match(/updates\.lock/) === null) {
+		return [];
+	}
+	const serviceName = parts.pop();
+	const appId = parts.pop();
+	return [appId, serviceName];
+}
+
+type LockedEntity = { appId: number; services: string[] };
+
+/**
+ * A map of locked services by appId.
+ * Exported for tests only; getServicesLockedByAppId is the public generator interface.
+ */
+export class LocksTakenMap extends Map<number, Set<string>> {
+	constructor(lockedEntities: LockedEntity[] = []) {
+		// Construct a Map<number, Set<string>> from user-friendly input args
+		super(
+			lockedEntities.map(({ appId, services }) => [appId, new Set(services)]),
+		);
+	}
+
+	// Add one or more locked services to an appId
+	public add(appId: number, services: string | string[]): void {
+		if (typeof services === 'string') {
+			services = [services];
+		}
+		if (this.has(appId)) {
+			const lockedSvcs = this.get(appId)!;
+			services.forEach((s) => lockedSvcs.add(s));
+		} else {
+			this.set(appId, new Set(services));
+		}
+	}
+
+	/**
+	 * @private Use this.getServices instead as there is no need to return
+	 * a mutable reference to the internal Set data structure.
+	 */
+	public get(appId: number): Set<string> | undefined {
+		return super.get(appId);
+	}
+
+	// Return an array copy of locked services under an appId
+	public getServices(appId: number): string[] {
+		return this.has(appId) ? Array.from(this.get(appId)!) : [];
+	}
+
+	// Return whether a service is locked under an appId
+	public isLocked(appId: number, service: string): boolean {
+		return this.has(appId) && this.get(appId)!.has(service);
+	}
+}
+
+/**
+ * Return a list of services that are locked by the Supervisor under each appId.
+ */
+export function getServicesLockedByAppId(): LocksTakenMap {
+	const locksTaken = lockfile.getLocksTaken();
+	const servicesByAppId = new LocksTakenMap();
+	for (const lockTakenPath of locksTaken) {
+		const [appId, serviceName] = getIdentifiersFromPath(lockTakenPath);
+		if (!StringIdentifier.is(appId) || !DockerName.is(serviceName)) {
+			continue;
+		}
+		const numAppId = +appId;
+		servicesByAppId.add(numAppId, serviceName);
+	}
+	return servicesByAppId;
+}
+
 /**
  * Try to take the locks for an application. If force is set, it will remove
  * all existing lockfiles before performing the operation

test/integration/lib/update-lock.spec.ts

@@ -2,6 +2,7 @@ import { expect } from 'chai';
 import * as path from 'path';
 import { promises as fs } from 'fs';
 import { testfs } from 'mocha-pod';
+import type { TestFs } from 'mocha-pod';
 
 import * as updateLock from '~/lib/update-lock';
 import { UpdatesLockedError } from '~/lib/errors';
@@ -281,4 +282,75 @@ describe('lib/update-lock', () => {
 			);
 		});
 	});
+
+	describe('getServicesLockedByAppId', () => {
+		const validPaths = [
+			'/tmp/123/one/updates.lock',
+			'/tmp/123/two/updates.lock',
+			'/tmp/123/three/updates.lock',
+			'/tmp/balena-supervisor/services/456/server/updates.lock',
+			'/tmp/balena-supervisor/services/456/client/updates.lock',
+			'/tmp/balena-supervisor/services/789/main/resin-updates.lock',
+		];
+		const invalidPaths = [
+			'/tmp/balena-supervisor/services/456/updates.lock',
+			'/tmp/balena-supervisor/services/server/updates.lock',
+			'/tmp/test/updates.lock',
+		];
+		let tFs: TestFs.Enabled;
+		beforeEach(async () => {
+			tFs = await testfs({ '/tmp': {} }).enable();
+			// TODO: mocha-pod should support empty directories
+			await Promise.all(
+				validPaths
+					.concat(invalidPaths)
+					.map((p) => fs.mkdir(path.dirname(p), { recursive: true })),
+			);
+		});
+		afterEach(async () => {
+			await Promise.all(
+				validPaths
+					.concat(invalidPaths)
+					.map((p) => fs.rm(path.dirname(p), { recursive: true })),
+			);
+			await tFs.restore();
+		});
+
+		it('should return locks taken by appId', async () => {
+			// Set up lockfiles
+			await Promise.all(
+				validPaths.map((p) => lockfile.lock(p, updateLock.LOCKFILE_UID)),
+			);
+
+			const locksTakenMap = updateLock.getServicesLockedByAppId();
+			expect([...locksTakenMap.keys()]).to.deep.include.members([
+				123, 456, 789,
+			]);
+			// Should register as locked if only `updates.lock` is present
+			expect(locksTakenMap.getServices(123)).to.deep.include.members([
+				'one',
+				'two',
+				'three',
+			]);
+			expect(locksTakenMap.getServices(456)).to.deep.include.members([
+				'server',
+				'client',
+			]);
+			// Should register as locked if only `resin-updates.lock` is present
+			expect(locksTakenMap.getServices(789)).to.deep.include.members(['main']);
+
+			// Cleanup lockfiles
+			await Promise.all(validPaths.map((p) => lockfile.unlock(p)));
+		});
+
+		it('should ignore invalid lockfile locations', async () => {
+			// Set up lockfiles
+			await Promise.all(invalidPaths.map((p) => lockfile.lock(p)));
+
+			expect(updateLock.getServicesLockedByAppId().size).to.equal(0);
+
+			// Cleanup lockfiles
+			await Promise.all(invalidPaths.map((p) => lockfile.unlock(p)));
+		});
+	});
 });

test/unit/lib/update-lock.spec.ts

@@ -14,4 +14,45 @@ describe('lib/update-lock: unit tests', () => {
 			);
 		});
 	});
+
+	describe('LocksTakenMap', () => {
+		it('should be an instance of Map<number, Set<string>>', () => {
+			const map = new updateLock.LocksTakenMap();
+			expect(map).to.be.an.instanceof(Map);
+		});
+
+		it('should add services while ignoring duplicates', () => {
+			const map = new updateLock.LocksTakenMap();
+			map.add(123, 'main');
+			expect(map.getServices(123)).to.deep.include.members(['main']);
+
+			map.add(123, 'main');
+			expect(map.getServices(123)).to.deep.include.members(['main']);
+
+			map.add(123, ['main', 'aux']);
+			expect(map.getServices(123)).to.deep.include.members(['main', 'aux']);
+		});
+
+		it('should track any number of appIds', () => {
+			const map = new updateLock.LocksTakenMap();
+			map.add(123, 'main');
+			map.add(456, ['aux', 'dep']);
+			expect(map.getServices(123)).to.deep.include.members(['main']);
+			expect(map.getServices(456)).to.deep.include.members(['aux', 'dep']);
+			expect(map.size).to.equal(2);
+		});
+
+		it('should return empty array for non-existent appIds', () => {
+			const map = new updateLock.LocksTakenMap();
+			expect(map.getServices(123)).to.deep.equal([]);
+		});
+
+		it('should return whether a service is locked under an appId', () => {
+			const map = new updateLock.LocksTakenMap();
+			map.add(123, 'main');
+			expect(map.isLocked(123, 'main')).to.be.true;
+			expect(map.isLocked(123, 'aux')).to.be.false;
+			expect(map.isLocked(456, 'main')).to.be.false;
+		});
+	});
 });