Ensure local mode switch runs before target state

This change makes DeviceState to wait until local mode switch is definitely
completed before actually applying the state, which avoids races in state cleanup.

Change-type: patch
Signed-off-by: Roman Mazur <roman@balena.io>

src/application-manager.coffee

@@ -983,3 +983,6 @@ module.exports = class ApplicationManager extends EventEmitter
 				return _.find app.services, (svc) ->
 					svc.serviceId == serviceId
 		.get('serviceName')
+
+
+	localModeSwitchCompletion: => @localModeManager.switchCompletion()

src/device-state.coffee

@@ -598,6 +598,8 @@ module.exports = class DeviceState extends EventEmitter
 		Promise.try =>
 			if !intermediate
 				@applyBlocker
+		.then =>
+			@applications.localModeSwitchCompletion()
 		.then =>
 			@usingInferStepsLock =>
 				@applications.getExtraStateForComparison()

src/local-mode.ts

@@ -1,6 +1,7 @@
 import * as Bluebird from 'bluebird';
 import * as Docker from 'dockerode';
 import * as _ from 'lodash';
+
 import Config from './config';
 import Database from './db';
 import log from './lib/supervisor-console';
@@ -39,10 +40,10 @@ export class EngineSnapshot {
 
 	public toString(): string {
 		return (
-			`containers [${this.containers}], ` +
-			`images [${this.images}], ` +
-			`volumes [${this.volumes}], ` +
-			`networks [${this.networks}]`
+			`${this.containers.length} containers, ` +
+			`${this.images.length} images, ` +
+			`${this.volumes.length} volumes, ` +
+			`${this.networks.length} networks`
 		);
 	}
 }
@@ -71,6 +72,9 @@ export class LocalModeManager {
 		public db: Database,
 	) {}
 
+	// Indicates that switch from or to the local mode is not complete.
+	private switchInProgress: Bluebird<void> | null = null;
+
 	public async init() {
 		// Setup a listener to catch state changes relating to local mode
 		this.config.on('change', changed => {
@@ -80,7 +84,7 @@ export class LocalModeManager {
 				// First switch the logger to it's correct state
 				this.logger.switchBackend(local);
 
-				this.handleLocalModeStateChange(local);
+				this.startLocalModeChangeHandling(local);
 			}
 		});
 
@@ -103,6 +107,14 @@ export class LocalModeManager {
 		}
 	}
 
