Skip to content

Commit bc34202

Browse files
Merge pull request #2464 from github/robertbrignull/invocation-rate-limiter
Move the invocation rate limiter class out of helpers.ts
2 parents 157210f + 50c46b6 commit bc34202

File tree

6 files changed

+243
-355
lines changed

6 files changed

+243
-355
lines changed

extensions/ql-vscode/src/codeql-cli/distribution.ts

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,7 @@ import * as semver from "semver";
66
import { URL } from "url";
77
import { ExtensionContext, Event } from "vscode";
88
import { DistributionConfig } from "../config";
9-
import {
10-
InvocationRateLimiter,
11-
InvocationRateLimiterResultKind,
12-
showAndLogErrorMessage,
13-
showAndLogWarningMessage,
14-
} from "../helpers";
9+
import { showAndLogErrorMessage, showAndLogWarningMessage } from "../helpers";
1510
import { extLogger } from "../common";
1611
import { getCodeQlCliVersion } from "./cli-version";
1712
import {
@@ -24,6 +19,10 @@ import {
2419
extractZipArchive,
2520
getRequiredAssetName,
2621
} from "../pure/distribution";
22+
import {
23+
InvocationRateLimiter,
24+
InvocationRateLimiterResultKind,
25+
} from "../common/invocation-rate-limiter";
2726

2827
/**
2928
* distribution.ts
@@ -76,7 +75,7 @@ export class DistributionManager implements DistributionProvider {
7675
extensionContext,
7776
);
7877
this.updateCheckRateLimiter = new InvocationRateLimiter(
79-
extensionContext,
78+
extensionContext.globalState,
8079
"extensionSpecificDistributionUpdateCheck",
8180
() =>
8281
this.extensionSpecificDistributionManager.checkForUpdatesToDistribution(),
Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import { Memento } from "./memento";
2+
3+
/**
4+
* Provides a utility method to invoke a function only if a minimum time interval has elapsed since
5+
* the last invocation of that function.
6+
*/
7+
export class InvocationRateLimiter<T> {
8+
constructor(
9+
private readonly globalState: Memento,
10+
private readonly funcIdentifier: string,
11+
private readonly func: () => Promise<T>,
12+
private readonly createDate: (dateString?: string) => Date = (s) =>
13+
s ? new Date(s) : new Date(),
14+
) {}
15+
16+
/**
17+
* Invoke the function if `minSecondsSinceLastInvocation` seconds have elapsed since the last invocation.
18+
*/
19+
public async invokeFunctionIfIntervalElapsed(
20+
minSecondsSinceLastInvocation: number,
21+
): Promise<InvocationRateLimiterResult<T>> {
22+
const updateCheckStartDate = this.createDate();
23+
const lastInvocationDate = this.getLastInvocationDate();
24+
if (
25+
minSecondsSinceLastInvocation &&
26+
lastInvocationDate &&
27+
lastInvocationDate <= updateCheckStartDate &&
28+
lastInvocationDate.getTime() + minSecondsSinceLastInvocation * 1000 >
29+
updateCheckStartDate.getTime()
30+
) {
31+
return createRateLimitedResult();
32+
}
33+
const result = await this.func();
34+
await this.setLastInvocationDate(updateCheckStartDate);
35+
return createInvokedResult(result);
36+
}
37+
38+
private getLastInvocationDate(): Date | undefined {
39+
const maybeDateString: string | undefined = this.globalState.get(
40+
InvocationRateLimiter._invocationRateLimiterPrefix + this.funcIdentifier,
41+
);
42+
return maybeDateString ? this.createDate(maybeDateString) : undefined;
43+
}
44+
45+
private async setLastInvocationDate(date: Date): Promise<void> {
46+
return await this.globalState.update(
47+
InvocationRateLimiter._invocationRateLimiterPrefix + this.funcIdentifier,
48+
date,
49+
);
50+
}
51+
52+
private static readonly _invocationRateLimiterPrefix =
53+
"invocationRateLimiter_lastInvocationDate_";
54+
}
55+
56+
export enum InvocationRateLimiterResultKind {
57+
Invoked,
58+
RateLimited,
59+
}
60+
61+
/**
62+
* The function was invoked and returned the value `result`.
63+
*/
64+
interface InvokedResult<T> {
65+
kind: InvocationRateLimiterResultKind.Invoked;
66+
result: T;
67+
}
68+
69+
/**
70+
* The function was not invoked as the minimum interval since the last invocation had not elapsed.
71+
*/
72+
interface RateLimitedResult {
73+
kind: InvocationRateLimiterResultKind.RateLimited;
74+
}
75+
76+
type InvocationRateLimiterResult<T> = InvokedResult<T> | RateLimitedResult;
77+
78+
function createInvokedResult<T>(result: T): InvokedResult<T> {
79+
return {
80+
kind: InvocationRateLimiterResultKind.Invoked,
81+
result,
82+
};
83+
}
84+
85+
function createRateLimitedResult(): RateLimitedResult {
86+
return {
87+
kind: InvocationRateLimiterResultKind.RateLimited,
88+
};
89+
}

extensions/ql-vscode/src/helpers.ts

Lines changed: 1 addition & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,7 @@ import { glob } from "glob";
1010
import { load } from "js-yaml";
1111
import { join, basename, dirname } from "path";
1212
import { dirSync } from "tmp-promise";
13-
import {
14-
ExtensionContext,
15-
Uri,
16-
window as Window,
17-
workspace,
18-
env,
19-
WorkspaceFolder,
20-
} from "vscode";
13+
import { Uri, window as Window, workspace, env, WorkspaceFolder } from "vscode";
2114
import { CodeQLCliServer, QlpacksInfo } from "./codeql-cli/cli";
2215
import { UserCancellationException } from "./common/vscode/progress";
2316
import { extLogger, OutputChannelLogger } from "./common";
@@ -363,106 +356,6 @@ export async function prepareCodeTour(
363356
}
364357
}
365358

