Skip to content

Commit c5e08f9

Browse files
MrtenzMrtenz
committed
Debounce state updates
1 parent be0d97c commit c5e08f9

File tree

4 files changed

+148
-65
lines changed

4 files changed

+148
-65
lines changed

‎packages/snaps-controllers/src/snaps/SnapController.test.tsx‎

Lines changed: 52 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -9228,6 +9228,14 @@ describe('SnapController', () => {
92289228
});
92299229

92309230
describe('SnapController:getSnapState', () => {
9231+
beforeAll(() => {
9232+
jest.useFakeTimers();
9233+
});
9234+
9235+
afterAll(() => {
9236+
jest.useRealTimers();
9237+
});
9238+
92319239
it(`gets the snap's state`, async () => {
92329240
const messenger = getSnapControllerMessenger();
92339241

@@ -9316,6 +9324,7 @@ describe('SnapController', () => {
93169324
DEFAULT_ENCRYPTION_KEY_DERIVATION_OPTIONS,
93179325
);
93189326

9327+
jest.advanceTimersByTime(5_000);
93199328
await promise;
93209329

93219330
const result = await messenger.call(
@@ -9367,6 +9376,7 @@ describe('SnapController', () => {
93679376
true,
93689377
);
93699378

9379+
jest.advanceTimersByTime(5_000);
93709380
await promise;
93719381

93729382
const encryptedState1 = await encrypt(
@@ -9557,6 +9567,14 @@ describe('SnapController', () => {
95579567
});
95589568

95599569
describe('SnapController:updateSnapState', () => {
9570+
beforeAll(() => {
9571+
jest.useFakeTimers();
9572+
});
9573+
9574+
afterAll(() => {
9575+
jest.useRealTimers();
9576+
});
9577+
95609578
it(`updates the snap's state`, async () => {
95619579
const messenger = getSnapControllerMessenger();
95629580

@@ -9587,6 +9605,7 @@ describe('SnapController', () => {
95879605
true,
95889606
);
95899607

9608+
jest.advanceTimersByTime(5_000);
95909609
await promise;
95919610

95929611
expect(updateSnapStateSpy).toHaveBeenCalledTimes(1);
@@ -9611,6 +9630,8 @@ describe('SnapController', () => {
96119630

96129631
const updateSnapStateSpy = jest.spyOn(snapController, 'updateSnapState');
96139632
const state = { foo: 'bar' };
9633+
9634+
const promise = waitForStateChange(messenger);
96149635
await messenger.call(
96159636
'SnapController:updateSnapState',
96169637
MOCK_SNAP_ID,
@@ -9619,6 +9640,10 @@ describe('SnapController', () => {
96199640
);
96209641

96219642
expect(updateSnapStateSpy).toHaveBeenCalledTimes(1);
9643+
9644+
jest.advanceTimersByTime(5_000);
9645+
await promise;
9646+
96229647
expect(
96239648
snapController.state.unencryptedSnapStates[MOCK_SNAP_ID],
96249649
).toStrictEqual(JSON.stringify(state));
@@ -9657,30 +9682,19 @@ describe('SnapController', () => {
96579682
true,
96589683
);
96599684

9685+
jest.advanceTimersByTime(5_000);
96609686
await promise;
96619687

96629688
expect(hmacSha512).toHaveBeenCalledTimes(10);
96639689

96649690
snapController.destroy();
96659691
});
96669692

9667-
it('queues multiple state updates', async () => {
9693+
it('debounces multiple state updates', async () => {
96689694
const messenger = getSnapControllerMessenger();
96699695

9670-
jest.useFakeTimers();
9671-
96729696
const encryptor = getSnapControllerEncryptor();
9673-
const { promise, resolve } = createDeferredPromise();
9674-
const encryptWithKey = jest
9675-
.fn<
9676-
ReturnType<typeof encryptor.encryptWithKey>,
9677-
Parameters<typeof encryptor.encryptWithKey>
9678-
>()
9679-
.mockImplementation(async (...args) => {
9680-
resolve();
9681-
await sleep(1);
9682-
return await encryptor.encryptWithKey(...args);
9683-
});
9697+
const encryptWithKey = jest.spyOn(encryptor, 'encryptWithKey');
96849698

96859699
const snapController = getSnapController(
96869700
getSnapControllerOptions({
@@ -9696,42 +9710,36 @@ describe('SnapController', () => {
96969710
}),
96979711
);
96989712

9699-
const firstStateChange = waitForStateChange(messenger);
9713+
const promise = waitForStateChange(messenger);
97009714
await messenger.call(
97019715
'SnapController:updateSnapState',
97029716
MOCK_SNAP_ID,
97039717
{ foo: 'bar' },
97049718
true,
97059719
);
97069720

9721+
expect(
9722+
await messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
9723+
).toStrictEqual({ foo: 'bar' });
9724+
97079725
await messenger.call(
97089726
'SnapController:updateSnapState',
97099727
MOCK_SNAP_ID,
97109728
{ bar: 'baz' },
97119729
true,
97129730
);
97139731

9714-
// We await this promise to ensure the timer is queued.
9715-
await promise;
9716-
jest.advanceTimersByTime(1);
9717-
9718-
// After this point the second update should be queued.
9719-
await firstStateChange;
9720-
const secondStateChange = waitForStateChange(messenger);
9721-
9722-
expect(encryptWithKey).toHaveBeenCalledTimes(1);
9723-
9724-
// This is a bit hacky, but we can't simply advance the timer by 1ms
9725-
// because the second timer is not running yet.
9726-
jest.useRealTimers();
9727-
await secondStateChange;
9728-
9729-
expect(encryptWithKey).toHaveBeenCalledTimes(2);
9730-
97319732
expect(
97329733
await messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
97339734
).toStrictEqual({ bar: 'baz' });
97349735

9736+
expect(encryptWithKey).not.toHaveBeenCalled();
9737+
9738+
jest.advanceTimersByTime(5_000);
9739+
9740+
await promise;
9741+
expect(encryptWithKey).toHaveBeenCalledTimes(1);
9742+
97359743
snapController.destroy();
97369744
});
97379745

@@ -9763,14 +9771,24 @@ describe('SnapController', () => {
97639771
true,
97649772
);
97659773

9774+
jest.advanceTimersByTime(5_000);
97669775
await promise;
9776+
97679777
expect(error).toHaveBeenCalledWith(errorValue);
97689778

97699779
snapController.destroy();
97709780
});
97719781
});
97729782

97739783
describe('SnapController:clearSnapState', () => {
9784+
beforeAll(() => {
9785+
jest.useFakeTimers();
9786+
});
9787+
9788+
afterAll(() => {
9789+
jest.useRealTimers();
9790+
});
9791+
97749792
it('clears the state of a snap', async () => {
97759793
const messenger = getSnapControllerMessenger();
97769794

@@ -9859,7 +9877,9 @@ describe('SnapController', () => {
98599877
// eslint-disable-next-line @typescript-eslint/await-thenable
98609878
await messenger.call('SnapController:clearSnapState', MOCK_SNAP_ID, true);
98619879

9880+
jest.advanceTimersByTime(5_000);
98629881
await promise;
9882+
98639883
expect(error).toHaveBeenCalledWith(errorValue);
98649884

98659885
snapController.destroy();

‎packages/snaps-controllers/src/snaps/SnapController.ts‎

Lines changed: 30 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,7 @@ import type {
167167
KeyDerivationOptions,
168168
} from '../types';
169169
import {
170+
debounce,
170171
fetchSnap,
171172
hasTimedOut,
172173
permissionsDiff,
@@ -1966,30 +1967,35 @@ export class SnapController extends BaseController<
19661967
* @param encrypted - A flag to indicate whether to use encrypted storage or
19671968
* not.
19681969
*/
1969-
async #persistSnapState(
1970-
snapId: SnapId,
1971-
newSnapState: Record<string, Json> | null,
1972-
encrypted: boolean,
1973-
) {
1974-
const runtime = this.#getRuntimeExpect(snapId);
1975-
await runtime.stateMutex.runExclusive(async () => {
1976-
const newState = await this.#getStateToPersist(
1977-
snapId,
1978-
newSnapState,
1979-
encrypted,
1980-
);
1970+
readonly #persistSnapState = debounce(
1971+
(
1972+
snapId: SnapId,
1973+
newSnapState: Record<string, Json> | null,
1974+
encrypted: boolean,
1975+
) => {
1976+
const runtime = this.#getRuntimeExpect(snapId);
1977+
runtime.stateMutex
1978+
.runExclusive(async () => {
1979+
const newState = await this.#getStateToPersist(
1980+
snapId,
1981+
newSnapState,
1982+
encrypted,
1983+
);
19811984

1982-
if (encrypted) {
1983-
return this.update((state) => {
1984-
state.snapStates[snapId] = newState;
1985-
});
1986-
}
1985+
if (encrypted) {
1986+
return this.update((state) => {
1987+
state.snapStates[snapId] = newState;
1988+
});
1989+
}
19871990

1988-
return this.update((state) => {
1989-
state.unencryptedSnapStates[snapId] = newState;
1990-
});
1991-
});
1992-
}
1991+
return this.update((state) => {
1992+
state.unencryptedSnapStates[snapId] = newState;
1993+
});
1994+
})
1995+
.catch(logError);
1996+
},
1997+
5_000,
1998+
);
19931999

19942000
/**
19952001
* Updates the own state of the snap with the given id.
@@ -2012,11 +2018,7 @@ export class SnapController extends BaseController<
20122018
runtime.unencryptedState = newSnapState;
20132019
}
20142020

2015-
// This is intentionally run asynchronously to avoid blocking the main
2016-
// thread.
2017-
this.#persistSnapState(snapId, newSnapState, encrypted).catch((error) => {
2018-
logError(error);
2019-
});
2021+
this.#persistSnapState(snapId, newSnapState, encrypted);
20202022
}
20212023

20222024
/**
@@ -2034,11 +2036,7 @@ export class SnapController extends BaseController<
20342036
runtime.unencryptedState = null;
20352037
}
20362038

2037-
// This is intentionally run asynchronously to avoid blocking the main
2038-
// thread.
2039-
this.#persistSnapState(snapId, null, encrypted).catch((error) => {
2040-
logError(error);
2041-
});
2039+
this.#persistSnapState(snapId, null, encrypted);
20422040
}
20432041

20442042
/**

‎packages/snaps-controllers/src/utils.test.ts‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import {
1313
MOCK_RPC_ORIGINS_PERMISSION,
1414
MOCK_SNAP_DIALOG_PERMISSION,
1515
} from './test-utils';
16-
import { getSnapFiles, permissionsDiff, setDiff } from './utils';
16+
import { debounce, getSnapFiles, permissionsDiff, setDiff } from './utils';
1717
import { SnapEndowments } from '../../snaps-rpc-methods/src/endowments';
1818

1919
describe('setDiff', () => {
@@ -180,3 +180,30 @@ describe('getSnapFiles', () => {
180180
]);
181181
});
182182
});
183+
184+
describe('debounce', () => {
185+
beforeAll(() => {
186+
jest.useFakeTimers();
187+
});
188+
189+
afterAll(() => {
190+
jest.useRealTimers();
191+
});
192+
193+
it('debounces a function based on a key', () => {
194+
const fn = jest.fn();
195+
const debounced = debounce(fn, 100);
196+
197+
expect(debounced('foo')).toBeUndefined();
198+
expect(debounced('foo')).toBeUndefined();
199+
expect(debounced('bar')).toBeUndefined();
200+
expect(debounced('bar')).toBeUndefined();
201+
202+
expect(fn).toHaveBeenCalledTimes(0);
203+
204+
jest.advanceTimersByTime(100);
205+
expect(fn).toHaveBeenCalledTimes(2);
206+
expect(fn).toHaveBeenNthCalledWith(1, 'foo');
207+
expect(fn).toHaveBeenNthCalledWith(2, 'bar');
208+
});
209+
});

‎packages/snaps-controllers/src/utils.ts‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,3 +329,41 @@ export async function fetchSnap(snapId: SnapId, location: SnapLocation) {
329329
);
330330
}
331331
}
332+
333+
/**
334+
* Debounce a function based on the given key, i.e., the function will only be
335+
* called after the timeout has passed since the last call for the same key.
336+
*
337+
* @param fn - The function to debounce.
338+
* @param timeout - The timeout in milliseconds. Defaults to 1000.
339+
* @returns A debounced function.
340+
* @template Key - The key to debounce the function on.
341+
* @template Args - The arguments of the function.
342+
* @example
343+
* const originalFunction = (key: string, value: number) => {
344+
* console.log(`Called with key: ${key} and value: ${value}`);
345+
* };
346+
*
347+
* const debouncedFunction = debounce(originalFunction);
348+
* debouncedFunction('foo', 1);
349+
*/
350+
export function debounce<Key, Args extends unknown[]>(
351+
fn: (key: Key, ...args: Args) => void,
352+
timeout = 1000,
353+
) {
354+
const timeouts = new Map<Key, NodeJS.Timeout>();
355+
356+
return (key: Key, ...args: Args): void => {
357+
if (timeouts.has(key)) {
358+
clearTimeout(timeouts.get(key));
359+
}
360+
361+
timeouts.set(
362+
key,
363+
setTimeout(() => {
364+
fn(key, ...args);
365+
timeouts.delete(key);
366+
}, timeout),
367+
);
368+
};
369+
}

0 commit comments

Comments
 (0)