Skip to content

Commit f289ddb

Browse files
vinistockvinistock
committed
Make TCP server persistent
1 parent fa4d84d commit f289ddb

File tree

2 files changed

+139
-115
lines changed

2 files changed

+139
-115
lines changed

‎vscode/src/streamingRunner.ts

+133-109
Original file line numberDiff line numberDiff line change
@@ -30,163 +30,180 @@ export enum Mode {
3030

3131
// The StreamingRunner class is responsible for executing the test process or launching the debugger while handling the
3232
// streaming events to update the test explorer status
33-
export class StreamingRunner {
34-
private readonly promises: Promise<void>[] = [];
35-
private readonly disposables: vscode.Disposable[] = [];
36-
private readonly run: vscode.TestRun;
33+
export class StreamingRunner implements vscode.Disposable {
34+
private promises: Promise<void>[] = [];
35+
private disposables: vscode.Disposable[] = [];
3736
private readonly findTestItem: (
3837
id: string,
3938
uri: vscode.Uri,
4039
) => Promise<vscode.TestItem | undefined>;
4140

41+
private readonly tcpServer: net.Server;
42+
private tcpPort: string | undefined;
43+
private connection: rpc.MessageConnection | undefined;
44+
private executionPromise:
45+
| { resolve: () => void; reject: (error: Error) => void }
46+
| undefined;
47+
48+
private run: vscode.TestRun | undefined;
49+
4250
constructor(
43-
run: vscode.TestRun,
4451
findTestItem: (
4552
id: string,
4653
uri: vscode.Uri,
4754
) => Promise<vscode.TestItem | undefined>,
4855
) {
49-
this.run = run;
5056
this.findTestItem = findTestItem;
57+
this.tcpServer = this.startServer();
5158
}
5259

5360
async execute(
61+
currentRun: vscode.TestRun,
5462
command: string,
5563
env: NodeJS.ProcessEnv,
5664
workspace: Workspace,
5765
mode: Mode,
5866
linkedCancellationSource: LinkedCancellationSource,
5967
) {
68+
this.run = currentRun;
69+
6070
await new Promise<void>((resolve, reject) => {
61-
const server = net.createServer();
62-
server.on("error", reject);
63-
server.unref();
64-
65-
server.listen(0, "localhost", async () => {
66-
const address = server.address();
67-
const serverPort =
68-
typeof address === "string" ? address : address?.port.toString();
69-
70-
if (!serverPort) {
71-
reject(
72-
new Error(
73-
"Failed to set up TCP server to communicate with test process",
74-
),
75-
);
76-
return;
77-
}
78-
79-
const abortController = new AbortController();
80-
81-
server.on("connection", (socket) => {
82-
const connection = rpc.createMessageConnection(
83-
new rpc.StreamMessageReader(socket),
84-
new rpc.StreamMessageWriter(socket),
85-
);
86-
const finalize = () => {
87-
Promise.all(this.promises)
88-
.then(() => {
89-
this.disposables.forEach((disposable) => disposable.dispose());
90-
connection.end();
91-
connection.dispose();
92-
server.close();
93-
resolve();
94-
})
95-
.catch(reject);
96-
};
97-
98-
// We resolve the promise and perform cleanup on two occasions: if the test run finished normally, then we
99-
// should receive the finish event. The other case is when the run is cancelled and the abort controller gets
100-
// triggered, in which case we will not receive the finish event
101-
linkedCancellationSource.onCancellationRequested(() => {
102-
this.run.appendOutput("\r\nTest run cancelled.");
103-
abortController.abort();
104-
finalize();
105-
});
106-
107-
this.disposables.push(
108-
connection.onNotification(NOTIFICATION_TYPES.finish, finalize),
109-
);
110-
111-
this.registerStreamingEvents(connection);
112-
113-
// Start listening for events
114-
connection.listen();
115-
});
116-
117-
if (mode === Mode.Run) {
118-
this.spawnTestProcess(
119-
command,
120-
env,
121-
workspace.workspaceFolder.uri.fsPath,
122-
serverPort,
123-
abortController,
124-
);
125-
} else {
126-
await this.launchDebugger(command, env, workspace, serverPort);
127-
}
71+
this.executionPromise = { resolve, reject };
72+
const abortController = new AbortController();
73+
74+
linkedCancellationSource.onCancellationRequested(async () => {
75+
this.run!.appendOutput("\r\nTest run cancelled.");
76+
abortController.abort();
77+
await this.finalize();
12878
});
79+
80+
if (mode === Mode.Run) {
81+
this.spawnTestProcess(
82+
command,
83+
env,
84+
workspace.workspaceFolder.uri.fsPath,
85+
abortController,
86+
);
87+
} else {
88+
// eslint-disable-next-line @typescript-eslint/no-floating-promises
89+
this.launchDebugger(command, env, workspace);
90+
}
12991
});
13092
}
13193

94+
dispose() {
95+
this.tcpServer.close();
96+
this.connection?.dispose();
97+
}
98+
13299
// Launches the debugger with streaming updates
133100
private async launchDebugger(
134101
command: string,
135102
env: NodeJS.ProcessEnv,
136103
workspace: Workspace,
137-
serverPort: string,
138104
) {
139-
await vscode.debug
140-
.startDebugging(
141-
workspace.workspaceFolder,
142-
{
143-
type: "ruby_lsp",
144-
name: "Debug",
145-
request: "launch",
146-
program: command,
147-
env: {
148-
...env,
149-
DISABLE_SPRING: "1",
150-
RUBY_LSP_REPORTER_PORT: serverPort,
151-
},
105+
const successFullyStarted = await vscode.debug.startDebugging(
106+
workspace.workspaceFolder,
107+
{
108+
type: "ruby_lsp",
109+
name: "Debug",
110+
request: "launch",
111+
program: command,
112+
env: {
113+
...env,
114+
DISABLE_SPRING: "1",
115+
RUBY_LSP_REPORTER_PORT: this.tcpPort,
152116
},
153-
{ testRun: this.run },
154-
)
155-
.then((successFullyStarted) => {
156-
if (!successFullyStarted) {
157-
throw new Error("Failed to start debugging session");
158-
}
159-
});
117+
},
118+
{ testRun: this.run },
119+
);
120+
121+
if (!successFullyStarted) {
122+
this.executionPromise!.reject(
123+
new Error("Failed to start debugging session"),
124+
);
125+
}
160126
}
161127

162128
// Spawns the test process and redirects any stdout or stderr output to the test run output
163129
private spawnTestProcess(
164130
command: string,
165131
env: NodeJS.ProcessEnv,
166132
cwd: string,
167-
serverPort: string,
168133
abortController: AbortController,
169134
) {
170135
const testProcess = spawn(command, {
171-
env: { ...env, RUBY_LSP_REPORTER_PORT: serverPort },
136+
env: { ...env, RUBY_LSP_REPORTER_PORT: this.tcpPort },
172137
stdio: ["pipe", "pipe", "pipe"],
173138
shell: true,
174139
signal: abortController.signal,
175140
cwd,
176141
});
177142

178143
testProcess.stdout.on("data", (data) => {
179-
this.run.appendOutput(data.toString().replace(/\n/g, "\r\n"));
144+
this.run!.appendOutput(data.toString().replace(/\n/g, "\r\n"));
180145
});
181146

182147
testProcess.stderr.on("data", (data) => {
183-
this.run.appendOutput(data.toString().replace(/\n/g, "\r\n"));
148+
this.run!.appendOutput(data.toString().replace(/\n/g, "\r\n"));
149+
});
150+
}
151+
152+
private startServer() {
153+
const server = net.createServer();
154+
server.on("error", (error) => {
155+
throw error;
156+
});
157+
server.unref();
158+
159+
server.listen(0, "localhost", () => {
160+
const address = server.address();
161+
162+
if (!address) {
163+
throw new Error("Failed setup TCP server for streaming updates");
164+
}
165+
this.tcpPort =
166+
typeof address === "string" ? address : address.port.toString();
167+
168+
// On any new connection to the TCP server, attach the JSON RPC reader and the events we defined
169+
server.on("connection", (socket) => {
170+
this.connection = rpc.createMessageConnection(
171+
new rpc.StreamMessageReader(socket),
172+
new rpc.StreamMessageWriter(socket),
173+
);
174+
175+
// Register and start listening for events
176+
this.registerStreamingEvents();
177+
this.connection.listen();
178+
});
184179
});
180+
181+
return server;
182+
}
183+
184+
private async finalize() {
185+
await Promise.all(this.promises);
186+
187+
this.disposables.forEach((disposable) => disposable.dispose());
188+
189+
this.promises = [];
190+
this.disposables = [];
191+
192+
if (this.connection) {
193+
this.connection.end();
194+
this.connection.dispose();
195+
}
196+
197+
this.executionPromise!.resolve();
185198
}
186199

187200
// Registers all streaming events that we will receive from the server except for the finish event, which is
188201
// registered to resolve the execute promise
189-
private registerStreamingEvents(connection: rpc.MessageConnection) {
202+
private registerStreamingEvents() {
203+
if (!this.connection) {
204+
return;
205+
}
206+
190207
const startTimestamps = new Map<string, number>();
191208
const withDuration = (
192209
id: string,
@@ -199,12 +216,19 @@ export class StreamingRunner {
199216

200217
// Handle the JSON events being emitted by the tests
201218
this.disposables.push(
202-
connection.onNotification(NOTIFICATION_TYPES.start, (params) => {
219+
this.connection.onNotification(
220+
NOTIFICATION_TYPES.finish,
221+
this.finalize.bind(this),
222+
),
223+
);
224+
225+
this.disposables.push(
226+
this.connection.onNotification(NOTIFICATION_TYPES.start, (params) => {
203227
this.promises.push(
204228
this.findTestItem(params.id, vscode.Uri.parse(params.uri)).then(
205229
(test) => {
206230
if (test) {
207-
this.run.started(test);
231+
this.run!.started(test);
208232
startTimestamps.set(test.id, Date.now());
209233
}
210234
},
@@ -214,13 +238,13 @@ export class StreamingRunner {
214238
);
215239

216240
this.disposables.push(
217-
connection.onNotification(NOTIFICATION_TYPES.pass, (params) => {
241+
this.connection.onNotification(NOTIFICATION_TYPES.pass, (params) => {
218242
this.promises.push(
219243
this.findTestItem(params.id, vscode.Uri.parse(params.uri)).then(
220244
(test) => {
221245
if (test) {
222246
withDuration(test.id, (duration) =>
223-
this.run.passed(test, duration),
247+
this.run!.passed(test, duration),
224248
);
225249
}
226250
},
@@ -230,13 +254,13 @@ export class StreamingRunner {
230254
);
231255

232256
this.disposables.push(
233-
connection.onNotification(NOTIFICATION_TYPES.fail, (params) => {
257+
this.connection.onNotification(NOTIFICATION_TYPES.fail, (params) => {
234258
this.promises.push(
235259
this.findTestItem(params.id, vscode.Uri.parse(params.uri)).then(
236260
(test) => {
237261
if (test) {
238262
withDuration(test.id, (duration) =>
239-
this.run.failed(
263+
this.run!.failed(
240264
test,
241265
new vscode.TestMessage(params.message),
242266
duration,
@@ -250,13 +274,13 @@ export class StreamingRunner {
250274
);
251275

252276
this.disposables.push(
253-
connection.onNotification(NOTIFICATION_TYPES.error, (params) => {
277+
this.connection.onNotification(NOTIFICATION_TYPES.error, (params) => {
254278
this.promises.push(
255279
this.findTestItem(params.id, vscode.Uri.parse(params.uri)).then(
256280
(test) => {
257281
if (test) {
258282
withDuration(test.id, (duration) =>
259-
this.run.errored(
283+
this.run!.errored(
260284
test,
261285
new vscode.TestMessage(params.message),
262286
duration,
@@ -270,12 +294,12 @@ export class StreamingRunner {
270294
);
271295

272296
this.disposables.push(
273-
connection.onNotification(NOTIFICATION_TYPES.skip, (params) => {
297+
this.connection.onNotification(NOTIFICATION_TYPES.skip, (params) => {
274298
this.promises.push(
275299
this.findTestItem(params.id, vscode.Uri.parse(params.uri)).then(
276300
(test) => {
277301
if (test) {
278-
this.run.skipped(test);
302+
this.run!.skipped(test);
279303
}
280304
},
281305
),

‎vscode/src/testController.ts

+6-6
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ export class TestController {
5959
vscode.FileCoverageDetail[]
6060
>();
6161

62+
private readonly runner = new StreamingRunner(this.findTestItem.bind(this));
63+
6264
constructor(
6365
context: vscode.ExtensionContext,
6466
telemetry: vscode.TelemetryLogger,
@@ -507,9 +509,8 @@ export class TestController {
507509

508510
for await (const command of response.commands) {
509511
try {
510-
const runner = new StreamingRunner(run, this.findTestItem.bind(this));
511-
512-
await runner.execute(
512+
await this.runner.execute(
513+
run,
513514
command,
514515
{
515516
...workspace.ruby.env,
@@ -552,9 +553,8 @@ export class TestController {
552553
? `${workspace.ruby.env.RUBYOPT} -rbundler/setup ${reporterPaths}`
553554
: `-rbundler/setup ${reporterPaths}`;
554555

555-
const runner = new StreamingRunner(run, this.findTestItem.bind(this));
556-
557-
await runner.execute(
556+
await this.runner.execute(
557+
run,
558558
command,
559559
{
560560
...workspace.ruby.env,

0 commit comments

Comments
 (0)