Skip to content

Commit

Permalink
Ensure database storage extension uses to track all storages does not…
Browse files Browse the repository at this point in the history 
… grow unnecessarily (#17598)

* Ensure central storage to track keys for other storage never contain duplicate entries

* News entry

* Add tests

* Add more test cases

* Fix storage for function
  • Loading branch information
@karthiknadig
Kartik Raj authored and karthiknadig committed Oct 5, 2021
1 parent b16f7fe commit ebdb4d0
Show file tree
Hide file tree
Showing 4 changed files with 113 additions and 29 deletions.
1 change: 1 addition & 0 deletions news/2 Fixes/17488.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Ensure database storage extension uses to track all storages does not grow unnecessarily.
87 changes: 63 additions & 24 deletions src/client/common/persistentState.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { Memento } from 'vscode';
import { IExtensionSingleActivationService } from '../activation/types';
import { ICommandManager } from './application/types';
import { Commands } from './constants';
import { traceError, traceVerbose } from './logger';
import {
GLOBAL_MEMENTO,
IExtensionContext,
Expand All @@ -16,6 +17,7 @@ import {
IPersistentStateFactory,
WORKSPACE_MEMENTO,
} from './types';
import { cache } from './utils/decorators';

export class PersistentState<T> implements IPersistentState<T> {
constructor(
Expand All @@ -39,30 +41,39 @@ export class PersistentState<T> implements IPersistentState<T> {
}

public async updateValue(newValue: T): Promise<void> {
if (this.expiryDurationMs) {
await this.storage.update(this.key, { data: newValue, expiry: Date.now() + this.expiryDurationMs });
} else {
await this.storage.update(this.key, newValue);
try {
if (this.expiryDurationMs) {
await this.storage.update(this.key, { data: newValue, expiry: Date.now() + this.expiryDurationMs });
} else {
await this.storage.update(this.key, newValue);
}
} catch (ex) {
traceError('Error while updating storage for key:', this.key, ex);
}
}
}

const GLOBAL_PERSISTENT_KEYS = 'PYTHON_EXTENSION_GLOBAL_STORAGE_KEYS';
const WORKSPACE_PERSISTENT_KEYS = 'PYTHON_EXTENSION_WORKSPACE_STORAGE_KEYS';
type keysStorage = { key: string; defaultValue: unknown };
const GLOBAL_PERSISTENT_KEYS_DEPRECATED = 'PYTHON_EXTENSION_GLOBAL_STORAGE_KEYS';
const WORKSPACE_PERSISTENT_KEYS_DEPRECATED = 'PYTHON_EXTENSION_WORKSPACE_STORAGE_KEYS';

const GLOBAL_PERSISTENT_KEYS = 'PYTHON_GLOBAL_STORAGE_KEYS';
const WORKSPACE_PERSISTENT_KEYS = 'PYTHON_WORKSPACE_STORAGE_KEYS';
type KeysStorageType = 'global' | 'workspace';
type KeysStorage = { key: string; defaultValue: unknown };

@injectable()
export class PersistentStateFactory implements IPersistentStateFactory, IExtensionSingleActivationService {
private readonly globalKeysStorage = new PersistentState<keysStorage[]>(
public readonly _globalKeysStorage = new PersistentState<KeysStorage[]>(
this.globalState,
GLOBAL_PERSISTENT_KEYS,
[],
);
private readonly workspaceKeysStorage = new PersistentState<keysStorage[]>(
public readonly _workspaceKeysStorage = new PersistentState<KeysStorage[]>(
this.workspaceState,
WORKSPACE_PERSISTENT_KEYS,
[],
);
private cleanedOnce = false;
constructor(
@inject(IMemento) @named(GLOBAL_MEMENTO) private globalState: Memento,
@inject(IMemento) @named(WORKSPACE_MEMENTO) private workspaceState: Memento,
Expand All @@ -71,16 +82,27 @@ export class PersistentStateFactory implements IPersistentStateFactory, IExtensi

public async activate(): Promise<void> {
this.cmdManager.registerCommand(Commands.ClearStorage, this.cleanAllPersistentStates.bind(this));
const globalKeysStorageDeprecated = this.createGlobalPersistentState(GLOBAL_PERSISTENT_KEYS_DEPRECATED, []);
const workspaceKeysStorageDeprecated = this.createWorkspacePersistentState(
WORKSPACE_PERSISTENT_KEYS_DEPRECATED,
[],
);
// Old storages have grown to be unusually large due to https://github.com/microsoft/vscode-python/issues/17488,
// so reset them. This line can be removed after a while.
if (globalKeysStorageDeprecated.value.length > 0) {
globalKeysStorageDeprecated.updateValue([]).ignoreErrors();
}
if (workspaceKeysStorageDeprecated.value.length > 0) {
workspaceKeysStorageDeprecated.updateValue([]).ignoreErrors();
}
}

public createGlobalPersistentState<T>(
key: string,
defaultValue?: T,
expiryDurationMs?: number,
): IPersistentState<T> {
if (!this.globalKeysStorage.value.includes({ key, defaultValue })) {
this.globalKeysStorage.updateValue([{ key, defaultValue }, ...this.globalKeysStorage.value]).ignoreErrors();
}
this.addKeyToStorage('global', key, defaultValue).ignoreErrors();
return new PersistentState<T>(this.globalState, key, defaultValue, expiryDurationMs);
}

Expand All @@ -89,29 +111,44 @@ export class PersistentStateFactory implements IPersistentStateFactory, IExtensi
defaultValue?: T,
expiryDurationMs?: number,
): IPersistentState<T> {
if (!this.workspaceKeysStorage.value.includes({ key, defaultValue })) {
this.workspaceKeysStorage
.updateValue([{ key, defaultValue }, ...this.workspaceKeysStorage.value])
.ignoreErrors();
}
this.addKeyToStorage('workspace', key, defaultValue).ignoreErrors();
return new PersistentState<T>(this.workspaceState, key, defaultValue, expiryDurationMs);
}

/**
* Note we use a decorator to cache the promise returned by this method, so it's only called once.
* It is only cached for the particular arguments passed, so the argument type is simplified here.
*/
@cache(-1, true)
private async addKeyToStorage<T>(keyStorageType: KeysStorageType, key: string, defaultValue?: T) {
const storage = keyStorageType === 'global' ? this._globalKeysStorage : this._workspaceKeysStorage;
const found = storage.value.find((value) => value.key === key && value.defaultValue === defaultValue);
if (!found) {
await storage.updateValue([{ key, defaultValue }, ...storage.value]);
}
}

private async cleanAllPersistentStates(): Promise<void> {
if (this.cleanedOnce) {
traceError('Storage can only be cleaned once per session, reload window.');
return;
}
await Promise.all(
this.globalKeysStorage.value.map(async (keyContent) => {
this._globalKeysStorage.value.map(async (keyContent) => {
const storage = this.createGlobalPersistentState(keyContent.key);
await storage.updateValue(keyContent.defaultValue);
}),
);
await Promise.all(
this.workspaceKeysStorage.value.map(async (keyContent) => {
this._workspaceKeysStorage.value.map(async (keyContent) => {
const storage = this.createWorkspacePersistentState(keyContent.key);
await storage.updateValue(keyContent.defaultValue);
}),
);
await this.globalKeysStorage.updateValue([]);
await this.workspaceKeysStorage.updateValue([]);
await this._globalKeysStorage.updateValue([]);
await this._workspaceKeysStorage.updateValue([]);
this.cleanedOnce = true;
traceVerbose('Finished clearing storage.');
}
}

Expand All @@ -128,9 +165,11 @@ interface IPersistentStorage<T> {
* Build a global storage object for the given key.
*/
export function getGlobalStorage<T>(context: IExtensionContext, key: string, defaultValue?: T): IPersistentStorage<T> {
const globalKeysStorage = new PersistentState<keysStorage[]>(context.globalState, GLOBAL_PERSISTENT_KEYS, []);
if (!globalKeysStorage.value.includes({ key, defaultValue: undefined })) {
globalKeysStorage.updateValue([{ key, defaultValue: undefined }, ...globalKeysStorage.value]).ignoreErrors();
const globalKeysStorage = new PersistentState<KeysStorage[]>(context.globalState, GLOBAL_PERSISTENT_KEYS, []);
const found = globalKeysStorage.value.find((value) => value.key === key && value.defaultValue === defaultValue);
if (!found) {
const newValue = [{ key, defaultValue }, ...globalKeysStorage.value];
globalKeysStorage.updateValue(newValue).ignoreErrors();
}
const raw = new PersistentState<T>(context.globalState, key, defaultValue);
return {
Expand Down
8 changes: 7 additions & 1 deletion src/client/common/utils/decorators.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,13 @@ export function cache(expiryDurationMs: number, cachePromise = false, expiryDura
if (isTestExecution()) {
return originalMethod.apply(this, args) as Promise<any>;
}
const key = getCacheKeyFromFunctionArgs(keyPrefix, args);
let key: string;
try {
key = getCacheKeyFromFunctionArgs(keyPrefix, args);
} catch (ex) {
traceError('Error while creating key for keyPrefix:', keyPrefix, ex);
return originalMethod.apply(this, args) as Promise<any>;
}
const cachedItem = cacheStoreForMethods.get(key);
if (cachedItem && (cachedItem.expiry > Date.now() || expiryDurationMs === -1)) {
traceVerbose(`Cached data exists ${key}`);
Expand Down
46 changes: 42 additions & 4 deletions src/test/common/persistentState.unit.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,19 @@

'use strict';

import { expect } from 'chai';
import { assert, expect } from 'chai';
import * as TypeMoq from 'typemoq';
import { Memento } from 'vscode';
import { IExtensionSingleActivationService } from '../../client/activation/types';
import { ICommandManager } from '../../client/common/application/types';
import { Commands } from '../../client/common/constants';
import { PersistentStateFactory } from '../../client/common/persistentState';
import { IDisposable, IPersistentStateFactory } from '../../client/common/types';
import { IDisposable } from '../../client/common/types';
import { sleep } from '../core';
import { MockMemento } from '../mocks/mementos';

suite('Persistent State', () => {
let cmdManager: TypeMoq.IMock<ICommandManager>;
let persistentStateFactory: IPersistentStateFactory & IExtensionSingleActivationService;
let persistentStateFactory: PersistentStateFactory;
let workspaceMemento: Memento;
let globalMemento: Memento;
setup(() => {
Expand Down Expand Up @@ -87,4 +87,42 @@ suite('Persistent State', () => {
expect(workspaceKey1State.value).to.equal(undefined);
expect(workspaceKey2State.value).to.equal('defaultKey2Value');
});

test('Ensure internal global storage extension uses to track other storages does not contain duplicate entries', async () => {
persistentStateFactory.createGlobalPersistentState('key1');
await sleep(1);
persistentStateFactory.createGlobalPersistentState('key2', 'defaultValue1');
await sleep(1);
persistentStateFactory.createGlobalPersistentState('key2', 'defaultValue1');
await sleep(1);
persistentStateFactory.createGlobalPersistentState('key1');
await sleep(1);
const { value } = persistentStateFactory._globalKeysStorage;
assert.deepEqual(
value.sort((k1, k2) => k1.key.localeCompare(k2.key)),
[
{ key: 'key1', defaultValue: undefined },
{ key: 'key2', defaultValue: 'defaultValue1' },
].sort((k1, k2) => k1.key.localeCompare(k2.key)),
);
});

test('Ensure internal workspace storage extension uses to track other storages does not contain duplicate entries', async () => {
persistentStateFactory.createWorkspacePersistentState('key2', 'defaultValue1');
await sleep(1);
persistentStateFactory.createWorkspacePersistentState('key1');
await sleep(1);
persistentStateFactory.createWorkspacePersistentState('key2', 'defaultValue1');
await sleep(1);
persistentStateFactory.createWorkspacePersistentState('key1');
await sleep(1);
const { value } = persistentStateFactory._workspaceKeysStorage;
assert.deepEqual(
value.sort((k1, k2) => k1.key.localeCompare(k2.key)),
[
{ key: 'key1', defaultValue: undefined },
{ key: 'key2', defaultValue: 'defaultValue1' },
].sort((k1, k2) => k1.key.localeCompare(k2.key)),
);
});
});

0 comments on commit ebdb4d0

Please sign in to comment.