Skip to content

Commit 0fc2c6e

Browse files
authored
Merge pull request #667 from balena-io/separate-worker-primary-init
Separate primary/worker initialization to avoid unnecessary imports
2 parents 9a744c7 + fd94ea7 commit 0fc2c6e

9 files changed

Lines changed: 280 additions & 223 deletions

File tree

bin/start.sh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,10 @@ fi
2929

3030
command="$(command -v node)"
3131
args=("--enable-source-maps")
32-
entrypoint="build/src/app.js"
32+
entrypoint="build/src/init.js"
3333
if [[ ${NODE_ENV} = "development" ]]; then
3434
args=('--enable-source-maps' '--import @swc-node/register/esm-register')
35-
entrypoint="src/app.ts"
35+
entrypoint="src/init.ts"
3636
fi
3737

3838
if [[ ! -x ${command} ]] || [[ ! -f ${entrypoint} ]]; then

src/init-instrumentation.ts

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/*
2+
Copyright (C) 2025 Balena Ltd.
3+
4+
This program is free software: you can redistribute it and/or modify
5+
it under the terms of the GNU Affero General Public License as published
6+
by the Free Software Foundation, either version 3 of the License, or
7+
(at your option) any later version.
8+
9+
This program is distributed in the hope that it will be useful,
10+
but WITHOUT ANY WARRANTY; without even the implied warranty of
11+
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12+
GNU Affero General Public License for more details.
13+
14+
You should have received a copy of the GNU Affero General Public License
15+
along with this program. If not, see <https://www.gnu.org/licenses/>.
16+
*/
17+
18+
import { set } from '@balena/es-version';
19+
// Set the desired es version for downstream modules that support it, before we import any
20+
set('es2021');
21+
22+
import { NodeSDK } from '@opentelemetry/sdk-node';
23+
import { ExpressInstrumentation } from '@opentelemetry/instrumentation-express';
24+
import { HttpInstrumentation } from '@opentelemetry/instrumentation-http';
25+
26+
const sdk = new NodeSDK({
27+
instrumentations: [new ExpressInstrumentation(), new HttpInstrumentation()],
28+
});
29+
sdk.start();

src/app.ts renamed to src/init-primary.ts

Lines changed: 89 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,6 @@
1515
along with this program. If not, see <https://www.gnu.org/licenses/>.
1616
*/
1717

18-
import './init.js';
19-
2018
import { metrics } from '@balena/node-metrics-gatherer';
2119
import cluster from 'cluster';
2220
import express from 'express';
@@ -35,15 +33,18 @@ import {
3533
VPN_VERBOSE_LOGS,
3634
} from './utils/config.js';
3735

38-
import proxyWorker from './proxy-worker.js';
39-
import vpnWorker from './vpn-worker.js';
40-
import { intVar } from '@balena/env-parsing';
41-
import { describeMetrics, Metrics } from './utils/metrics.js';
36+
import { describePrimaryMetrics, Metrics } from './utils/metrics.js';
4237
import { service } from './utils/service.js';
4338

39+
if (!cluster.isPrimary) {
40+
throw new Error(
41+
'init-primary should only be imported by the primary cluster process',
42+
);
43+
}
44+
4445
const masterLogger = getLogger('master');
4546

46-
describeMetrics();
47+
describePrimaryMetrics();
4748

4849
export interface BitrateMessage {
4950
type: 'bitrate';
@@ -54,90 +55,91 @@ export interface BitrateMessage {
5455
};
5556
}
5657

