state service - harden the save code to try to prevent issues

src/vs/platform/lifecycle/electron-main/lifecycleMainService.ts

@@ -287,13 +287,7 @@ export class LifecycleMainService extends Disposable implements ILifecycleMainSe
 
 		this.logService.trace('Lifecycle#onWillShutdown.fire()');
 
-		const joiners: Promise<void>[] = [
-			// Always make sure the state service
-			// is flushed. Since we depend on it
-			// we cannot really use the lifecycle
-			// service from the state service...
-			this.stateMainService.flush()
-		];
+		const joiners: Promise<void>[] = [];
 
 		this._onWillShutdown.fire({
 			join(promise) {
@@ -301,7 +295,23 @@ export class LifecycleMainService extends Disposable implements ILifecycleMainSe
 			}
 		});
 
-		this.pendingWillShutdownPromise = Promises.settled(joiners).then(() => undefined, err => this.logService.error(err));
+		this.pendingWillShutdownPromise = (async () => {
+
+			// Settle all shutdown event joiners
+			try {
+				await Promises.settled(joiners);
+			} catch (error) {
+				this.logService.error(error);
+			}
+
+			// Then, always make sure at the end
+			// the state service is flushed.
+			try {
+				await this.stateMainService.close();
+			} catch (error) {
+				this.logService.error(error);
+			}
+		})();
 
 		return this.pendingWillShutdownPromise;
 	}

src/vs/platform/state/electron-main/state.ts

@@ -19,5 +19,5 @@ export interface IStateMainService {
 
 	removeItem(key: string): void;
 
-	flush(): Promise<void>;
+	close(): Promise<void>;
 }

src/vs/platform/state/electron-main/stateMainService.ts

@@ -3,9 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { join } from 'vs/base/common/path';
 import { promises } from 'fs';
 import { writeFile } from 'vs/base/node/pfs';
+import { join } from 'vs/base/common/path';
 import { isUndefined, isUndefinedOrNull } from 'vs/base/common/types';
 import { IStateMainService } from 'vs/platform/state/electron-main/state';
 import { ILogService } from 'vs/platform/log/common/log';
@@ -20,6 +20,7 @@ export class FileStorage {
 	private lastFlushedSerializedDatabase: string | undefined = undefined;
 
 	private readonly flushDelayer = new ThrottledDelayer<void>(100 /* buffer saves over a short time */);
+	private closed = false;
 
 	constructor(
 		private readonly dbPath: string,
@@ -112,7 +113,11 @@ export class FileStorage {
 		}
 	}
 
-	private save(delay?: number): Promise<void> {
+	private async save(delay?: number): Promise<void> {
+		if (this.closed) {
+			return; // already closed
+		}
+
 		return this.flushDelayer.trigger(() => this.doSave(), delay);
 	}
 
@@ -130,8 +135,14 @@ export class FileStorage {
 		}
 	}
 
-	flush(): Promise<void> {
-		return this.save(0 /* as soon as possible */);
+	async close(): Promise<void> {
+		if (this.closed) {
+			return; // already closed
+		}
+
+		this.closed = true;
+
+		return this.flushDelayer.trigger(() => this.doSave(), 0 /* as soon as possible */);
 	}
 }
 
@@ -193,9 +204,9 @@ export class StateMainService implements IStateMainService {
 		this.fileStorage.removeItem(key);
 	}
 
-	flush(): Promise<void> {
+	close(): Promise<void> {
 		this.checkInit();
 
-		return this.fileStorage.flush();
+		return this.fileStorage.close();
 	}
 }

src/vs/platform/state/test/electron-main/state.test.ts

@@ -5,7 +5,7 @@
 
 import * as assert from 'assert';
 import { tmpdir } from 'os';
-import { promises } from 'fs';
+import { promises, readFileSync } from 'fs';
 import { join } from 'vs/base/common/path';
 import { flakySuite, getRandomTestPath } from 'vs/base/test/node/testUtils';
 import { FileStorage } from 'vs/platform/state/electron-main/stateMainService';
@@ -43,7 +43,7 @@ flakySuite('StateMainService', () => {
 
 		service.setItem('some.other.key', 'some.other.value');
 
-		await service.flush();
+		await service.close();
 
 		service = new FileStorage(storageFile, new NullLogService());
 		await service.init();
@@ -106,7 +106,7 @@ flakySuite('StateMainService', () => {
 		assert.strictEqual(service.getItem('some.key3'), 'some.value3');
 		assert.strictEqual(service.getItem('some.key4'), undefined);
 
-		await service.flush();
+		await service.close();
 
 		service = new FileStorage(storageFile, new NullLogService());
 		await service.init();
@@ -141,4 +141,28 @@ flakySuite('StateMainService', () => {
 		assert.strictEqual(service.getItem('some.key3'), undefined);
 		assert.strictEqual(service.getItem('some.key4'), undefined);
 	});
+
+	test('Used after close', async function () {
+		const storageFile = join(testDir, 'storage.json');
+		writeFileSync(storageFile, '');
+
+		const service = new FileStorage(storageFile, new NullLogService());
+
+		await service.init();
+
+		service.setItem('some.key1', 'some.value1');
+		service.setItem('some.key2', 'some.value2');
+		service.setItem('some.key3', 'some.value3');
+		service.setItem('some.key4', 'some.value4');
+
+		await service.close();
+
+		service.setItem('some.key5', 'some.marker');
+
+		const contents = readFileSync(storageFile).toString();
+		assert.ok(contents.includes('some.value1'));
+		assert.ok(!contents.includes('some.marker'));
+
+		await service.close();
+	});
 });