Skip to content

Commit 82bedfc

Browse files
lwin-kyawlwin-kyaw
committed
trying Cockatiel ServicePolicy
1 parent b078925 commit 82bedfc

File tree

4 files changed

+234
-1
lines changed

4 files changed

+234
-1
lines changed

‎packages/shield-controller/package.json‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,9 @@
4848
},
4949
"dependencies": {
5050
"@metamask/base-controller": "^8.4.1",
51-
"@metamask/utils": "^11.8.1"
51+
"@metamask/controller-utils": "^11.14.1",
52+
"@metamask/utils": "^11.8.1",
53+
"cockatiel": "^3.1.2"
5254
},
5355
"devDependencies": {
5456
"@babel/runtime": "^7.23.9",

‎packages/shield-controller/src/polling-with-policy.test.ts‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
import { PollingWithCockatielPolicy } from './polling-with-policy';
2+
import { delay } from '../tests/utils';
3+
4+
describe('PollingWithCockatielPolicy', () => {
5+
it('should return the success result', async () => {
6+
const policy = new PollingWithCockatielPolicy();
7+
const result = await policy.start('test', async () => {
8+
return 'test';
9+
});
10+
expect(result).toBe('test');
11+
});
12+
13+
it('should retry the request and complete successfully', async () => {
14+
const policy = new PollingWithCockatielPolicy();
15+
let invocationCount = 0;
16+
const mockRequestFn = jest
17+
.fn()
18+
.mockImplementation(async (_abortSignal: AbortSignal) => {
19+
invocationCount += 1;
20+
return new Promise((resolve, reject) => {
21+
setTimeout(() => {
22+
// eslint-disable-next-line jest/no-conditional-in-test
23+
if (invocationCount < 3) {
24+
reject(new Error('test error'));
25+
}
26+
resolve('test');
27+
}, 100);
28+
});
29+
});
30+
const result = await policy.start('test', mockRequestFn);
31+
expect(result).toBe('test');
32+
expect(mockRequestFn).toHaveBeenCalledTimes(3);
33+
});
34+
35+
it('should not retry when the error is not retryable', async () => {
36+
const policy = new PollingWithCockatielPolicy();
37+
const mockRequestFn = jest
38+
.fn()
39+
.mockImplementation(async (_abortSignal: AbortSignal) => {
40+
return new Promise((_resolve, reject) => {
41+
const error = new Error('Not retryable error') as {
42+
shouldRetry?: boolean;
43+
};
44+
error.shouldRetry = false;
45+
// eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
46+
reject(error);
47+
});
48+
});
49+
await expect(policy.start('test', mockRequestFn)).rejects.toThrow(
50+
'Not retryable error',
51+
);
52+
expect(mockRequestFn).toHaveBeenCalledTimes(1);
53+
});
54+
55+
it('should throw an error when the retry exceeds the max retries', async () => {
56+
const policy = new PollingWithCockatielPolicy({
57+
maxRetries: 3,
58+
});
59+
60+
const requestFn = jest
61+
.fn()
62+
.mockImplementation(async (_abortSignal: AbortSignal) => {
63+
return new Promise((_resolve, reject) => {
64+
setTimeout(() => {
65+
reject(new Error('test error'));
66+
}, 100);
67+
});
68+
});
69+
70+
const result = policy.start('test', requestFn);
71+
await expect(result).rejects.toThrow('test error');
72+
expect(requestFn).toHaveBeenCalledTimes(4);
73+
});
74+
75+
it('should throw a `Request Cancelled` error when the request is aborted', async () => {
76+
const policy = new PollingWithCockatielPolicy({
77+
maxRetries: 3,
78+
});
79+
80+
const requestFn = jest
81+
.fn()
82+
.mockImplementation(async (abortSignal: AbortSignal) => {
83+
return new Promise((_resolve, reject) => {
84+
setTimeout(() => {
85+
// eslint-disable-next-line jest/no-conditional-in-test
86+
if (abortSignal.aborted) {
87+
reject(new Error('test error'));
88+
}
89+
reject(new Error('test error'));
90+
}, 100);
91+
});
92+
});
93+
94+
const result = policy.start('test', requestFn);
95+
await delay(10);
96+
policy.abortPendingRequest('test');
97+
await expect(result).rejects.toThrow('Request cancelled');
98+
});
99+
100+
it('should throw a `Request Cancelled` error when a new request is started with the same request id', async () => {
101+
const policy = new PollingWithCockatielPolicy();
102+
103+
const requestFn = jest
104+
.fn()
105+
.mockImplementation(async (abortSignal: AbortSignal) => {
106+
return new Promise((resolve, reject) => {
107+
setTimeout(() => {
108+
// eslint-disable-next-line jest/no-conditional-in-test
109+
if (abortSignal.aborted) {
110+
reject(new Error('test error'));
111+
}
112+
resolve('test');
113+
}, 100);
114+
});
115+
});
116+
117+
const result = policy.start('test', requestFn);
118+
await delay(10);
119+
const secondResult = policy.start('test', requestFn);
120+
await expect(result).rejects.toThrow('Request cancelled');
121+
expect(await secondResult).toBe('test');
122+
});
123+
124+
it('should resolve the result when two requests are started with the different request ids', async () => {
125+
const policy = new PollingWithCockatielPolicy();
126+
127+
const requestFn = (result: string) =>
128+
jest.fn().mockImplementation(async (abortSignal: AbortSignal) => {
129+
return new Promise((resolve, reject) => {
130+
// eslint-disable-next-line jest/no-conditional-in-test
131+
if (abortSignal.aborted) {
132+
reject(new Error('test error'));
133+
}
134+
setTimeout(() => {
135+
resolve(result);
136+
}, 100);
137+
});
138+
});
139+
140+
const result = policy.start('test', requestFn('test'));
141+
const secondResult = policy.start('test2', requestFn('test2'));
142+
expect(await result).toBe('test');
143+
expect(await secondResult).toBe('test2');
144+
});
145+
});

‎packages/shield-controller/src/polling-with-policy.ts‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
import {
2+
createServicePolicy,
3+
type CreateServicePolicyOptions,
4+
type ServicePolicy,
5+
} from '@metamask/controller-utils';
6+
import { handleWhen } from 'cockatiel';
7+
8+
export type RequestFn<ReturnType> = (
9+
signal: AbortSignal,
10+
) => Promise<ReturnType>;
11+
12+
export class PollingWithCockatielPolicy {
13+
readonly #policy: ServicePolicy;
14+
15+
readonly #requestEntry = new Map<string, AbortController>();
16+
17+
constructor(policyOptions: CreateServicePolicyOptions = {}) {
18+
const retryFilterPolicy = handleWhen(this.#shouldRetry);
19+
this.#policy = createServicePolicy({
20+
...policyOptions,
21+
retryFilterPolicy,
22+
});
23+
}
24+
25+
async start<ReturnType>(requestId: string, requestFn: RequestFn<ReturnType>) {
26+
this.abortPendingRequest(requestId);
27+
const abortController = this.addNewRequestEntry(requestId);
28+
const disposableListeners = this.#registerListeners();
29+
30+
try {
31+
const result = await this.#policy.execute(
32+
async ({ signal: abortSignal }) => {
33+
return requestFn(abortSignal);
34+
},
35+
abortController.signal,
36+
);
37+
return result;
38+
} catch (error) {
39+
if (abortController.signal.aborted) {
40+
throw new Error('Request cancelled');
41+
}
42+
throw error;
43+
} finally {
44+
this.#unregisterListeners(disposableListeners);
45+
}
46+
}
47+
48+
addNewRequestEntry(requestId: string) {
49+
const abortController = new AbortController();
50+
this.#requestEntry.set(requestId, abortController);
51+
return abortController;
52+
}
53+
54+
abortPendingRequest(requestId: string) {
55+
const abortController = this.#requestEntry.get(requestId);
56+
abortController?.abort();
57+
}
58+
59+
// TODO: remove
60+
#registerListeners(): { dispose: () => void }[] {
61+
const disposableListeners = [];
62+
disposableListeners.push(
63+
this.#policy.onBreak((data) => {
64+
console.log('onBreak', data);
65+
}),
66+
);
67+
disposableListeners.push(
68+
this.#policy.circuitBreakerPolicy.onStateChange((data) => {
69+
console.log('onStateChange', data);
70+
}),
71+
);
72+
return disposableListeners;
73+
}
74+
75+
// TODO: remove
76+
#unregisterListeners(disposableListeners: { dispose: () => void }[]) {
77+
disposableListeners.forEach((disposable) => disposable.dispose());
78+
}
79+
80+
#shouldRetry(error: unknown): boolean {
81+
const errorWithRetryStatus = error as { shouldRetry?: boolean };
82+
return errorWithRetryStatus.shouldRetry !== false;
83+
}
84+
}

‎yarn.lock‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4706,11 +4706,13 @@ __metadata:
47064706
"@lavamoat/preinstall-always-fail": "npm:^2.1.0"
47074707
"@metamask/auto-changelog": "npm:^3.4.4"
47084708
"@metamask/base-controller": "npm:^8.4.1"
4709+
"@metamask/controller-utils": "npm:^11.14.1"
47094710
"@metamask/signature-controller": "npm:^34.0.1"
47104711
"@metamask/transaction-controller": "npm:^60.9.0"
47114712
"@metamask/utils": "npm:^11.8.1"
47124713
"@ts-bridge/cli": "npm:^0.6.1"
47134714
"@types/jest": "npm:^27.4.1"
4715+
cockatiel: "npm:^3.1.2"
47144716
deepmerge: "npm:^4.2.2"
47154717
jest: "npm:^27.5.1"
47164718
ts-jest: "npm:^27.1.4"

0 commit comments

Comments
 (0)