Skip to content

Commit 70dae85

Browse files
jacobwgillespiejacobwgillespie
authored
Merge pull request #123 from depot/health-stream
2 parents 02ad4e8 + 0dbfc57 commit 70dae85

File tree

7 files changed

+148
-50
lines changed

7 files changed

+148
-50
lines changed

‎proto/depot/cloud/v3/machine.proto

+9
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package depot.cloud.v3;
55
service MachineService {
66
rpc RegisterMachine(RegisterMachineRequest) returns (stream RegisterMachineResponse);
77
rpc PingMachineHealth(PingMachineHealthRequest) returns (PingMachineHealthResponse);
8+
rpc ReportMachineHealth(stream ReportMachineHealthRequest) returns (ReportMachineHealthResponse);
89
rpc Usage(UsageRequest) returns (UsageResponse);
910
}
1011

@@ -134,6 +135,14 @@ message PingMachineHealthResponse {
134135
bool should_terminate = 1;
135136
}
136137

138+
message ReportMachineHealthRequest {
139+
repeated DiskSpace disks = 1;
140+
}
141+
142+
message ReportMachineHealthResponse {
143+
bool should_terminate = 1;
144+
}
145+
137146
message Cert {
138147
string cert = 1;
139148
string key = 2;

‎src/gen/ts/depot/cloud/v3/machine_connect.ts

+12-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// @generated by protoc-gen-connect-es v0.12.0 with parameter "target=ts,import_extension=none"
1+
// @generated by protoc-gen-connect-es v1.4.0 with parameter "target=ts,import_extension=none"
22
// @generated from file depot/cloud/v3/machine.proto (package depot.cloud.v3, syntax proto3)
33
/* eslint-disable */
44
// @ts-nocheck
@@ -9,6 +9,8 @@ import {
99
PingMachineHealthResponse,
1010
RegisterMachineRequest,
1111
RegisterMachineResponse,
12+
ReportMachineHealthRequest,
13+
ReportMachineHealthResponse,
1214
UsageRequest,
1315
UsageResponse,
1416
} from './machine_pb'
@@ -37,6 +39,15 @@ export const MachineService = {
3739
O: PingMachineHealthResponse,
3840
kind: MethodKind.Unary,
3941
},
42+
/**
43+
* @generated from rpc depot.cloud.v3.MachineService.ReportMachineHealth
44+
*/
45+
reportMachineHealth: {
46+
name: 'ReportMachineHealth',
47+
I: ReportMachineHealthRequest,
48+
O: ReportMachineHealthResponse,
49+
kind: MethodKind.ClientStreaming,
50+
},
4051
/**
4152
* @generated from rpc depot.cloud.v3.MachineService.Usage
4253
*/

‎src/gen/ts/depot/cloud/v3/machine_pb.ts

+81-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// @generated by protoc-gen-es v1.3.0 with parameter "target=ts,import_extension=none"
1+
// @generated by protoc-gen-es v1.10.0 with parameter "target=ts,import_extension=none"
22
// @generated from file depot/cloud/v3/machine.proto (package depot.cloud.v3, syntax proto3)
33
/* eslint-disable */
44
// @ts-nocheck
@@ -900,6 +900,86 @@ export class PingMachineHealthResponse extends Message<PingMachineHealthResponse
900900
}
901901
}
902902