+	public startLocalModeChangeHandling(local: boolean) {
+		this.switchInProgress = Bluebird.resolve(
+			this.handleLocalModeStateChange(local),
+		).finally(() => {
+			this.switchInProgress = null;
+		});
+	}
+
 	// Query the engine to get currently running containers and installed images.
 	public async collectEngineSnapshot(): Promise<EngineSnapshotRecord> {
 		const containersPromise = this.docker
@@ -118,26 +130,30 @@ export class LocalModeManager {
 			.listNetworks()
 			.then(resp => _.map(resp, 'Id'));
 
-		const data = await Bluebird.all([
+		const [containers, images, volumes, networks] = await Bluebird.all([
 			containersPromise,
 			imagesPromise,
 			volumesPromise,
 			networksPromise,
 		]);
 		return new EngineSnapshotRecord(
-			new EngineSnapshot(data[0], data[1], data[2], data[3]),
+			new EngineSnapshot(containers, images, volumes, networks),
 			new Date(),
 		);
 	}
 
+	private async cleanEngineSnapshots() {
+		await this.db.models('engineSnapshot').delete();
+	}
+
 	// Store engine snapshot data in the local database.
 	public async storeEngineSnapshot(record: EngineSnapshotRecord) {
 		const timestamp = record.timestamp.toISOString();
 		log.debug(
 			`Storing engine snapshot in the database. Timestamp: ${timestamp}`,
 		);
-		await this.db.models('engineSnapshot').delete();
-		return this.db.models('engineSnapshot').insert({
+		await this.cleanEngineSnapshots();
+		await this.db.models('engineSnapshot').insert({
 			snapshot: JSON.stringify(record.snapshot),
 			timestamp,
 		});
@@ -210,23 +226,34 @@ export class LocalModeManager {
 
 			const previousRecord = await this.retrieveLatestSnapshot();
 			if (!previousRecord) {
-				log.warn('Previous engine snapshot was not stored. Skipping clanup.');
+				log.info('Previous engine snapshot was not stored. Skipping cleanup.');
 				return;
 			}
 
 			log.debug(
 				`Leaving local mode and cleaning up objects since ${previousRecord.timestamp.toISOString()}`,
 			);
-			return await this.removeLocalModeArtifacts(
+			await this.removeLocalModeArtifacts(
 				currentRecord.snapshot.diff(previousRecord.snapshot),
 			);
+			await this.cleanEngineSnapshots();
 		} catch (e) {
 			log.error(
 				`Problems managing engine state on local mode switch. Local mode: ${local}.`,
 				e,
 			);
+		} finally {
+			log.debug('Handling of local mode switch is completed');
 		}
 	}
+
+	// Returns a promise to await local mode switch completion started previously.
+	public async switchCompletion() {
+		if (this.switchInProgress == null) {
+			return;
+		}
+		await this.switchInProgress;
+	}
 }
 
 export default LocalModeManager;

test/local-mode.ts

@@ -2,6 +2,7 @@ import { expect } from 'chai';
 import * as Docker from 'dockerode';
 import * as sinon from 'sinon';
 import * as tmp from 'tmp';
+
 import Config from '../src/config';
 import DB from '../src/db';
 import log from '../src/lib/supervisor-console';
@@ -17,6 +18,13 @@ describe('LocalModeManager', () => {
 	let localMode: LocalModeManager;
 	let dockerStub: sinon.SinonStubbedInstance<Docker>;
 
+	const recordsCount = async () =>
+		await db
+			.models('engineSnapshot')
+			.count('* as cnt')
+			.first()
+			.then(r => r.cnt);
+
 	// Cleanup the database (to make sure nothing is left since last tests).
 	beforeEach(async () => {
 		await db.models('engineSnapshot').delete();
@@ -236,17 +244,43 @@ describe('LocalModeManager', () => {
 				expect(dockerStub.getNetwork.notCalled).to.be.true;
 				removeStubs.forEach(s => expect(s.remove.notCalled).to.be.true);
 			});
+
+			it('can be awaited', async () => {
+				const removeStubs = stubRemoveMethods(false);
+				await storeCurrentSnapshot([], [], [], []);
+
+				// Run asynchronously (like on config change).
+				localMode.startLocalModeChangeHandling(false);
+
+				// Await like it's done by DeviceState.
+				await localMode.switchCompletion();
+
+				removeStubs.forEach(s => expect(s.remove.calledTwice).to.be.true);
+			});
+
+			it('cleans the last snapshot so that nothing is done on restart', async () => {
+				const removeStubs = stubRemoveMethods(false);
+				await storeCurrentSnapshot([], [], [], []);
+
+				await localMode.handleLocalModeStateChange(false);
+
+				// The database should be empty now.
+				expect(await recordsCount()).to.be.equal(0);
+
+				// This has to be no ops.
+				await localMode.handleLocalModeStateChange(false);
+				// So our stubs must be called only once.
+				// We delete 2 objects of each type during the first call, so number of getXXX and remove calls is 2.
+				expect(dockerStub.getImage.callCount).to.be.equal(2);
+				expect(dockerStub.getContainer.callCount).to.be.equal(2);
+				expect(dockerStub.getVolume.callCount).to.be.equal(2);
+				expect(dockerStub.getNetwork.callCount).to.be.equal(2);
+				removeStubs.forEach(s => expect(s.remove.callCount).to.be.equal(2));
+			});
 		});
 	});
 
 	describe('engine snapshot storage', () => {
-		const recordsCount = async () =>
-			await db
-				.models('engineSnapshot')
-				.count('* as cnt')
-				.first()
-				.then(r => r.cnt);
-
 		const recordSample = new EngineSnapshotRecord(
 			new EngineSnapshot(
 				['c1', 'c2'],