Merge pull request #1953 from balena-os/shutdown-with-locks

Shutdown with locks

src/compose/composition-steps.ts

@@ -120,7 +120,7 @@ type Executors<T extends CompositionStepAction> = {
 type LockingFn = (
 	// TODO: Once the entire codebase is typescript, change
 	// this to number
-	app: number | null,
+	app: number | number[] | null,
 	args: BaseCompositionStepArgs,
 	fn: () => Promise<unknown>,
 ) => Promise<unknown>;

src/device-state.ts

@@ -684,24 +684,36 @@ export function reportCurrentState(newState: DeviceReport = {}) {
 	emitAsync('change', undefined);
 }
 
-export async function reboot(force?: boolean, skipLock?: boolean) {
-	await updateLock.abortIfHUPInProgress({ force });
-	await applicationManager.stopAll({ force, skipLock });
-	logger.logSystemMessage('Rebooting', {}, 'Reboot');
-	const $reboot = await dbus.reboot();
-	shuttingDown = true;
-	emitAsync('shutdown', undefined);
-	return await $reboot;
+export interface ShutdownOpts {
+	force?: boolean;
+	reboot?: boolean;
 }
 
-export async function shutdown(force?: boolean, skipLock?: boolean) {
+export async function shutdown({
+	force = false,
+	reboot = false,
+}: ShutdownOpts = {}) {
 	await updateLock.abortIfHUPInProgress({ force });
-	await applicationManager.stopAll({ force, skipLock });
-	logger.logSystemMessage('Shutting down', {}, 'Shutdown');
-	const $shutdown = await dbus.shutdown();
-	shuttingDown = true;
-	emitAsync('shutdown', undefined);
-	return $shutdown;
+	// Get current apps to create locks for
+	const apps = await applicationManager.getCurrentApps();
+	const appIds = Object.keys(apps).map((strId) => parseInt(strId, 10));
+	// Try to create a lock for all the services before shutting down
+	return updateLock.lock(appIds, { force }, async () => {
+		let dbusAction;
+		switch (reboot) {
+			case true:
+				logger.logSystemMessage('Rebooting', {}, 'Reboot');
+				dbusAction = await dbus.reboot();
+				break;
+			case false:
+				logger.logSystemMessage('Shutting down', {}, 'Shutdown');
+				dbusAction = await dbus.shutdown();
+				break;
+		}
+		shuttingDown = true;
+		emitAsync('shutdown', undefined);
+		return dbusAction;
+	});
 }
 
 export async function executeStepAction<T extends PossibleStepTargets>(
@@ -728,13 +740,13 @@ export async function executeStepAction<T extends PossibleStepTargets>(
 				// and if they do, we wouldn't know about it until after
 				// the response has been sent back to the API. Just return
 				// "OK" for this and the below action
-				await reboot(force, skipLock);
+				await shutdown({ force, reboot: true });
 				return {
 					Data: 'OK',
 					Error: null,
 				};
 			case 'shutdown':
-				await shutdown(force, skipLock);
+				await shutdown({ force, reboot: false });
 				return {
 					Data: 'OK',
 					Error: null,

src/lib/update-lock.ts

@@ -71,20 +71,20 @@ export const readLock: LockFn = Bluebird.promisify(locker.async.readLock, {
 });
 
 // Unlock all lockfiles, optionally of an appId | appUuid, then release resources.
-function dispose(
-	release: () => void,
+async function dispose(
 	appIdentifier: string | number,
-): Bluebird<void> {
-	return Bluebird.map(
-		lockfile.getLocksTaken((p: string) =>
-			p.includes(`${BASE_LOCK_DIR}/${appIdentifier}`),
-		),
-		(lockName) => {
-			return lockfile.unlock(lockName);
-		},
-	)
-		.finally(release)
-		.return();
+	release: () => void,
+): Promise<void> {
+	const locks = lockfile.getLocksTaken((p: string) =>
+		p.includes(`${BASE_LOCK_DIR}/${appIdentifier}`),
+	);
+	try {
+		// Try to unlock all locks taken
+		await Promise.all(locks.map((l) => lockfile.unlock(l)));
+	} finally {
+		// Release final resource
+		release();
+	}
 }
 
 /**
@@ -96,69 +96,77 @@ function dispose(
  * TODO: Remove skipLock as it's not a good interface. If lock is called it should try to take the lock
  * without an option to skip.
  */
-export function lock<T extends unknown>(
-	appId: number,
+export async function lock<T extends unknown>(
+	appId: number | number[],
 	{ force = false, skipLock = false }: { force: boolean; skipLock?: boolean },
 	fn: () => Resolvable<T>,
-): Bluebird<T> {
-	if (skipLock || appId == null) {
-		return Bluebird.resolve(fn());
+): Promise<T> {
+	const appIdsToLock = Array.isArray(appId) ? appId : [appId];
+	if (skipLock || !appId || !appIdsToLock.length) {
+		return fn();
 	}
 
-	const takeTheLock = () => {
-		return config
-			.get('lockOverride')
-			.then((lockOverride) => {
-				return writeLock(appId)
-					.tap((release: () => void) => {
-						const lockDir = getPathOnHost(lockPath(appId));
-						return Bluebird.resolve(fs.readdir(lockDir))
-							.catchReturn(ENOENT, [])
-							.mapSeries((serviceName) => {
-								return Bluebird.mapSeries(
-									lockFilesOnHost(appId, serviceName),
-									(tmpLockName) => {
-										return (
-											Bluebird.try(() => {
-												if (force || lockOverride) {
-													return lockfile.unlock(tmpLockName);
-												}
-											})
-												.then(() => {
-													return lockfile.lock(tmpLockName, LOCKFILE_UID);
-												})
-												// If lockfile exists, throw a user-friendly error.
-												// Otherwise throw the error as-is.
-												// This will interrupt the call to Bluebird.using, so
-												// dispose needs to be called even though it's referenced
-												// by .disposer later.
-												.catch((error) => {
-													return dispose(release, appId).throw(
-														lockfile.LockfileExistsError.is(error)
-															? new UpdatesLockedError(
-																	`Lockfile exists for ${JSON.stringify({
-																		serviceName,
-																		appId,
-																	})}`,
-															  )
-															: (error as Error),
-													);
-												})
-										);
-									},
-								);
-							});
-					})
-					.disposer((release: () => void) => dispose(release, appId));
-			})
-			.catch((err) => {
-				throw new InternalInconsistencyError(
-					`Error getting lockOverride config value: ${err?.message ?? err}`,
-				);
+	// Sort appIds so they are always locked in the same sequence
+	const sortedIds = appIdsToLock.sort();
+
+	let lockOverride: boolean;
+	try {
+		lockOverride = await config.get('lockOverride');
+	} catch (err) {
+		throw new InternalInconsistencyError(
+			`Error getting lockOverride config value: ${err?.message ?? err}`,
+		);
+	}
+
+	const releases = new Map<number, () => void>();
+	try {
+		for (const id of sortedIds) {
+			const lockDir = getPathOnHost(lockPath(id));
+			// Acquire write lock for appId
+			releases.set(id, await writeLock(id));
+			// Get list of service folders in lock directory
+			const serviceFolders = await fs.readdir(lockDir).catch((e) => {
+				if (ENOENT(e)) {
+					return [];
+				}
+				throw e;
 			});
-	};
-
-	const disposer = takeTheLock();
-
-	return Bluebird.using(disposer, fn as () => PromiseLike<T>);
+			// Attempt to create a lock for each service
+			for (const service of serviceFolders) {
+				await lockService(id, service, force || lockOverride);
+			}
+		}
+		// Resolve the function passed
+		return fn();
+	} finally {
+		for (const [id, release] of releases.entries()) {
+			// Try to dispose all the locks
+			await dispose(id, release);
+		}
+	}
+}
+
+async function lockService(
+	appId: number,
+	service: string,
+	force: boolean = false,
+): Promise<void> {
+	const serviceLockFiles = lockFilesOnHost(appId, service);
+	for await (const file of serviceLockFiles) {
+		try {
+			if (force) {
+				await lockfile.unlock(file);
+			}
+			await lockfile.lock(file, LOCKFILE_UID);
+		} catch (e) {
+			if (lockfile.LockfileExistsError.is(e)) {
+				// Throw more descriptive error
+				throw new UpdatesLockedError(
+					`Lockfile exists for { appId: ${appId}, service: ${service} }`,
+				);
+			}
+			// Otherwise just throw the error
+			throw e;
+		}
+	}
 }

test/05-device-state.spec.ts

@@ -548,7 +548,7 @@ describe('device-state', () => {
 			.stub(updateLock, 'abortIfHUPInProgress')
 			.throws(new UpdatesLockedError(testErrMsg));
 
-		await expect(deviceState.reboot())
+		await expect(deviceState.shutdown({ reboot: true }))
 			.to.eventually.be.rejectedWith(testErrMsg)
 			.and.be.an.instanceOf(UpdatesLockedError);
 		await expect(deviceState.shutdown())

test/41-device-api-v1.spec.ts

@@ -538,7 +538,54 @@ describe('SupervisorAPI [V1 Endpoints]', () => {
 			shutdownMock.resetHistory();
 		});
 
-		it('should return 423 and reject the reboot if no locks are set', async () => {
+		it('should lock all applications before trying to shutdown', async () => {
+			// Setup 2 applications running
+			const twoContainers = [
+				mockedAPI.mockService({
+					containerId: 'abc123',
+					appId: 1000,
+					releaseId: 55555,
+				}),
+				mockedAPI.mockService({
+					containerId: 'def456',
+					appId: 2000,
+					releaseId: 77777,
+				}),
+			];
+			const twoImages = [
+				mockedAPI.mockImage({
+					appId: 1000,
+				}),
+				mockedAPI.mockImage({
+					appId: 2000,
+				}),
+			];
+			appMock.mockManagers(twoContainers, [], []);
+			appMock.mockImages([], false, twoImages);
+
+			const lockSpy = spy(updateLock, 'lock');
+			await mockedDockerode.testWithData(
+				{ containers: twoContainers, images: twoImages },
+				async () => {
+					const response = await request
+						.post('/v1/shutdown')
+						.set('Accept', 'application/json')
+						.set('Authorization', `Bearer ${apiKeys.cloudApiKey}`)
+						.expect(202);
+
+					expect(lockSpy.callCount).to.equal(1);
+					// Check that lock was passed both application Ids
+					expect(lockSpy.lastCall.args[0]).to.deep.equal([1000, 2000]);
+					expect(response.body).to.have.property('Data').that.is.not.empty;
+					expect(shutdownMock).to.have.been.calledOnce;
+				},
+			);
+
+			shutdownMock.resetHistory();
+			lockSpy.restore();
+		});
+
+		it('should return 423 and reject the reboot if locks are set', async () => {
 			stub(updateLock, 'lock').callsFake((__, opts, fn) => {
 				if (opts.force) {
 					return Bluebird.resolve(fn());

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

@@ -242,6 +242,45 @@ describe('lib/update-lock', () => {
 			);
 		});
 
+		it('locks all applications before resolving input function', async () => {
+			const appIds = [111, 222, 333];
+			// Set up fake filesystem for lockfiles
+			mockFs({
+				[path.join(
+					constants.rootMountPoint,
+					updateLock.lockPath(111),
+					serviceName,
+				)]: {},
+				[path.join(
+					constants.rootMountPoint,
+					updateLock.lockPath(222),
+					serviceName,
+				)]: {},
+				[path.join(
+					constants.rootMountPoint,
+					updateLock.lockPath(333),
+					serviceName,
+				)]: {},
+			});
+
+			await expect(
+				updateLock.lock(appIds, { force: false }, async () => {
+					// At this point the locks should be taken and not removed
+					// until this function has been resolved
+					// Both `updates.lock` and `resin-updates.lock` should have been taken
+					expect(lockSpy.args).to.have.length(6);
+					// Make sure that no locks have been removed also
+					expect(unlockSpy).to.not.be.called;
+					return Promise.resolve();
+				}),
+			).to.eventually.be.fulfilled;
+
+			// Everything that was locked should have been unlocked after function resolves
+			expect(lockSpy.args.map(([lock]) => [lock])).to.deep.equal(
+				unlockSpy.args,
+			);
+		});
+
 		it('resolves input function without locking when appId is null', async () => {
 			mockLockDir({ createLockfile: true });