903+
/**
904+
* @generated from message depot.cloud.v3.ReportMachineHealthRequest
905+
*/
906+
export class ReportMachineHealthRequest extends Message<ReportMachineHealthRequest> {
907+
/**
908+
* @generated from field: repeated depot.cloud.v3.DiskSpace disks = 1;
909+
*/
910+
disks: DiskSpace[] = []
911+
912+
constructor(data?: PartialMessage<ReportMachineHealthRequest>) {
913+
super()
914+
proto3.util.initPartial(data, this)
915+
}
916+
917+
static readonly runtime: typeof proto3 = proto3
918+
static readonly typeName = 'depot.cloud.v3.ReportMachineHealthRequest'
919+
static readonly fields: FieldList = proto3.util.newFieldList(() => [
920+
{no: 1, name: 'disks', kind: 'message', T: DiskSpace, repeated: true},
921+
])
922+
923+
static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): ReportMachineHealthRequest {
924+
return new ReportMachineHealthRequest().fromBinary(bytes, options)
925+
}
926+
927+
static fromJson(jsonValue: JsonValue, options?: Partial<JsonReadOptions>): ReportMachineHealthRequest {
928+
return new ReportMachineHealthRequest().fromJson(jsonValue, options)
929+
}
930+
931+
static fromJsonString(jsonString: string, options?: Partial<JsonReadOptions>): ReportMachineHealthRequest {
932+
return new ReportMachineHealthRequest().fromJsonString(jsonString, options)
933+
}
934+
935+
static equals(
936+
a: ReportMachineHealthRequest | PlainMessage<ReportMachineHealthRequest> | undefined,
937+
b: ReportMachineHealthRequest | PlainMessage<ReportMachineHealthRequest> | undefined,
938+
): boolean {
939+
return proto3.util.equals(ReportMachineHealthRequest, a, b)
940+
}
941+
}
942+
943+
/**
944+
* @generated from message depot.cloud.v3.ReportMachineHealthResponse
945+
*/
946+
export class ReportMachineHealthResponse extends Message<ReportMachineHealthResponse> {
947+
/**
948+
* @generated from field: bool should_terminate = 1;
949+
*/
950+
shouldTerminate = false
951+
952+
constructor(data?: PartialMessage<ReportMachineHealthResponse>) {
953+
super()
954+
proto3.util.initPartial(data, this)
955+
}
956+
957+
static readonly runtime: typeof proto3 = proto3
958+
static readonly typeName = 'depot.cloud.v3.ReportMachineHealthResponse'
959+
static readonly fields: FieldList = proto3.util.newFieldList(() => [
960+
{no: 1, name: 'should_terminate', kind: 'scalar', T: 8 /* ScalarType.BOOL */},
961+
])
962+
963+
static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): ReportMachineHealthResponse {
964+
return new ReportMachineHealthResponse().fromBinary(bytes, options)
965+
}
966+
967+
static fromJson(jsonValue: JsonValue, options?: Partial<JsonReadOptions>): ReportMachineHealthResponse {
968+
return new ReportMachineHealthResponse().fromJson(jsonValue, options)
969+
}
970+
971+
static fromJsonString(jsonString: string, options?: Partial<JsonReadOptions>): ReportMachineHealthResponse {
972+
return new ReportMachineHealthResponse().fromJsonString(jsonString, options)
973+
}
974+
975+
static equals(
976+
a: ReportMachineHealthResponse | PlainMessage<ReportMachineHealthResponse> | undefined,
977+
b: ReportMachineHealthResponse | PlainMessage<ReportMachineHealthResponse> | undefined,
978+
): boolean {
979+
return proto3.util.equals(ReportMachineHealthResponse, a, b)
980+
}
981+
}
982+
903983
/**
904984
* @generated from message depot.cloud.v3.Cert
905985
*/

