Skip to content

Commit e90a076

Browse files
BrunoQuaresmaBrunoQuaresma
authored
fix(site): Fix websocket connections (#7187)
1 parent 7fa1112 commit e90a076

File tree

6 files changed

+229
-116
lines changed

6 files changed

+229
-116
lines changed

‎site/src/api/api.ts

Lines changed: 100 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1106,25 +1106,107 @@ export const watchAgentMetadata = (agentId: string): EventSource => {
11061106
)
11071107
}
11081108

1109-
export const watchBuildLogs = (
1109+
type WatchBuildLogsByTemplateVersionIdOptions = {
1110+
after?: number
1111+
onMessage: (log: TypesGen.ProvisionerJobLog) => void
1112+
onDone: () => void
1113+
onError: (error: Error) => void
1114+
}
1115+
export const watchBuildLogsByTemplateVersionId = (
11101116
versionId: string,
1111-
onMessage: (log: TypesGen.ProvisionerJobLog) => void,
1117+
{
1118+
onMessage,
1119+
onDone,
1120+
onError,
1121+
after,
1122+
}: WatchBuildLogsByTemplateVersionIdOptions,
11121123
) => {
1113-
return new Promise<void>((resolve, reject) => {
1114-
const proto = location.protocol === "https:" ? "wss:" : "ws:"
1115-
const socket = new WebSocket(
1116-
`${proto}//${location.host}/api/v2/templateversions/${versionId}/logs?follow=true`,
1117-
)
1118-
socket.binaryType = "blob"
1119-
socket.addEventListener("message", (event) =>
1120-
onMessage(JSON.parse(event.data) as TypesGen.ProvisionerJobLog),
1121-
)
1122-
socket.addEventListener("error", () => {
1123-
reject(new Error("Connection for logs failed."))
1124-
})
1125-
socket.addEventListener("close", () => {
1126-
// When the socket closes, logs have finished streaming!
1127-
resolve()
1128-
})
1124+
const searchParams = new URLSearchParams({ follow: "true" })
1125+
if (after !== undefined) {
1126+
searchParams.append("after", after.toString())
1127+
}
1128+
const proto = location.protocol === "https:" ? "wss:" : "ws:"
1129+
const socket = new WebSocket(
1130+
`${proto}//${
1131+
location.host
1132+
}/api/v2/templateversions/${versionId}/logs?${searchParams.toString()}`,
1133+
)
1134+
socket.binaryType = "blob"
1135+
socket.addEventListener("message", (event) =>
1136+
onMessage(JSON.parse(event.data) as TypesGen.ProvisionerJobLog),
1137+
)
1138+
socket.addEventListener("error", () => {
1139+
onError(new Error("Connection for logs failed."))
1140+
socket.close()
1141+
})
1142+
socket.addEventListener("close", () => {
1143+
// When the socket closes, logs have finished streaming!
1144+
onDone()
1145+
})
1146+
return socket
1147+
}
1148+
1149+
type WatchStartupLogsOptions = {
1150+
after: number
1151+
onMessage: (logs: TypesGen.WorkspaceAgentStartupLog[]) => void
1152+
onDone: () => void
1153+
onError: (error: Error) => void
1154+
}
1155+
1156+
export const watchStartupLogs = (
1157+
agentId: string,
1158+
{ after, onMessage, onDone, onError }: WatchStartupLogsOptions,
1159+
) => {
1160+
const proto = location.protocol === "https:" ? "wss:" : "ws:"
1161+
const socket = new WebSocket(
1162+
`${proto}//${location.host}/api/v2/workspaceagents/${agentId}/startup-logs?follow&after=${after}`,
1163+
)
1164+
socket.binaryType = "blob"
1165+
socket.addEventListener("message", (event) => {
1166+
const logs = JSON.parse(event.data) as TypesGen.WorkspaceAgentStartupLog[]
1167+
onMessage(logs)
1168+
})
1169+
socket.addEventListener("error", () => {
1170+
onError(new Error("socket errored"))
1171+
})
1172+
socket.addEventListener("close", () => {
1173+
onDone()
1174+
})
1175+
1176+
return socket
1177+
}
1178+
1179+
type WatchBuildLogsByBuildIdOptions = {
1180+
after?: number
1181+
onMessage: (log: TypesGen.ProvisionerJobLog) => void
1182+
onDone: () => void
1183+
onError: (error: Error) => void
1184+
}
1185+
export const watchBuildLogsByBuildId = (
1186+
buildId: string,
1187+
{ onMessage, onDone, onError, after }: WatchBuildLogsByBuildIdOptions,
1188+
) => {
1189+
const searchParams = new URLSearchParams({ follow: "true" })
1190+
if (after !== undefined) {
1191+
searchParams.append("after", after.toString())
1192+
}
1193+
const proto = location.protocol === "https:" ? "wss:" : "ws:"
1194+
const socket = new WebSocket(
1195+
`${proto}//${
1196+
location.host
1197+
}/api/v2/workspacebuilds/${buildId}/logs?${searchParams.toString()}`,
1198+
)
1199+
socket.binaryType = "blob"
1200+
socket.addEventListener("message", (event) =>
1201+
onMessage(JSON.parse(event.data) as TypesGen.ProvisionerJobLog),
1202+
)
1203+
socket.addEventListener("error", () => {
1204+
onError(new Error("Connection for logs failed."))
1205+
socket.close()
1206+
})
1207+
socket.addEventListener("close", () => {
1208+
// When the socket closes, logs have finished streaming!
1209+
onDone()
11291210
})
1211+
return socket
11301212
}

‎site/src/pages/TemplateVersionPage/TemplateVersionEditorPage/TemplateVersionEditorPage.test.tsx

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,13 @@ test("Use custom name and set it as active when publishing", async () => {
3636
jest
3737
.spyOn(api, "getTemplateVersion")
3838
.mockResolvedValue({ ...MockTemplateVersion, id: "new-version-id" })
39-
jest.spyOn(api, "watchBuildLogs").mockImplementation((_, onMessage) => {
40-
onMessage(MockWorkspaceBuildLogs[0])
41-
return Promise.resolve()
42-
})
39+
jest
40+
.spyOn(api, "watchBuildLogsByTemplateVersionId")
41+
.mockImplementation((_, options) => {
42+
options.onMessage(MockWorkspaceBuildLogs[0])
43+
options.onDone()
44+
return jest.fn() as never
45+
})
4346
const buildButton = within(topbar).getByRole("button", {
4447
name: "Build template",
4548
})
@@ -97,10 +100,13 @@ test("Do not mark as active if promote is not checked", async () => {
97100
jest
98101
.spyOn(api, "getTemplateVersion")
99102
.mockResolvedValue({ ...MockTemplateVersion, id: "new-version-id" })
100-
jest.spyOn(api, "watchBuildLogs").mockImplementation((_, onMessage) => {
101-
onMessage(MockWorkspaceBuildLogs[0])
102-
return Promise.resolve()
103-
})
103+
jest
104+
.spyOn(api, "watchBuildLogsByTemplateVersionId")
105+
.mockImplementation((_, options) => {
106+
options.onMessage(MockWorkspaceBuildLogs[0])
107+
options.onDone()
108+
return jest.fn() as never
109+
})
104110
const buildButton = within(topbar).getByRole("button", {
105111
name: "Build template",
106112
})
@@ -153,10 +159,13 @@ test("Patch request is not send when the name is not updated", async () => {
153159
jest
154160
.spyOn(api, "getTemplateVersion")
155161
.mockResolvedValue({ ...MockTemplateVersion, id: "new-version-id" })
156-
jest.spyOn(api, "watchBuildLogs").mockImplementation((_, onMessage) => {
157-
onMessage(MockWorkspaceBuildLogs[0])
158-
return Promise.resolve()
159-
})
162+
jest
163+
.spyOn(api, "watchBuildLogsByTemplateVersionId")
164+
.mockImplementation((_, options) => {
165+
options.onMessage(MockWorkspaceBuildLogs[0])
166+
options.onDone()
167+
return jest.fn() as never
168+
})
160169
const buildButton = within(topbar).getByRole("button", {
161170
name: "Build template",
162171
})

‎site/src/xServices/templateVersionEditor/templateVersionEditorXService.ts

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ export const templateVersionEditorMachine = createMachine(
4949
| { type: "SET_MISSING_VARIABLE_VALUES"; values: VariableValue[] }
5050
| { type: "CANCEL_MISSING_VARIABLE_VALUES" }
5151
| { type: "ADD_BUILD_LOG"; log: ProvisionerJobLog }
52+
| { type: "BUILD_DONE" }
5253
| { type: "PUBLISH" }
5354
| ({ type: "CONFIRM_PUBLISH" } & PublishVersionData)
5455
| { type: "CANCEL_PUBLISH" },
@@ -157,14 +158,12 @@ export const templateVersionEditorMachine = createMachine(
157158
invoke: {
158159
id: "watchBuildLogs",
159160
src: "watchBuildLogs",
160-
onDone: {
161-
target: "fetchingVersion",
162-
},
163161
},
164162
on: {
165163
ADD_BUILD_LOG: {
166164
actions: "addBuildLog",
167165
},
166+
BUILD_DONE: "fetchingVersion",
168167
CANCEL_VERSION: {
169168
target: "cancelingBuild",
170169
},
@@ -360,9 +359,21 @@ export const templateVersionEditorMachine = createMachine(
360359
throw new Error("version must be set")
361360
}
362361

363-
return API.watchBuildLogs(version.id, (log) => {
364-
callback({ type: "ADD_BUILD_LOG", log })
362+
const socket = API.watchBuildLogsByTemplateVersionId(version.id, {
363+
onMessage: (log) => {
364+
callback({ type: "ADD_BUILD_LOG", log })
365+
},
366+
onDone: () => {
367+
callback({ type: "BUILD_DONE" })
368+
},
369+
onError: (error) => {
370+
console.error(error)
371+
},
365372
})
373+
374+
return () => {
375+
socket.close()
376+
}
366377
},
367378
getResources: (ctx) => {
368379
if (!ctx.version) {

‎site/src/xServices/workspace/workspaceXService.ts

Lines changed: 27 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -117,31 +117,32 @@ export const checks = {
117117
const permissionsToCheck = (
118118
workspace: TypesGen.Workspace,
119119
template: TypesGen.Template,
120-
) => ({
121-
[checks.readWorkspace]: {
122-
object: {
123-
resource_type: "workspace",
124-
resource_id: workspace.id,
125-
owner_id: workspace.owner_id,
120+
) =>
121+
({
122+
[checks.readWorkspace]: {
123+
object: {
124+
resource_type: "workspace",
125+
resource_id: workspace.id,
126+
owner_id: workspace.owner_id,
127+
},
128+
action: "read",
126129
},
127-
action: "read",
128-
},
129-
[checks.updateWorkspace]: {
130-
object: {
131-
resource_type: "workspace",
132-
resource_id: workspace.id,
133-
owner_id: workspace.owner_id,
130+
[checks.updateWorkspace]: {
131+
object: {
132+
resource_type: "workspace",
133+
resource_id: workspace.id,
134+
owner_id: workspace.owner_id,
135+
},
136+
action: "update",
134137
},
135-
action: "update",
136-
},
137-
[checks.updateTemplate]: {
138-
object: {
139-
resource_type: "template",
140-
resource_id: template.id,
138+
[checks.updateTemplate]: {
139+
object: {
140+
resource_type: "template",
141+
resource_id: template.id,
142+
},
143+
action: "update",
141144
},
142-
action: "update",
143-
},
144-
})
145+
} as const)
145146

146147
export const workspaceMachine = createMachine(
147148
{
@@ -851,6 +852,10 @@ export const workspaceMachine = createMachine(
851852
context.eventSource.onerror = () => {
852853
send({ type: "EVENT_SOURCE_ERROR", error: "sse error" })
853854
}
855+
856+
return () => {
857+
context.eventSource?.close()
858+
}
854859
},
855860
getBuilds: async (context) => {
856861
if (context.workspace) {

‎site/src/xServices/workspaceAgentLogs/workspaceAgentLogsXService.ts

Lines changed: 30 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ export const workspaceAgentLogsMachine = createMachine(
2222
}
2323
| {
2424
type: "FETCH_STARTUP_LOGS"
25+
}
26+
| {
27+
type: "STARTUP_DONE"
2528
},
2629
context: {} as {
2730
agentID: string
@@ -56,16 +59,19 @@ export const workspaceAgentLogsMachine = createMachine(
5659
id: "streamStartupLogs",
5760
src: "streamStartupLogs",
5861
},
62+
on: {
63+
ADD_STARTUP_LOGS: {
64+
actions: "addStartupLogs",
65+
},
66+
STARTUP_DONE: {
67+
target: "loaded",
68+
},
69+
},
5970
},
6071
loaded: {
6172
type: "final",
6273
},
6374
},
64-
on: {
65-
ADD_STARTUP_LOGS: {
66-
actions: "addStartupLogs",
67-
},
68-
},
6975
},
7076
{
7177
services: {
@@ -79,20 +85,14 @@ export const workspaceAgentLogsMachine = createMachine(
7985
})),
8086
),
8187
streamStartupLogs: (ctx) => async (callback) => {
82-
return new Promise<void>((resolve, reject) => {
83-
const proto = location.protocol === "https:" ? "wss:" : "ws:"
84-
let after = 0
85-
if (ctx.startupLogs && ctx.startupLogs.length > 0) {
86-
after = ctx.startupLogs[ctx.startupLogs.length - 1].id
87-
}
88-
const socket = new WebSocket(
89-
`${proto}//${location.host}/api/v2/workspaceagents/${ctx.agentID}/startup-logs?follow&after=${after}`,
90-
)
91-
socket.binaryType = "blob"
92-
socket.addEventListener("message", (event) => {
93-
const logs = JSON.parse(
94-
event.data,
95-
) as TypesGen.WorkspaceAgentStartupLog[]
88+
let after = 0
89+
if (ctx.startupLogs && ctx.startupLogs.length > 0) {
90+
after = ctx.startupLogs[ctx.startupLogs.length - 1].id
91+
}
92+
93+
const socket = API.watchStartupLogs(ctx.agentID, {
94+
after,
95+
onMessage: (logs) => {
9696
callback({
9797
type: "ADD_STARTUP_LOGS",
9898
logs: logs.map((log) => ({
@@ -102,14 +102,18 @@ export const workspaceAgentLogsMachine = createMachine(
102102
time: log.created_at,
103103
})),
104104
})
105-
})
106-
socket.addEventListener("error", () => {
107-
reject(new Error("socket errored"))
108-
})
109-
socket.addEventListener("open", () => {
110-
resolve()
111-
})
105+
},
106+
onDone: () => {
107+
callback({ type: "STARTUP_DONE" })
108+
},
109+
onError: (error) => {
110+
console.error(error)
111+
},
112112
})
113+
114+
return () => {
115+
socket.close()
116+
}
113117
},
114118
},
115119
actions: {

0 commit comments

Comments
 (0)