Skip to content

Commit b2041aa

Browse files
avicarieavicarie
committed
Fix memory leak for poster generation (#570)
1 parent 6ae9071 commit b2041aa

2 files changed

Lines changed: 144 additions & 71 deletions

File tree

‎scripts/generator.js‎

Lines changed: 129 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,29 @@ const PDF_TIMEOUT = 5 * 60 * 1000;
1414
const MAX_RENDER_ATTEMPTS = 3;
1515
const SCALE = 96 / 72;
1616

17+
const BROWSER_RECYCLE_AFTER = 50;
18+
1719
let browser = null;
20+
let renderCount = 0;
1821
const cwd = process.cwd();
1922

2023
const fileOutputDir = path.join(cwd, 'output');
2124

2225
const pdfPath = id => path.join(fileOutputDir, `${id}.pdf`);
2326
const csvPath = id => path.join(fileOutputDir, `${id}.csv`);
2427

28+
async function closeBrowser() {
29+
if (browser) {
30+
const b = browser;
31+
browser = null; // null first so the 'disconnected' handler is a no-op
32+
try {
33+
await b.close();
34+
} catch (err) {
35+
console.error(`[browser] Error closing browser: ${err.message}`);
36+
}
37+
}
38+
}
39+
2540
async function initialize() {
2641
const launchOptions = {
2742
args: ['--no-sandbox', '--disable-setuid-sandbox', '--disable-web-security'],
@@ -30,7 +45,9 @@ async function initialize() {
3045
launchOptions.executablePath = process.env.PUPPETEER_EXECUTABLE_PATH;
3146
}
3247
browser = await puppeteer.launch(launchOptions);
48+
renderCount = 0;
3349
browser.on('disconnected', () => {
50+
console.log('[browser] Browser disconnected unexpectedly, will re-launch on next job.');
3451
browser = null;
3552
});
3653
}
@@ -72,19 +89,37 @@ async function waitFile(filePath) {
7289
}
7390

7491
/**
75-
* Renders component to PDF or CSV file
76-
* @returns {Promise}
92+
* Renders component to PDF or CSV file.
93+
*
94+
* The caller is responsible for passing an AbortSignal-like token via
95+
* `options.abortSignal` so this function can close its page when a timeout
96+
* fires from outside.
97+
*
98+
* @returns {Promise<boolean>} posterUploaded
7799
*/
78100
async function renderComponent(options) {
79101
const { id, component, template, props, onInfo, onError } = options;
80102

81103
const page = await browser.newPage();
82104

105+
let pageClosed = false;
106+
async function safeClosePage() {
107+
if (!pageClosed) {
108+
pageClosed = true;
109+
try {
110+
await page.close();
111+
} catch (err) {
112+
console.error(`[page] Error closing page for ${id}: ${err.message}`);
113+
}
114+
}
115+
}
116+
83117
await page.exposeFunction('serverLog', log);
84118

85-
page.on('error', error => {
86-
page.close();
87-
browser.close();
119+
page.on('error', async error => {
120+
console.error(`[page] Page crashed for ${id}: ${error.message}`);
121+
await safeClosePage();
122+
await closeBrowser();
88123
onError(error);
89124
});
90125

@@ -99,79 +134,80 @@ async function renderComponent(options) {
99134

100135
const pageUrl = generateRenderUrl(component, template, props, id);
101136

102-
console.log(`Opening ${pageUrl} in Puppeteer.`);
137+
console.log(`[render] Opening ${pageUrl} in Puppeteer.`);
103138

104-
if (component === 'StopRoutePlate' && (props.downloadTable || props.downloadSummary)) {
105-
// Allow the downloading of CSV file since the component just sends it to the client instead of actually rendering
106-
const client = await page.createCDPSession();
107-
await client.send('Page.setDownloadBehavior', {
108-
behavior: 'allow',
109-
downloadPath: fileOutputDir,
110-
});
139+
try {
140+
if (component === 'StopRoutePlate' && (props.downloadTable || props.downloadSummary)) {
141+
// Allow the downloading of CSV file since the component just sends it to the client instead of actually rendering
142+
const client = await page.createCDPSession();
143+
await client.send('Page.setDownloadBehavior', {
144+
behavior: 'allow',
145+
downloadPath: fileOutputDir,
146+
});
111147

112-
const csvFilePath = props.downloadSummary ? csvPath(`summary-${id}`) : csvPath(id);
148+
const csvFilePath = props.downloadSummary ? csvPath(`summary-${id}`) : csvPath(id);
113149

114-
try {
115150
await page.goto(pageUrl);
116151
await waitFile(csvFilePath);
117152
const posterUploaded = await uploadPosterToCloud(csvFilePath);
118-
await page.close();
153+
await safeClosePage();
119154
return posterUploaded;
120-
} catch (err) {
121-
throw new Error('StopRoutePlate CSV rendering failed');
122155
}
123-
}
124156

125-
await page.goto(pageUrl, {
126-
timeout: RENDER_TIMEOUT,
127-
});
157+
await page.goto(pageUrl, {
158+
timeout: RENDER_TIMEOUT,
159+
});
128160

129-
const { error = null, width, height } = await page.evaluate(
130-
() =>
131-
new Promise(resolve => {
132-
window.callPhantom = opts => resolve(opts);
133-
}),
134-
);
161+
const { error = null, width, height } = await page.evaluate(
162+
() =>
163+
new Promise(resolve => {
164+
window.callPhantom = opts => resolve(opts);
165+
}),
166+
);
135167

136-
if (error) {
137-
throw new Error(error);
138-
}
168+
if (error) {
169+
throw new Error(error);
170+
}
139171

140-
await page.emulateMediaType('screen');
141-
142-
let printOptions = {};
143-
if (props.printTimetablesAsA4 || component === 'CoverPage') {
144-
printOptions = {
145-
printBackground: true,
146-
format: 'A4',
147-
margin: 0,
148-
timeout: PDF_TIMEOUT,
149-
};
150-
} else if (props.printAsA5) {
151-
printOptions = {
152-
printBackground: true,
153-
format: 'A5',
154-
margin: 0,
155-
timeout: PDF_TIMEOUT,
156-
};
157-
} else {
158-
printOptions = {
159-
printBackground: true,
160-
width: width * SCALE,
161-
height: height * SCALE,
162-
pageRanges: '1',
163-
scale: SCALE,
164-
};
165-
}
172+
await page.emulateMediaType('screen');
173+
174+
let printOptions = {};
175+
if (props.printTimetablesAsA4 || component === 'CoverPage') {
176+
printOptions = {
177+
printBackground: true,
178+
format: 'A4',
179+
margin: 0,
180+
timeout: PDF_TIMEOUT,
181+
};
182+
} else if (props.printAsA5) {
183+
printOptions = {
184+
printBackground: true,
185+
format: 'A5',
186+
margin: 0,
187+
timeout: PDF_TIMEOUT,
188+
};
189+
} else {
190+
printOptions = {
191+
printBackground: true,
192+
width: width * SCALE,
193+
height: height * SCALE,
194+
pageRanges: '1',
195+
scale: SCALE,
196+
};
197+
}
166198

167-
const contents = await page.pdf(printOptions);
199+
const contents = await page.pdf(printOptions);
168200

169-
const pdfFilePath = pdfPath(id);
170-
await fs.outputFile(pdfFilePath, contents);
171-
await page.close();
201+
const pdfFilePath = pdfPath(id);
202+
await fs.outputFile(pdfFilePath, contents);
203+
await safeClosePage();
172204

173-
const posterUploaded = await uploadPosterToCloud(pdfFilePath);
174-
return posterUploaded;
205+
const posterUploaded = await uploadPosterToCloud(pdfFilePath);
206+
return posterUploaded;
207+
} catch (err) {
208+
await safeClosePage();
209+
throw err;
210+
}
175211
}
176212

177213
async function generate(options) {
@@ -185,15 +221,31 @@ async function generate(options) {
185221
onInfo('Creating new browser instance');
186222
await initialize();
187223
}
224+
225+
// Recycle the browser periodically to reclaim Chromium-internal heap.
226+
if (renderCount > 0 && renderCount % BROWSER_RECYCLE_AFTER === 0) {
227+
await closeBrowser();
228+
await initialize();
229+
}
230+
231+
let timeoutHandle;
188232
const timeout = new Promise((resolve, reject) => {
189-
setTimeout(reject, RENDER_TIMEOUT, new Error('Render timeout'));
233+
timeoutHandle = setTimeout(reject, RENDER_TIMEOUT, new Error('Render timeout'));
190234
});
191235

192-
const posterUploaded = await Promise.race([renderComponent(options), timeout]);
236+
let posterUploaded;
237+
try {
238+
posterUploaded = await Promise.race([renderComponent(options), timeout]);
239+
} finally {
240+
// Always clear the timeout so the closure is released immediately.
241+
clearTimeout(timeoutHandle);
242+
}
243+
193244
const uploadFailed = !posterUploaded && AZURE_STORAGE_ACCOUNT && AZURE_STORAGE_KEY;
194245

195246
if (!uploadFailed) {
196247
onInfo('Rendered successfully.');
248+
renderCount += 1;
197249
} else {
198250
const err = { message: 'Rendered successfully but uploading poster failed.', stack: '' };
199251
throw err;
@@ -202,6 +254,17 @@ async function generate(options) {
202254
return { success: true, uploaded: !uploadFailed };
203255
} catch (error) {
204256
onError(error);
257+
if (browser) {
258+
try {
259+
const pages = await browser.pages();
260+
console.log(`[browser] Pages open after error: ${pages.length}`);
261+
} catch (inspectErr) {
262+
console.error(
263+
`[browser] Cannot inspect pages after error, closing browser: ${inspectErr.message}`,
264+
);
265+
await closeBrowser();
266+
}
267+
}
205268
}
206269
}
207270

@@ -212,4 +275,5 @@ module.exports = {
212275
generate,
213276
generateRenderUrl,
214277
csvPath,
278+
closeBrowser,
215279
};

‎scripts/worker.js‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
const { Worker, QueueScheduler } = require('bullmq');
22
const Redis = require('ioredis');
33

4-
const { generate } = require('./generator');
4+
const { generate, closeBrowser } = require('./generator');
55
const { addEvent, updatePoster } = require('./store');
66
const { REDIS_CONNECTION_STRING } = require('../constants');
77

@@ -75,9 +75,18 @@ worker.on('failed', (job, err) => {
7575

7676
worker.on('drained', () => console.log('Job queue empty! Waiting for new jobs...'));
7777

78-
process.on('SIGINT', () => {
79-
console.log('Shutting down worker...');
80-
worker.close(true);
81-
queueScheduler.close();
78+
async function shutdown(signal) {
79+
console.log(`Shutting down worker (${signal})...`);
80+
try {
81+
await worker.close(true);
82+
await queueScheduler.close();
83+
await closeBrowser();
84+
} catch (err) {
85+
console.error(`Error during shutdown: ${err.message}`);
86+
}
8287
process.exit(0);
83-
});
88+
}
89+
90+
process.on('SIGINT', () => shutdown('SIGINT'));
91+
// Docker / Kubernetes sends SIGTERM; handle it so the process drains cleanly.
92+
process.on('SIGTERM', () => shutdown('SIGTERM'));

0 commit comments

Comments
 (0)