Skip to content

Commit 2dec537

Browse files
committed
fix: avoid WASM heap overflow when saving/loading large labelmaps
Use a fresh temporary web worker for each labelmap read/write instead of the shared ITK-wasm worker. Each worker gets a clean WASM heap, preventing the 2GB signed pointer overflow in Emscripten that occurred when the shared worker's heap accumulated from loading large base images. Also fixes Array.from OOM when saving large labelmaps by using subarray copy instead.
1 parent baf4b50 commit 2dec537

File tree

6 files changed

+446
-9
lines changed

6 files changed

+446
-9
lines changed

src/components/SaveSession.vue

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import { saveAs } from 'file-saver';
3636
import { onKeyDown } from '@vueuse/core';
3737
3838
import { serialize } from '../io/state-file/serialize';
39+
import { useMessageStore } from '../store/messages';
3940
4041
const DEFAULT_FILENAME = 'session.volview.zip';
4142
@@ -58,6 +59,11 @@ export default defineComponent({
5859
const blob = await serialize();
5960
saveAs(blob, fileName.value);
6061
props.close();
62+
} catch (err) {
63+
const messageStore = useMessageStore();
64+
messageStore.addError('Failed to save session', {
65+
error: err instanceof Error ? err : new Error(String(err)),
66+
});
6167
} finally {
6268
saving.value = false;
6369
}

src/io/readWriteImage.ts

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,23 +8,29 @@ import {
88
import { vtiReader, vtiWriter } from '@/src/io/vtk/async';
99
import { getWorker } from '@/src/io/itk/worker';
1010

11-
export const readImage = async (file: File) => {
11+
export const readImage = async (file: File, webWorker?: Worker | null) => {
1212
if (file.name.endsWith('.vti'))
1313
return (await vtiReader(file)) as vtkImageData;
1414

15-
const { image } = await readImageItk(file, { webWorker: getWorker() });
15+
const { image } = await readImageItk(file, {
16+
webWorker: webWorker ?? getWorker(),
17+
});
1618
return vtkITKHelper.convertItkToVtkImage(image);
1719
};
1820

19-
export const writeImage = async (format: string, image: vtkImageData) => {
21+
export const writeImage = async (
22+
format: string,
23+
image: vtkImageData,
24+
webWorker?: Worker | null
25+
) => {
2026
if (format === 'vti') {
2127
return vtiWriter(image);
2228
}
2329
// copyImage so writeImage does not detach live data when passing to worker
2430
const itkImage = copyImage(vtkITKHelper.convertVtkToItkImage(image));
2531

2632
const result = await writeImageItk(itkImage, `image.${format}`, {
27-
webWorker: getWorker(),
33+
webWorker: webWorker ?? getWorker(),
2834
});
2935
return result.serializedImage.data as Uint8Array<ArrayBuffer>;
3036
};

src/io/vtk/async.ts

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,39 @@ import vtkDataSet from '@kitware/vtk.js/Common/DataModel/DataSet';
55
import { vtkObject } from '@kitware/vtk.js/interfaces';
66
import { StateObject } from './common';
77

8+
// VTK.js DataArray.getState() calls Array.from() on typed arrays,
9+
// which OOMs for large images (>~180M voxels). This helper temporarily
10+
// swaps each array's data with empty before getState(), then injects
11+
// the original TypedArrays into the resulting state. Structured clone
12+
// (postMessage) handles TypedArrays efficiently, and vtk()
13+
// reconstruction accepts them in DataArray.extend().
14+
const getStateWithTypedArrays = (dataSet: vtkDataSet) => {
15+
const pointData = (dataSet as any).getPointData?.();
16+
const arrays: any[] = pointData?.getArrays?.() ?? [];
17+
18+
const typedArrays = arrays.map((arr: any) => arr.getData());
19+
20+
// Swap to empty so Array.from runs on [] instead of huge TypedArray
21+
arrays.forEach((arr: any) => arr.setData(new Uint8Array(0)));
22+
23+
let state: any;
24+
try {
25+
state = dataSet.getState();
26+
} finally {
27+
arrays.forEach((arr: any, i: number) => arr.setData(typedArrays[i]));
28+
}
29+
30+
// Inject original TypedArrays into the serialized state
31+
state?.pointData?.arrays?.forEach((entry: any, i: number) => {
32+
if (entry?.data) {
33+
entry.data.values = typedArrays[i];
34+
entry.data.size = typedArrays[i].length;
35+
}
36+
});
37+
38+
return state;
39+
};
40+
841
interface SuccessReadResult {
942
status: 'success';
1043
obj: StateObject;
@@ -52,7 +85,7 @@ export const runAsyncVTKWriter =
5285
);
5386
const worker = new PromiseWorker(asyncWorker);
5487
const result = (await worker.postMessage({
55-
obj: dataSet.getState(),
88+
obj: getStateWithTypedArrays(dataSet),
5689
writerName,
5790
})) as WriteResult;
5891
asyncWorker.terminate();

src/store/segmentGroups.ts

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { onImageDeleted } from '@/src/composables/onImageDeleted';
1010
import { normalizeForStore, removeFromArray } from '@/src/utils';
1111
import { SegmentMask } from '@/src/types/segment';
1212
import { DEFAULT_SEGMENT_MASKS, CATEGORICAL_COLORS } from '@/src/config';
13+
import { createWebWorker } from 'itk-wasm';
1314
import { readImage, writeImage } from '@/src/io/readWriteImage';
1415
import {
1516
type DataSelection,
@@ -480,12 +481,21 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
480481

481482
state.manifest.segmentGroups = serialized;
482483

483-
// save labelmap images
484+
// save labelmap images — fresh worker per write to avoid heap accumulation
484485
await Promise.all(
485486
serialized.map(async ({ id, path }) => {
486487
const vtkImage = dataIndex[id];
487-
const serializedImage = await writeImage(saveFormat.value, vtkImage);
488-
zip.file(path, serializedImage);
488+
const worker = await createWebWorker(null);
489+
try {
490+
const serializedImage = await writeImage(
491+
saveFormat.value,
492+
vtkImage,
493+
worker
494+
);
495+
zip.file(path, serializedImage);
496+
} finally {
497+
worker.terminate();
498+
}
489499
})
490500
);
491501
}
@@ -527,7 +537,14 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
527537
const file = stateFiles.find(
528538
(entry) => entry.archivePath === normalize(segmentGroup.path!)
529539
)?.file;
530-
return { image: await readImage(file!) };
540+
// Use a fresh worker per labelmap to avoid WASM heap accumulation.
541+
// The shared worker may already have a large heap from base images.
542+
const worker = await createWebWorker(null);
543+
try {
544+
return { image: await readImage(file!, worker) };
545+
} finally {
546+
worker.terminate();
547+
}
531548
}
532549

533550
const labelmapResults = await Promise.all(
Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
import * as fs from 'node:fs';
2+
import * as path from 'node:path';
3+
import * as zlib from 'node:zlib';
4+
import { cleanuptotal } from 'wdio-cleanuptotal-service';
5+
import { volViewPage } from '../pageobjects/volview.page';
6+
import { DOWNLOAD_TIMEOUT, TEMP_DIR } from '../../wdio.shared.conf';
7+
import { writeManifestToFile } from './utils';
8+
9+
// 268M voxels — labelmap at this size triggers Array.from OOM
10+
const DIM_X = 1024;
11+
const DIM_Y = 1024;
12+
const DIM_Z = 256;
13+
14+
const writeBufferToFile = async (data: Buffer, fileName: string) => {
15+
const filePath = path.join(TEMP_DIR, fileName);
16+
await fs.promises.writeFile(filePath, data);
17+
cleanuptotal.addCleanup(async () => {
18+
if (fs.existsSync(filePath)) fs.unlinkSync(filePath);
19+
});
20+
return filePath;
21+
};
22+
23+
// UInt8 base image — small compressed size, fast to load
24+
const createUint8NiftiGz = () => {
25+
const header = Buffer.alloc(352);
26+
header.writeInt32LE(348, 0);
27+
header.writeInt16LE(3, 40);
28+
header.writeInt16LE(DIM_X, 42);
29+
header.writeInt16LE(DIM_Y, 44);
30+
header.writeInt16LE(DIM_Z, 46);
31+
header.writeInt16LE(1, 48);
32+
header.writeInt16LE(1, 50);
33+
header.writeInt16LE(1, 52);
34+
header.writeInt16LE(2, 70); // datatype: UINT8
35+
header.writeInt16LE(8, 72); // bitpix
36+
header.writeFloatLE(1, 76);
37+
header.writeFloatLE(1, 80);
38+
header.writeFloatLE(1, 84);
39+
header.writeFloatLE(1, 88);
40+
header.writeFloatLE(352, 108);
41+
header.writeFloatLE(1, 112);
42+
header.writeInt16LE(1, 254);
43+
header.writeFloatLE(1, 280);
44+
header.writeFloatLE(0, 284);
45+
header.writeFloatLE(0, 288);
46+
header.writeFloatLE(0, 292);
47+
header.writeFloatLE(0, 296);
48+
header.writeFloatLE(1, 300);
49+
header.writeFloatLE(0, 304);
50+
header.writeFloatLE(0, 308);
51+
header.writeFloatLE(0, 312);
52+
header.writeFloatLE(0, 316);
53+
header.writeFloatLE(1, 320);
54+
header.writeFloatLE(0, 324);
55+
header.write('n+1\0', 344, 'binary');
56+
57+
const imageData = Buffer.alloc(DIM_X * DIM_Y * DIM_Z);
58+
return zlib.gzipSync(Buffer.concat([header, imageData]), { level: 1 });
59+
};
60+
61+
const waitForFileExists = (filePath: string, timeout: number) =>
62+
new Promise<void>((resolve, reject) => {
63+
const dir = path.dirname(filePath);
64+
const basename = path.basename(filePath);
65+
66+
const watcher = fs.watch(dir, (eventType, filename) => {
67+
if (eventType === 'rename' && filename === basename) {
68+
clearTimeout(timerId);
69+
watcher.close();
70+
resolve();
71+
}
72+
});
73+
74+
const timerId = setTimeout(() => {
75+
watcher.close();
76+
reject(
77+
new Error(`File ${filePath} not created within ${timeout}ms timeout`)
78+
);
79+
}, timeout);
80+
81+
fs.access(filePath, fs.constants.R_OK, (err) => {
82+
if (!err) {
83+
clearTimeout(timerId);
84+
watcher.close();
85+
resolve();
86+
}
87+
});
88+
});
89+
90+
describe('Save large labelmap', function () {
91+
this.timeout(180_000);
92+
93+
it('saves session without error when labelmap exceeds 200M voxels', async () => {
94+
const prefix = `save-large-${Date.now()}`;
95+
const baseFileName = `${prefix}-u8.nii.gz`;
96+
97+
await writeBufferToFile(createUint8NiftiGz(), baseFileName);
98+
99+
const manifest = { resources: [{ url: `/tmp/${baseFileName}` }] };
100+
const manifestFileName = `${prefix}-manifest.json`;
101+
await writeManifestToFile(manifest, manifestFileName);
102+
103+
await volViewPage.open(`?urls=[tmp/${manifestFileName}]`);
104+
await volViewPage.waitForViews(DOWNLOAD_TIMEOUT * 6);
105+
106+
// Activate paint tool — creates a segment group
107+
await volViewPage.activatePaint();
108+
109+
// Paint a stroke to allocate the labelmap
110+
const views2D = await volViewPage.getViews2D();
111+
const canvas = await views2D[0].$('canvas');
112+
const location = await canvas.getLocation();
113+
const size = await canvas.getSize();
114+
const cx = Math.round(location.x + size.width / 2);
115+
const cy = Math.round(location.y + size.height / 2);
116+
117+
await browser
118+
.action('pointer')
119+
.move({ x: cx, y: cy })
120+
.down()
121+
.move({ x: cx + 20, y: cy })
122+
.up()
123+
.perform();
124+
125+
const notificationsBefore = await volViewPage.getNotificationsCount();
126+
127+
// Save session — before fix, this throws RangeError: Invalid array length
128+
const sessionFileName = await volViewPage.saveSession();
129+
const downloadedPath = path.join(TEMP_DIR, sessionFileName);
130+
131+
// Wait for either the file to appear (success) or notification (error)
132+
const saveResult = await Promise.race([
133+
waitForFileExists(downloadedPath, 90_000).then(() => 'saved' as const),
134+
browser
135+
.waitUntil(
136+
async () => {
137+
const count = await volViewPage.getNotificationsCount();
138+
return count > notificationsBefore;
139+
},
140+
{ timeout: 90_000, interval: 1000 }
141+
)
142+
.then(() => 'error' as const),
143+
]);
144+
145+
if (saveResult === 'error') {
146+
const errorDetails = await browser.execute(() => {
147+
const app = document.querySelector('#app') as any;
148+
const pinia = app?.__vue_app__?.config?.globalProperties?.$pinia;
149+
if (!pinia) return 'no pinia';
150+
const store = pinia.state.value.message;
151+
if (!store) return 'no message store';
152+
return store.msgList
153+
.map((id: string) => {
154+
const msg = store.byID[id];
155+
return `[${msg.type}] ${msg.title}: ${msg.options?.details?.slice(0, 300)}`;
156+
})
157+
.join('\n');
158+
});
159+
throw new Error(`Save error:\n${errorDetails}`);
160+
}
161+
162+
// Wait for the file to be fully written (Chrome may create it before flushing)
163+
await browser.waitUntil(
164+
() => {
165+
try {
166+
return fs.statSync(downloadedPath).size > 0;
167+
} catch {
168+
return false;
169+
}
170+
},
171+
{
172+
timeout: 30_000,
173+
interval: 500,
174+
timeoutMsg: 'Downloaded file remained 0 bytes',
175+
}
176+
);
177+
const stat = fs.statSync(downloadedPath);
178+
expect(stat.size).toBeGreaterThan(0);
179+
});
180+
});

0 commit comments

Comments
 (0)