Skip to content

Commit c97136f

Browse files
brollbbrollb
authored
Add support for forking interactive sessions. Closes #1886 (#1906)
* Add support for forking interactive sessions. Closes #1886 * Add multiplexing test * Update unit tests for interactive client files
1 parent ef84001 commit c97136f

File tree

7 files changed

+179
-110
lines changed

7 files changed

+179
-110
lines changed

‎src/common/compute/interactive/message.js‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,25 +20,26 @@
2020
}
2121

2222
class Message {
23-
constructor(type, data) {
23+
constructor(sessionID, type, data) {
24+
this.sessionID = sessionID;
2425
this.type = type;
2526
this.data = data;
2627
}
2728

2829
static decode(serialized) {
29-
const {type, data} = JSON.parse(serialized);
30-
return new Message(type, data);
30+
const {sessionID, type, data} = JSON.parse(serialized);
31+
return new Message(sessionID, type, data);
3132
}
3233

3334
encode() {
34-
return Message.encode(this.type, this.data);
35+
return Message.encode(this.sessionID, this.type, this.data);
3536
}
3637

37-
static encode(type, data=0) {
38+
static encode(sessionID, type, data=0) {
3839
if (typeof Buffer !== 'undefined' && data instanceof Buffer) {
3940
data = data.toString();
4041
}
41-
return JSON.stringify({type, data});
42+
return JSON.stringify({sessionID, type, data});
4243
}
4344
}
4445
Object.assign(Message, Constants);

‎src/common/compute/interactive/session.js‎

Lines changed: 89 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -16,53 +16,14 @@ define([
1616
const {CommandFailedError} = Errors;
1717
const isNodeJs = typeof window === 'undefined';
1818
const WebSocket = isNodeJs ? require('ws') : window.WebSocket;
19+
let numSessions = 1;
1920

2021
class InteractiveSession {
21-
constructor(computeID, config={}) {
22+
constructor(channel) {
2223
this.currentTask = null;
23-
const address = gmeConfig.extensions.InteractiveComputeHost ||
24-
this.getDefaultServerURL();
25-
this.ws = new WebSocket(address);
26-
this.connected = defer();
27-
this.ws.onopen = () => {
28-
this.ws.send(JSON.stringify([computeID, config, this.getGMEToken()]));
29-
this.ws.onmessage = async (wsMsg) => {
30-
const data = await Task.getMessageData(wsMsg);
31-
32-
const msg = Message.decode(data);
33-
if (msg.type === Message.COMPLETE) {
34-
const err = msg.data;
35-
this.channel = new MessageChannel(this.ws);
36-
if (err) {
37-
this.connected.reject(err);
38-
} else {
39-
this.connected.resolve();
40-
this.checkReady();
41-
}
42-
}
43-
};
44-
};
45-
46-
this.ready = null;
47-
}
48-
49-
getDefaultServerURL() {
50-
const isSecure = !isNodeJs && location.protocol.includes('s');
51-
const protocol = isSecure ? 'wss' : 'ws';
52-
const defaultHost = isNodeJs ? '127.0.0.1' :
53-
location.origin
54-
.replace(location.protocol + '//', '')
55-
.replace(/:[0-9]+$/, '');
56-
return `${protocol}://${defaultHost}:${gmeConfig.server.port + 1}`;
57-
}
58-
59-
getGMEToken() {
60-
if (isNodeJs) {
61-
return '';
62-
}
63-
64-
const [, token] = (document.cookie || '').split('=');
65-
return token;
24+
this.id = numSessions++;
25+
this.channel = channel;
26+
this.channel.onClientConnect(this.id);
6627
}
6728

6829
checkReady() {
@@ -72,7 +33,7 @@ define([
7233
}
7334

7435
isIdle() {
75-
return !this.currentTask && this.ws.readyState === WebSocket.OPEN;
36+
return !this.currentTask && this.channel.isOpen();
7637
}
7738

7839
ensureIdle(action) {
@@ -84,7 +45,7 @@ define([
8445
spawn(cmd) {
8546
this.ensureIdle('spawn a task');
8647

87-
const msg = new Message(Message.RUN, cmd);
48+
const msg = new Message(this.id, Message.RUN, cmd);
8849
const task = new Task(this.channel, msg);
8950
this.runTask(task);
9051
return task;
@@ -111,7 +72,7 @@ define([
11172

11273
async exec(cmd) {
11374
this.ensureIdle('exec a task');
114-
const msg = new Message(Message.RUN, cmd);
75+
const msg = new Message(this.id, Message.RUN, cmd);
11576
const task = new Task(this.channel, msg);
11677
const result = {
11778
stdout: '',
@@ -131,14 +92,14 @@ define([
13192
async addArtifact(name, dataInfo, type, auth) {
13293
auth = auth || {};
13394
this.ensureIdle('add artifact');
134-
const msg = new Message(Message.ADD_ARTIFACT, [name, dataInfo, type, auth]);
95+
const msg = new Message(this.id, Message.ADD_ARTIFACT, [name, dataInfo, type, auth]);
13596
const task = new Task(this.channel, msg);
13697
await this.runTask(task);
13798
}
13899

139100
async saveArtifact(/*path, name, storageId, config*/) {
140101
this.ensureIdle('save artifact');
141-
const msg = new Message(Message.SAVE_ARTIFACT, [...arguments]);
102+
const msg = new Message(this.id, Message.SAVE_ARTIFACT, [...arguments]);
142103
const task = new Task(this.channel, msg);
143104
const [exitCode, dataInfo] = await this.runTask(task);
144105
if (exitCode) {
@@ -149,21 +110,21 @@ define([
149110

150111
async addFile(filepath, content) {
151112
this.ensureIdle('add file');
152-
const msg = new Message(Message.ADD_FILE, [filepath, content]);
113+
const msg = new Message(this.id, Message.ADD_FILE, [filepath, content]);
153114
const task = new Task(this.channel, msg);
154115
await this.runTask(task);
155116
}
156117

157118
async removeFile(filepath) {
158119
this.ensureIdle('remove file');
159-
const msg = new Message(Message.REMOVE_FILE, [filepath]);
120+
const msg = new Message(this.id, Message.REMOVE_FILE, [filepath]);
160121
const task = new Task(this.channel, msg);
161122
await this.runTask(task);
162123
}
163124

164125
async setEnvVar(name, value) {
165126
this.ensureIdle('set env var');
166-
const msg = new Message(Message.SET_ENV, [name, value]);
127+
const msg = new Message(this.id, Message.SET_ENV, [name, value]);
167128
const task = new Task(this.channel, msg);
168129
await this.runTask(task);
169130
}
@@ -174,24 +135,75 @@ define([
174135
'Cannot kill task. Must be a RUN task.'
175136
);
176137
if (task === this.currentTask) {
177-
const msg = new Message(Message.KILL, task.msg.data);
138+
const msg = new Message(this.id, Message.KILL, task.msg.data);
178139
const killTask = new Task(this.channel, msg);
179140
await killTask.run();
180141
this.checkReady();
181142
}
182143
}
183144

184145
close() {
185-
this.ws.close();
146+
this.channel.onClientExit(this.id);
147+
}
148+
149+
fork() {
150+
const Session = this.constructor;
151+
return new Session(this.channel);
186152
}
187153

188154
static async new(computeID, config={}, SessionClass=InteractiveSession) {
189-
const session = new SessionClass(computeID, config);
190-
await session.whenConnected();
155+
const channel = await createMessageChannel(computeID, config);
156+
const session = new SessionClass(channel);
191157
return session;
192158
}
193159
}
194160

161+
async function createMessageChannel(computeID, config) {
162+
const address = gmeConfig.extensions.InteractiveComputeHost ||
163+
getDefaultServerURL();
164+
165+
const connectedWs = await new Promise((resolve, reject) => {
166+
const ws = new WebSocket(address);
167+
ws.onopen = () => {
168+
ws.send(JSON.stringify([computeID, config, getGMEToken()]));
169+
ws.onmessage = async (wsMsg) => {
170+
const data = await Task.getMessageData(wsMsg);
171+
172+
const msg = Message.decode(data);
173+
if (msg.type === Message.COMPLETE) {
174+
const err = msg.data;
175+
if (err) {
176+
reject(err);
177+
} else {
178+
resolve(ws);
179+
}
180+
}
181+
};
182+
};
183+
});
184+
185+
return new MessageChannel(connectedWs);
186+
}
187+
188+
function getDefaultServerURL() {
189+
const isSecure = !isNodeJs && location.protocol.includes('s');
190+
const protocol = isSecure ? 'wss' : 'ws';
191+
const defaultHost = isNodeJs ? '127.0.0.1' :
192+
location.origin
193+
.replace(location.protocol + '//', '')
194+
.replace(/:[0-9]+$/, '');
195+
return `${protocol}://${defaultHost}:${gmeConfig.server.port + 1}`;
196+
}
197+
198+
function getGMEToken() {
199+
if (isNodeJs) {
200+
return '';
201+
}
202+
203+
const [, token] = (document.cookie || '').split('=');
204+
return token;
205+
}
206+
195207
function assert(cond, msg) {
196208
if (!cond) {
197209
throw new Error(msg);
@@ -208,6 +220,7 @@ define([
208220
this.ws.onmessage = message => {
209221
this.listeners.forEach(fn => fn(message));
210222
};
223+
this.clients = [];
211224
}
212225

213226
send(data) {
@@ -224,6 +237,26 @@ define([
224237
this.listeners.splice(index, 1);
225238
}
226239
}
240+
241+
isOpen() {
242+
return this.ws.readyState === WebSocket.OPEN;
243+
}
244+
245+
onClientConnect(id) {
246+
this.clients.push(id);
247+
}
248+
249+
onClientExit(id) {
250+
const index = this.clients.indexOf(id);
251+
if (index === -1) {
252+
throw new Error(`Client not found: ${id}`);
253+
}
254+
this.clients.splice(index, 1);
255+
256+
if (this.clients.length === 0) {
257+
this.ws.close();
258+
}
259+
}
227260
}
228261

229262
return InteractiveSession;

‎src/routers/InteractiveCompute/InteractiveCompute.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ class ComputeBroker {
4848
const session = new InteractiveSession(blobClient, client, ws);
4949
this.initSessions.push(session);
5050
} catch (err) {
51-
ws.send(Message.encode(Message.COMPLETE, err.message));
51+
ws.send(Message.encode(-1, Message.COMPLETE, err.message));
5252
this.logger.warn(`Error creating session: ${err}`);
5353
ws.close();
5454
}

‎src/routers/InteractiveCompute/Session.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ class Session extends EventEmitter {
2525
this.workerSocket = socket;
2626
this.emit('connected');
2727

28-
this.clientSocket.send(Message.encode(Message.COMPLETE));
28+
this.clientSocket.send(Message.encode(-1, Message.COMPLETE));
2929
this.queuedMsgs.forEach(msg => this.workerSocket.send(msg));
3030
this.wsChannel = new Channel(this.clientSocket, this.workerSocket);
3131
this.wsChannel.on(Channel.CLOSE, () => this.close());

0 commit comments

Comments
 (0)