Skip to content

Commit ebdb4d0

Browse files
Kartik Rajkarthiknadig
Kartik Raj
authored and
karthiknadig
committed
Ensure database storage extension uses to track all storages does not 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
1 parent b16f7fe commit ebdb4d0

4 files changed

Lines changed: 113 additions & 29 deletions

File tree

‎news/2 Fixes/17488.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Ensure database storage extension uses to track all storages does not grow unnecessarily.

‎src/client/common/persistentState.ts‎

Lines changed: 63 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { Memento } from 'vscode';
88
import { IExtensionSingleActivationService } from '../activation/types';
99
import { ICommandManager } from './application/types';
1010
import { Commands } from './constants';
11+
import { traceError, traceVerbose } from './logger';
1112
import {
1213
GLOBAL_MEMENTO,
1314
IExtensionContext,
@@ -16,6 +17,7 @@ import {
1617
IPersistentStateFactory,
1718
WORKSPACE_MEMENTO,
1819
} from './types';
20+
import { cache } from './utils/decorators';
1921

2022
export class PersistentState<T> implements IPersistentState<T> {
2123
constructor(
@@ -39,30 +41,39 @@ export class PersistentState<T> implements IPersistentState<T> {
3941
}
4042

4143
public async updateValue(newValue: T): Promise<void> {
42-
if (this.expiryDurationMs) {
43-
await this.storage.update(this.key, { data: newValue, expiry: Date.now() + this.expiryDurationMs });
44-
} else {
45-
await this.storage.update(this.key, newValue);
44+
try {
45+
if (this.expiryDurationMs) {
46+
await this.storage.update(this.key, { data: newValue, expiry: Date.now() + this.expiryDurationMs });
47+
} else {
48+
await this.storage.update(this.key, newValue);
49+
}
50+
} catch (ex) {
51+
traceError('Error while updating storage for key:', this.key, ex);
4652
}
4753
}
4854
}
4955

50-
const GLOBAL_PERSISTENT_KEYS = 'PYTHON_EXTENSION_GLOBAL_STORAGE_KEYS';
51-
const WORKSPACE_PERSISTENT_KEYS = 'PYTHON_EXTENSION_WORKSPACE_STORAGE_KEYS';
52-
type keysStorage = { key: string; defaultValue: unknown };
56+
const GLOBAL_PERSISTENT_KEYS_DEPRECATED = 'PYTHON_EXTENSION_GLOBAL_STORAGE_KEYS';
57+
const WORKSPACE_PERSISTENT_KEYS_DEPRECATED = 'PYTHON_EXTENSION_WORKSPACE_STORAGE_KEYS';
58+
59+
const GLOBAL_PERSISTENT_KEYS = 'PYTHON_GLOBAL_STORAGE_KEYS';
60+
const WORKSPACE_PERSISTENT_KEYS = 'PYTHON_WORKSPACE_STORAGE_KEYS';
61+
type KeysStorageType = 'global' | 'workspace';
62+
type KeysStorage = { key: string; defaultValue: unknown };
5363

5464
@injectable()
5565
export class PersistentStateFactory implements IPersistentStateFactory, IExtensionSingleActivationService {
56-
private readonly globalKeysStorage = new PersistentState<keysStorage[]>(
66+
public readonly _globalKeysStorage = new PersistentState<KeysStorage[]>(
5767
this.globalState,
5868
GLOBAL_PERSISTENT_KEYS,
5969
[],
6070
);
61-
private readonly workspaceKeysStorage = new PersistentState<keysStorage[]>(
71+
public readonly _workspaceKeysStorage = new PersistentState<KeysStorage[]>(
6272
this.workspaceState,
6373
WORKSPACE_PERSISTENT_KEYS,
6474
[],
6575
);
76+
private cleanedOnce = false;
6677
constructor(
6778
@inject(IMemento) @named(GLOBAL_MEMENTO) private globalState: Memento,
6879
@inject(IMemento) @named(WORKSPACE_MEMENTO) private workspaceState: Memento,
@@ -71,16 +82,27 @@ export class PersistentStateFactory implements IPersistentStateFactory, IExtensi
7182

7283
public async activate(): Promise<void> {
7384
this.cmdManager.registerCommand(Commands.ClearStorage, this.cleanAllPersistentStates.bind(this));
85+
const globalKeysStorageDeprecated = this.createGlobalPersistentState(GLOBAL_PERSISTENT_KEYS_DEPRECATED, []);
86+
const workspaceKeysStorageDeprecated = this.createWorkspacePersistentState(
87+
WORKSPACE_PERSISTENT_KEYS_DEPRECATED,
88+
[],
89+
);
90+
// Old storages have grown to be unusually large due to https://github.com/microsoft/vscode-python/issues/17488,
91+
// so reset them. This line can be removed after a while.
92+
if (globalKeysStorageDeprecated.value.length > 0) {
93+
globalKeysStorageDeprecated.updateValue([]).ignoreErrors();
94+
}
95+
if (workspaceKeysStorageDeprecated.value.length > 0) {
96+
workspaceKeysStorageDeprecated.updateValue([]).ignoreErrors();
97+
}
7498
}
7599

76100
public createGlobalPersistentState<T>(
77101
key: string,
78102
defaultValue?: T,
79103
expiryDurationMs?: number,
80104
): IPersistentState<T> {
81-
if (!this.globalKeysStorage.value.includes({ key, defaultValue })) {
82-
this.globalKeysStorage.updateValue([{ key, defaultValue }, ...this.globalKeysStorage.value]).ignoreErrors();
83-
}
105+
this.addKeyToStorage('global', key, defaultValue).ignoreErrors();
84106
return new PersistentState<T>(this.globalState, key, defaultValue, expiryDurationMs);
85107
}
86108

@@ -89,29 +111,44 @@ export class PersistentStateFactory implements IPersistentStateFactory, IExtensi
89111
defaultValue?: T,
90112
expiryDurationMs?: number,
91113
): IPersistentState<T> {
92-
if (!this.workspaceKeysStorage.value.includes({ key, defaultValue })) {
93-
this.workspaceKeysStorage
94-
.updateValue([{ key, defaultValue }, ...this.workspaceKeysStorage.value])
95-
.ignoreErrors();
96-
}
114+
this.addKeyToStorage('workspace', key, defaultValue).ignoreErrors();
97115
return new PersistentState<T>(this.workspaceState, key, defaultValue, expiryDurationMs);
98116
}
99117

118+
/**
119+
* Note we use a decorator to cache the promise returned by this method, so it's only called once.
120+
* It is only cached for the particular arguments passed, so the argument type is simplified here.
121+
*/
122+
@cache(-1, true)
123+
private async addKeyToStorage<T>(keyStorageType: KeysStorageType, key: string, defaultValue?: T) {
124+
const storage = keyStorageType === 'global' ? this._globalKeysStorage : this._workspaceKeysStorage;
125+
const found = storage.value.find((value) => value.key === key && value.defaultValue === defaultValue);
126+
if (!found) {
127+
await storage.updateValue([{ key, defaultValue }, ...storage.value]);
128+
}
129+
}
130+
100131
private async cleanAllPersistentStates(): Promise<void> {
132+
if (this.cleanedOnce) {
133+
traceError('Storage can only be cleaned once per session, reload window.');
134+
return;
135+
}
101136
await Promise.all(
102-
this.globalKeysStorage.value.map(async (keyContent) => {
137+
this._globalKeysStorage.value.map(async (keyContent) => {
103138
const storage = this.createGlobalPersistentState(keyContent.key);
104139
await storage.updateValue(keyContent.defaultValue);
105140
}),
106141
);
107142
await Promise.all(
108-
this.workspaceKeysStorage.value.map(async (keyContent) => {
143+
this._workspaceKeysStorage.value.map(async (keyContent) => {
109144
const storage = this.createWorkspacePersistentState(keyContent.key);
110145
await storage.updateValue(keyContent.defaultValue);
111146
}),
112147
);
113-
await this.globalKeysStorage.updateValue([]);
114-
await this.workspaceKeysStorage.updateValue([]);
148+
await this._globalKeysStorage.updateValue([]);
149+
await this._workspaceKeysStorage.updateValue([]);
150+
this.cleanedOnce = true;
151+
traceVerbose('Finished clearing storage.');
115152
}
116153
}
117154

@@ -128,9 +165,11 @@ interface IPersistentStorage<T> {
128165
* Build a global storage object for the given key.
129166
*/
130167
export function getGlobalStorage<T>(context: IExtensionContext, key: string, defaultValue?: T): IPersistentStorage<T> {
131-
const globalKeysStorage = new PersistentState<keysStorage[]>(context.globalState, GLOBAL_PERSISTENT_KEYS, []);
132-
if (!globalKeysStorage.value.includes({ key, defaultValue: undefined })) {
133-
globalKeysStorage.updateValue([{ key, defaultValue: undefined }, ...globalKeysStorage.value]).ignoreErrors();
168+
const globalKeysStorage = new PersistentState<KeysStorage[]>(context.globalState, GLOBAL_PERSISTENT_KEYS, []);
169+
const found = globalKeysStorage.value.find((value) => value.key === key && value.defaultValue === defaultValue);
170+
if (!found) {
171+
const newValue = [{ key, defaultValue }, ...globalKeysStorage.value];
172+
globalKeysStorage.updateValue(newValue).ignoreErrors();
134173
}
135174
const raw = new PersistentState<T>(context.globalState, key, defaultValue);
136175
return {

‎src/client/common/utils/decorators.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,13 @@ export function cache(expiryDurationMs: number, cachePromise = false, expiryDura
153153
if (isTestExecution()) {
154154
return originalMethod.apply(this, args) as Promise<any>;
155155
}
156-
const key = getCacheKeyFromFunctionArgs(keyPrefix, args);
156+
let key: string;
157+
try {
158+
key = getCacheKeyFromFunctionArgs(keyPrefix, args);
159+
} catch (ex) {
160+
traceError('Error while creating key for keyPrefix:', keyPrefix, ex);
161+
return originalMethod.apply(this, args) as Promise<any>;
162+
}
157163
const cachedItem = cacheStoreForMethods.get(key);
158164
if (cachedItem && (cachedItem.expiry > Date.now() || expiryDurationMs === -1)) {
159165
traceVerbose(`Cached data exists ${key}`);

‎src/test/common/persistentState.unit.test.ts‎

Lines changed: 42 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,19 +3,19 @@
33

44
'use strict';
55

6-
import { expect } from 'chai';
6+
import { assert, expect } from 'chai';
77
import * as TypeMoq from 'typemoq';
88
import { Memento } from 'vscode';
9-
import { IExtensionSingleActivationService } from '../../client/activation/types';
109
import { ICommandManager } from '../../client/common/application/types';
1110
import { Commands } from '../../client/common/constants';
1211
import { PersistentStateFactory } from '../../client/common/persistentState';
13-
import { IDisposable, IPersistentStateFactory } from '../../client/common/types';
12+
import { IDisposable } from '../../client/common/types';
13+
import { sleep } from '../core';
1414
import { MockMemento } from '../mocks/mementos';
1515

1616
suite('Persistent State', () => {
1717
let cmdManager: TypeMoq.IMock<ICommandManager>;
18-
let persistentStateFactory: IPersistentStateFactory & IExtensionSingleActivationService;
18+
let persistentStateFactory: PersistentStateFactory;
1919
let workspaceMemento: Memento;
2020
let globalMemento: Memento;
2121
setup(() => {
@@ -87,4 +87,42 @@ suite('Persistent State', () => {
8787
expect(workspaceKey1State.value).to.equal(undefined);
8888
expect(workspaceKey2State.value).to.equal('defaultKey2Value');
8989
});
90+
91+
test('Ensure internal global storage extension uses to track other storages does not contain duplicate entries', async () => {
92+
persistentStateFactory.createGlobalPersistentState('key1');
93+
await sleep(1);
94+
persistentStateFactory.createGlobalPersistentState('key2', 'defaultValue1');
95+
await sleep(1);
96+
persistentStateFactory.createGlobalPersistentState('key2', 'defaultValue1');
97+
await sleep(1);
98+
persistentStateFactory.createGlobalPersistentState('key1');
99+
await sleep(1);
100+
const { value } = persistentStateFactory._globalKeysStorage;
101+
assert.deepEqual(
102+
value.sort((k1, k2) => k1.key.localeCompare(k2.key)),
103+
[
104+
{ key: 'key1', defaultValue: undefined },
105+
{ key: 'key2', defaultValue: 'defaultValue1' },
106+
].sort((k1, k2) => k1.key.localeCompare(k2.key)),
107+
);
108+
});
109+
110+
test('Ensure internal workspace storage extension uses to track other storages does not contain duplicate entries', async () => {
111+
persistentStateFactory.createWorkspacePersistentState('key2', 'defaultValue1');
112+
await sleep(1);
113+
persistentStateFactory.createWorkspacePersistentState('key1');
114+
await sleep(1);
115+
persistentStateFactory.createWorkspacePersistentState('key2', 'defaultValue1');
116+
await sleep(1);
117+
persistentStateFactory.createWorkspacePersistentState('key1');
118+
await sleep(1);
119+
const { value } = persistentStateFactory._workspaceKeysStorage;
120+
assert.deepEqual(
121+
value.sort((k1, k2) => k1.key.localeCompare(k2.key)),
122+
[
123+
{ key: 'key1', defaultValue: undefined },
124+
{ key: 'key2', defaultValue: 'defaultValue1' },
125+
].sort((k1, k2) => k1.key.localeCompare(k2.key)),
126+
);
127+
});
90128
});

0 commit comments

Comments
 (0)