57-
if (cluster.isPrimary) {
58-
interface WorkerMetric {
59-
rxBitrate: Array<BitrateMessage['data']['rxBitrate']>;
60-
txBitrate: Array<BitrateMessage['data']['txBitrate']>;
61-
}
62-
const workerMetrics = new Map<string, WorkerMetric>();
63-
64-
let verbose = VPN_VERBOSE_LOGS;
58+
interface WorkerMetric {
59+
rxBitrate: Array<BitrateMessage['data']['rxBitrate']>;
60+
txBitrate: Array<BitrateMessage['data']['txBitrate']>;
61+
}
62+
const workerMetrics = new Map<string, WorkerMetric>();
6563

66-
type WorkerState = {
67-
instanceId: number;
68-
finished: boolean;
69-
};
70-
const workerStates: { [instanceId: number]: WorkerState } = {};
64+
let verbose = VPN_VERBOSE_LOGS;
7165

72-
process.on('SIGUSR2', () => {
73-
masterLogger.notice('caught SIGUSR2, toggling log verbosity');
74-
verbose = !verbose;
75-
for (const clusterWorker of Object.values(cluster.workers ?? {})) {
76-
clusterWorker?.send('toggleVerbosity');
77-
}
78-
});
66+
type WorkerState = {
67+
instanceId: number;
68+
finished: boolean;
69+
};
70+
const workerStates: { [instanceId: number]: WorkerState } = {};
7971

80-
process.on('SIGTERM', () => {
81-
masterLogger.notice('received SIGTERM');
82-
for (const clusterWorker of Object.values(cluster.workers ?? {})) {
83-
clusterWorker?.send('prepareShutdown');
84-
}
85-
masterLogger.notice(
86-
`waiting ${DEFAULT_SIGTERM_TIMEOUT}ms for workers to finish`,
87-
);
88-
});
72+
process.on('SIGUSR2', () => {
73+
masterLogger.notice('caught SIGUSR2, toggling log verbosity');
74+
verbose = !verbose;
75+
for (const clusterWorker of Object.values(cluster.workers ?? {})) {
76+
clusterWorker?.send('toggleVerbosity');
77+
}
78+
});
8979

90-
cluster.on('message', (_worker, { data, type }: BitrateMessage) => {
91-
if (type === 'bitrate') {
92-
let workerMetric = workerMetrics.get(data.uuid);
93-
if (workerMetric == null) {
94-
workerMetric = {
95-
rxBitrate: [],
96-
txBitrate: [],
97-
};
98-
workerMetrics.set(data.uuid, workerMetric);
99-
}
100-
workerMetric.rxBitrate.push(data.rxBitrate);
101-
workerMetric.txBitrate.push(data.txBitrate);
80+
process.on('SIGTERM', () => {
81+
masterLogger.notice('received SIGTERM');
82+
for (const clusterWorker of Object.values(cluster.workers ?? {})) {
83+
clusterWorker?.send('prepareShutdown');
84+
}
85+
masterLogger.notice(
86+
`waiting ${DEFAULT_SIGTERM_TIMEOUT}ms for workers to finish`,
87+
);
88+
});
89+
90+
cluster.on('message', (_worker, { data, type }: BitrateMessage) => {
91+
if (type === 'bitrate') {
92+
let workerMetric = workerMetrics.get(data.uuid);
93+
if (workerMetric == null) {
94+
workerMetric = {
95+
rxBitrate: [],
96+
txBitrate: [],
97+
};
98+
workerMetrics.set(data.uuid, workerMetric);
10299
}
103-
});
104-
105-
cluster.on('message', (_worker, msg: { type: string; data: WorkerState }) => {
106-
const { data, type } = msg;
107-
108-
// worker finished connection draining
109-
if (type === 'drain') {
110-
try {
111-
workerStates[data.instanceId] = data;
112-
const drainCount = Object.keys(workerStates).length;
113-
masterLogger.notice(
114-
`total: ${VPN_INSTANCE_COUNT} drained: ${drainCount}`,
115-
);
116-
for (const key in workerStates) {
117-
if (key != null) {
118-
const value: WorkerState = workerStates[key];
119-
const workerState = value.finished;
120-
masterLogger.notice(`instanceId:${key} finished:${workerState}`);
121-
}
122-
}
123-
if (drainCount >= VPN_INSTANCE_COUNT) {
124-
masterLogger.notice(`all ${drainCount} worker(s) drained`);
125-
process.exit(0);
100+
workerMetric.rxBitrate.push(data.rxBitrate);
101+
workerMetric.txBitrate.push(data.txBitrate);
102+
}
103+
});
104+
105+
cluster.on('message', (_worker, msg: { type: string; data: WorkerState }) => {
106+
const { data, type } = msg;
107+
108+
// worker finished connection draining
109+
if (type === 'drain') {
110+
try {
111+
workerStates[data.instanceId] = data;
112+
const drainCount = Object.keys(workerStates).length;
113+
masterLogger.notice(
114+
`total: ${VPN_INSTANCE_COUNT} drained: ${drainCount}`,
115+
);
116+
for (const key in workerStates) {
117+
if (key != null) {
118+
const value: WorkerState = workerStates[key];
119+
const workerState = value.finished;
120+
masterLogger.notice(`instanceId:${key} finished:${workerState}`);
126121
}
127-
} catch (err) {
128-
masterLogger.warning(`${err} handling message from worker`);
129122
}
130-
} else {
131-
return;
123+
if (drainCount >= VPN_INSTANCE_COUNT) {
124+
masterLogger.notice(`all ${drainCount} worker(s) drained`);
125+
process.exit(0);
126+
}
127+
} catch (err) {
128+
masterLogger.warning(`${err} handling message from worker`);
132129
}
133-
});
134-
135-
masterLogger.notice(
136-
`open-balena-vpn@${VERSION} process started with pid=${process.pid}`,
137-
);
138-
masterLogger.debug('registering as service instance...');
139-
service
140-
.wrap({ ipAddress: VPN_SERVICE_ADDRESS }, async (serviceInstance) => {
130+
} else {
131+
return;
132+
}
133+
});
134+
135+
masterLogger.notice(
136+
`open-balena-vpn@${VERSION} process started with pid=${process.pid}`,
137+
);
138+
masterLogger.debug('registering as service instance...');
139+
try {
140+
await service.wrap(
141+
{ ipAddress: VPN_SERVICE_ADDRESS },
142+
async (serviceInstance) => {
141143
const serviceLogger = getLogger('master', serviceInstance.getId());
142144
serviceLogger.info(
143145
`registered as service instance with id=${serviceInstance.getId()} ipAddress=${VPN_SERVICE_ADDRESS}`,
@@ -232,27 +234,9 @@ if (cluster.isPrimary) {
232234
.listen(8080);
233235

234236
return [app, metrics];
235-
})
236-
.catch((err) => {
237-
console.error('Error starting master:', err);
238-
process.exit(1);
239-
});
240-
}
241-
242-
if (cluster.isWorker) {
243-
// Ensure the prom-client worker listener is registered by instantiating the class
244-
// tslint:disable-next-line:no-unused-expression-chai
245-
new prometheus.AggregatorRegistry();
246-
247-
const instanceId = intVar('WORKER_ID');
248-
const serviceId = intVar('SERVICE_ID');
249-
getLogger('worker', serviceId, instanceId).notice(
250-
`process started with pid=${process.pid}`,
237+
},
251238
);
252-
vpnWorker(instanceId, serviceId)
253-
.then(() => proxyWorker(instanceId, serviceId))
254-
.catch((err) => {
255-
console.error('Error starting worker:', err);
256-
process.exit(1);
257-
});
239+
} catch (err) {
240+
console.error('Error starting master:', err);
241+
process.exit(1);
258242
}

src/init-worker.ts

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
Copyright (C) 2018 Balena Ltd.
3+
4+
This program is free software: you can redistribute it and/or modify
5+
it under the terms of the GNU Affero General Public License as published
6+
by the Free Software Foundation, either version 3 of the License, or
7+
(at your option) any later version.
8+
9+
This program is distributed in the hope that it will be useful,
10+
but WITHOUT ANY WARRANTY; without even the implied warranty of
11+
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12+
GNU Affero General Public License for more details.
13+
14+
You should have received a copy of the GNU Affero General Public License
15+
along with this program. If not, see <https://www.gnu.org/licenses/>.
16+
*/
17+
18+
import cluster from 'cluster';
19+
import prometheus from 'prom-client';
20+
21+
import { getLogger } from './utils/index.js';
22+
23+
import proxyWorker from './proxy-worker.js';
24+
import vpnWorker from './vpn-worker.js';
25+
import { intVar } from '@balena/env-parsing';
26+
import { describeWorkerMetrics } from './utils/metrics.js';
27+
28+
if (!cluster.isWorker) {
29+
throw new Error('init-worker should only be imported by a worker process');
30+
}
31+
32+
describeWorkerMetrics();
33+
34+
// Ensure the prom-client worker listener is registered by instantiating the class
35+
// tslint:disable-next-line:no-unused-expression-chai
36+
new prometheus.AggregatorRegistry();
37+
38+
const instanceId = intVar('WORKER_ID');
39+
const serviceId = intVar('SERVICE_ID');
40+
getLogger('worker', serviceId, instanceId).notice(
41+
`process started with pid=${process.pid}`,
42+
);
43+
try {
44+
await vpnWorker(instanceId, serviceId);
45+
proxyWorker(instanceId, serviceId);
46+
} catch (err) {
47+
console.error('Error starting worker:', err);
48+
process.exit(1);
49+
}

src/init.ts

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
Copyright (C) 2025 Balena Ltd.
2+
Copyright (C) 2018 Balena Ltd.
33
44
This program is free software: you can redistribute it and/or modify
55
it under the terms of the GNU Affero General Public License as published
@@ -15,15 +15,14 @@
1515
along with this program. If not, see <https://www.gnu.org/licenses/>.
1616
*/
1717

18-
import { set } from '@balena/es-version';
19-
// Set the desired es version for downstream modules that support it, before we import any
20-
set('es2021');
18+
import './init-instrumentation.js';
2119

22-
import { NodeSDK } from '@opentelemetry/sdk-node';
23-
import { ExpressInstrumentation } from '@opentelemetry/instrumentation-express';
24-
import { HttpInstrumentation } from '@opentelemetry/instrumentation-http';
20+
import cluster from 'cluster';
2521

26-
const sdk = new NodeSDK({
27-
instrumentations: [new ExpressInstrumentation(), new HttpInstrumentation()],
28-
});
29-
sdk.start();
22+
if (cluster.isPrimary) {
23+
await import('./init-primary.js');
24+
}
25+
26+
if (cluster.isWorker) {
27+
await import('./init-worker.js');
28+
}

src/proxy-worker.ts

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -302,8 +302,6 @@ const worker = (instanceId: number, serviceId: number) => {
302302
`process started with pid=${process.pid}`,
303303
);
304304

305-
const tunnel = new Tunnel(instanceId, serviceId);
306-
tunnel.start(VPN_CONNECT_PROXY_PORT);
307-
return tunnel;
305+
new Tunnel(instanceId, serviceId).start(VPN_CONNECT_PROXY_PORT);
308306
};
309307
export default worker;

0 commit comments

Comments
 (0)