366-
/**
367-
* Provides a utility method to invoke a function only if a minimum time interval has elapsed since
368-
* the last invocation of that function.
369-
*/
370-
export class InvocationRateLimiter<T> {
371-
constructor(
372-
extensionContext: ExtensionContext,
373-
funcIdentifier: string,
374-
func: () => Promise<T>,
375-
createDate: (dateString?: string) => Date = (s) =>
376-
s ? new Date(s) : new Date(),
377-
) {
378-
this._createDate = createDate;
379-
this._extensionContext = extensionContext;
380-
this._func = func;
381-
this._funcIdentifier = funcIdentifier;
382-
}
383-
384-
/**
385-
* Invoke the function if `minSecondsSinceLastInvocation` seconds have elapsed since the last invocation.
386-
*/
387-
public async invokeFunctionIfIntervalElapsed(
388-
minSecondsSinceLastInvocation: number,
389-
): Promise<InvocationRateLimiterResult<T>> {
390-
const updateCheckStartDate = this._createDate();
391-
const lastInvocationDate = this.getLastInvocationDate();
392-
if (
393-
minSecondsSinceLastInvocation &&
394-
lastInvocationDate &&
395-
lastInvocationDate <= updateCheckStartDate &&
396-
lastInvocationDate.getTime() + minSecondsSinceLastInvocation * 1000 >
397-
updateCheckStartDate.getTime()
398-
) {
399-
return createRateLimitedResult();
400-
}
401-
const result = await this._func();
402-
await this.setLastInvocationDate(updateCheckStartDate);
403-
return createInvokedResult(result);
404-
}
405-
406-
private getLastInvocationDate(): Date | undefined {
407-
const maybeDateString: string | undefined =
408-
this._extensionContext.globalState.get(
409-
InvocationRateLimiter._invocationRateLimiterPrefix +
410-
this._funcIdentifier,
411-
);
412-
return maybeDateString ? this._createDate(maybeDateString) : undefined;
413-
}
414-
415-
private async setLastInvocationDate(date: Date): Promise<void> {
416-
return await this._extensionContext.globalState.update(
417-
InvocationRateLimiter._invocationRateLimiterPrefix + this._funcIdentifier,
418-
date,
419-
);
420-
}
421-
422-
private readonly _createDate: (dateString?: string) => Date;
423-
private readonly _extensionContext: ExtensionContext;
424-
private readonly _func: () => Promise<T>;
425-
private readonly _funcIdentifier: string;
426-
427-
private static readonly _invocationRateLimiterPrefix =
428-
"invocationRateLimiter_lastInvocationDate_";
429-
}
430-
431-
export enum InvocationRateLimiterResultKind {
432-
Invoked,
433-
RateLimited,
434-
}
435-
436-
/**
437-
* The function was invoked and returned the value `result`.
438-
*/
439-
interface InvokedResult<T> {
440-
kind: InvocationRateLimiterResultKind.Invoked;
441-
result: T;
442-
}
443-
444-
/**
445-
* The function was not invoked as the minimum interval since the last invocation had not elapsed.
446-
*/
447-
interface RateLimitedResult {
448-
kind: InvocationRateLimiterResultKind.RateLimited;
449-
}
450-
451-
type InvocationRateLimiterResult<T> = InvokedResult<T> | RateLimitedResult;
452-
453-
function createInvokedResult<T>(result: T): InvokedResult<T> {
454-
return {
455-
kind: InvocationRateLimiterResultKind.Invoked,
456-
result,
457-
};
458-
}
459-
460-
function createRateLimitedResult(): RateLimitedResult {
461-
return {
462-
kind: InvocationRateLimiterResultKind.RateLimited,
463-
};
464-
}
465-
466359
export interface QlPacksForLanguage {
467360
/** The name of the pack containing the dbscheme. */
468361
dbschemePack: string;
Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
import type { Memento } from "vscode";
2+
import { InvocationRateLimiter } from "../../../src/common/invocation-rate-limiter";
3+
4+
describe("Invocation rate limiter", () => {
5+
// 1 January 2020
6+
let currentUnixTime = 1577836800;
7+
8+
function createDate(dateString?: string): Date {
9+
if (dateString) {
10+
return new Date(dateString);
11+
}
12+
const numMillisecondsPerSecond = 1000;
13+
return new Date(currentUnixTime * numMillisecondsPerSecond);
14+
}
15+
16+
function createInvocationRateLimiter<T>(
17+
funcIdentifier: string,
18+
func: () => Promise<T>,
19+
): InvocationRateLimiter<T> {
20+
return new InvocationRateLimiter(
21+
new MockMemento(),
22+
funcIdentifier,
23+
func,
24+
(s) => createDate(s),
25+
);
26+
}
27+
28+
class MockMemento implements Memento {
29+
keys(): readonly string[] {
30+
throw new Error("Method not implemented.");
31+
}
32+
map = new Map<any, any>();
33+
34+
/**
35+
* Return a value.
36+
*
37+
* @param key A string.
38+
* @param defaultValue A value that should be returned when there is no
39+
* value (`undefined`) with the given key.
40+
* @return The stored value or the defaultValue.
41+
*/
42+
get<T>(key: string, defaultValue?: T): T {
43+
return this.map.has(key) ? this.map.get(key) : defaultValue;
44+
}
45+
46+
/**
47+
* Store a value. The value must be JSON-stringifyable.
48+
*
49+
* @param key A string.
50+
* @param value A value. MUST not contain cyclic references.
51+
*/
52+
async update(key: string, value: any): Promise<void> {
53+
this.map.set(key, value);
54+
}
55+
}
56+
57+
it("initially invokes function", async () => {
58+
let numTimesFuncCalled = 0;
59+
const invocationRateLimiter = createInvocationRateLimiter(
60+
"funcid",
61+
async () => {
62+
numTimesFuncCalled++;
63+
},
64+
);
65+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
66+
expect(numTimesFuncCalled).toBe(1);
67+
});
68+
69+
it("doesn't invoke function again if no time has passed", async () => {
70+
let numTimesFuncCalled = 0;
71+
const invocationRateLimiter = createInvocationRateLimiter(
72+
"funcid",
73+
async () => {
74+
numTimesFuncCalled++;
75+
},
76+
);
77+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
78+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
79+
expect(numTimesFuncCalled).toBe(1);
80+
});
81+
82+
it("doesn't invoke function again if requested time since last invocation hasn't passed", async () => {
83+
let numTimesFuncCalled = 0;
84+
const invocationRateLimiter = createInvocationRateLimiter(
85+
"funcid",
86+
async () => {
87+
numTimesFuncCalled++;
88+
},
89+
);
90+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(100);
91+
currentUnixTime += 1;
92+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(2);
93+
expect(numTimesFuncCalled).toBe(1);
94+
});
95+
96+
it("invokes function again immediately if requested time since last invocation is 0 seconds", async () => {
97+
let numTimesFuncCalled = 0;
98+
const invocationRateLimiter = createInvocationRateLimiter(
99+
"funcid",
100+
async () => {
101+
numTimesFuncCalled++;
102+
},
103+
);
104+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(0);
105+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(0);
106+
expect(numTimesFuncCalled).toBe(2);
107+
});
108+
109+
it("invokes function again after requested time since last invocation has elapsed", async () => {
110+
let numTimesFuncCalled = 0;
111+
const invocationRateLimiter = createInvocationRateLimiter(
112+
"funcid",
113+
async () => {
114+
numTimesFuncCalled++;
115+
},
116+
);
117+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(1);
118+
currentUnixTime += 1;
119+
await invocationRateLimiter.invokeFunctionIfIntervalElapsed(1);
120+
expect(numTimesFuncCalled).toBe(2);
121+
});
122+
123+
it("invokes functions with different rate limiters", async () => {
124+
let numTimesFuncACalled = 0;
125+
const invocationRateLimiterA = createInvocationRateLimiter(
126+
"funcid",
127+
async () => {
128+
numTimesFuncACalled++;
129+
},
130+
);
131+
let numTimesFuncBCalled = 0;
132+
const invocationRateLimiterB = createInvocationRateLimiter(
133+
"funcid",
134+
async () => {
135+
numTimesFuncBCalled++;
136+
},
137+
);
138+
await invocationRateLimiterA.invokeFunctionIfIntervalElapsed(100);
139+
await invocationRateLimiterB.invokeFunctionIfIntervalElapsed(100);
140+
expect(numTimesFuncACalled).toBe(1);
141+
expect(numTimesFuncBCalled).toBe(1);
142+
});
143+
});

0 commit comments

Comments
 (0)