‎src/tasks/buildkit.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -258,7 +258,7 @@ keepBytes = ${cacheSizeBytes}
258258
try {
259259
await Promise.all([
260260
buildkit,
261-
reportHealth({machineId, signal, headers, path: rootDir}),
261+
reportHealth({signal, headers, path: rootDir}),
262262
reportUsage({machineId, signal, headers}),
263263
])
264264
} catch (error) {

‎src/tasks/engine.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ export async function startEngine(message: RegisterMachineResponse, task: Regist
110110
try {
111111
await Promise.all([
112112
engine,
113-
reportEngineHealth({machineId, signal, headers, path: '/var/lib/engine'}),
113+
reportEngineHealth({signal, headers, path: '/var/lib/engine'}),
114114
// reportUsage({machineId, signal, headers}),
115115
])
116116
} catch (error) {

‎src/tasks/engineHealth.ts

+22-23
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,40 @@
1+
import {PlainMessage} from '@bufbuild/protobuf'
12
import {execa} from 'execa'
3+
import {DiskSpace} from '../gen/ts/depot/cloud/v3/machine_pb'
24
import {sleep} from '../utils/common'
35
import {stats} from '../utils/disk'
46
import {client} from '../utils/grpc'
57

68
export interface ReportHealthParams {
7-
machineId: string
89
signal: AbortSignal
910
headers: HeadersInit
1011
path: string
1112
}
1213

13-
export async function reportEngineHealth({machineId, signal, headers, path}: ReportHealthParams) {
14-
while (true) {
15-
if (signal.aborted) return
16-
14+
export async function reportEngineHealth({signal, headers, path}: ReportHealthParams) {
15+
while (!signal.aborted) {
1716
await waitForWorkers(signal)
1817

1918
try {
20-
while (true) {
21-
if (signal.aborted) return
22-
23-
const disk_stats = await stats(path)
24-
const disk_space = disk_stats
25-
? [
26-
{
27-
path,
28-
freeMb: disk_stats.freeMb,
29-
totalMb: disk_stats.totalMb,
30-
freeInodes: disk_stats.freeInodes,
31-
totalInodes: disk_stats.totalInodes,
32-
},
33-
]
34-
: undefined
35-
36-
await client.pingMachineHealth({machineId, disks: disk_space}, {headers, signal})
37-
await sleep(1000)
19+
async function* stream() {
20+
while (!signal.aborted) {
21+
const diskStats = await stats(path)
22+
if (diskStats) {
23+
const diskSpace: PlainMessage<DiskSpace> = {
24+
path,
25+
freeMb: diskStats.freeMb,
26+
totalMb: diskStats.totalMb,
27+
freeInodes: diskStats.freeInodes,
28+
totalInodes: diskStats.totalInodes,
29+
}
30+
yield {disks: [diskSpace]}
31+
}
32+
await sleep(1000)
33+
}
3834
}
35+
36+
const res = await client.reportMachineHealth(stream(), {headers, signal})
37+
if (res.shouldTerminate) return
3938
} catch (error) {
4039
console.log('Error reporting health:', error)
4140
}

‎src/tasks/health.ts

+22-23
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,40 @@
1+
import {PlainMessage} from '@bufbuild/protobuf'
12
import {execa} from 'execa'
3+
import {DiskSpace} from '../gen/ts/depot/cloud/v3/machine_pb'
24
import {sleep} from '../utils/common'
35
import {stats} from '../utils/disk'
46
import {client} from '../utils/grpc'
57

68
export interface ReportHealthParams {
7-
machineId: string
89
signal: AbortSignal
910
headers: HeadersInit
1011
path: string
1112
}
1213

13-
export async function reportHealth({machineId, signal, headers, path}: ReportHealthParams) {
14-
while (true) {
15-
if (signal.aborted) return
16-
14+
export async function reportHealth({signal, headers, path}: ReportHealthParams) {
15+
while (!signal.aborted) {
1716
await waitForBuildKitWorkers(signal)
1817

1918
try {
20-
while (true) {
21-
if (signal.aborted) return
22-
23-
const disk_stats = await stats(path)
24-
const disk_space = disk_stats
25-
? [
26-
{
27-
path,
28-
freeMb: disk_stats.freeMb,
29-
totalMb: disk_stats.totalMb,
30-
freeInodes: disk_stats.freeInodes,
31-
totalInodes: disk_stats.totalInodes,
32-
},
33-
]
34-
: undefined
35-
36-
await client.pingMachineHealth({machineId, disks: disk_space}, {headers, signal})
37-
await sleep(1000)
19+
async function* stream() {
20+
while (!signal.aborted) {
21+
const diskStats = await stats(path)
22+
if (diskStats) {
23+
const diskSpace: PlainMessage<DiskSpace> = {
24+
path,
25+
freeMb: diskStats.freeMb,
26+
totalMb: diskStats.totalMb,
27+
freeInodes: diskStats.freeInodes,
28+
totalInodes: diskStats.totalInodes,
29+
}
30+
yield {disks: [diskSpace]}
31+
}
32+
await sleep(1000)
33+
}
3834
}
35+
36+
const res = await client.reportMachineHealth(stream(), {headers, signal})
37+
if (res.shouldTerminate) return
3938
} catch (error) {
4039
console.log('Error reporting health:', error)
4140
}

0 commit comments

Comments
 (0)