Fix memory leak in image optimization

packages/next/next-server/server/lib/squoosh/codecs.ts

@@ -1,7 +1,6 @@
 import { promises as fsp } from 'fs'
 import * as path from 'path'
 import { instantiateEmscriptenWasm, pathify } from './emscripten-utils.js'
-import { execOnce } from '../../../lib/utils.js'
 
 // MozJPEG
 // @ts-ignore
@@ -23,25 +22,20 @@ const webpDecWasm = path.resolve(__dirname, './webp/webp_node_dec.wasm')
 // @ts-ignore
 import * as pngEncDec from './png/squoosh_png.js'
 const pngEncDecWasm = path.resolve(__dirname, './png/squoosh_png_bg.wasm')
-const pngEncDecInit = execOnce(() =>
+const pngEncDecInit = () =>
   pngEncDec.default(fsp.readFile(pathify(pngEncDecWasm)))
-)
 
 // OxiPNG
 // @ts-ignore
 import * as oxipng from './png/squoosh_oxipng.js'
 const oxipngWasm = path.resolve(__dirname, './png/squoosh_oxipng_bg.wasm')
-const oxipngInit = execOnce(() =>
-  oxipng.default(fsp.readFile(pathify(oxipngWasm)))
-)
+const oxipngInit = () => oxipng.default(fsp.readFile(pathify(oxipngWasm)))
 
 // Resize
 // @ts-ignore
 import * as resize from './resize/squoosh_resize.js'
 const resizeWasm = path.resolve(__dirname, './resize/squoosh_resize_bg.wasm')
-const resizeInit = execOnce(() =>
-  resize.default(fsp.readFile(pathify(resizeWasm)))
-)
+const resizeInit = () => resize.default(fsp.readFile(pathify(resizeWasm)))
 
 // rotate
 const rotateWasm = path.resolve(__dirname, './rotate/rotate.wasm')
@@ -128,7 +122,7 @@ export const preprocessors = {
           target_width: width,
           target_height: height,
         }))
-        return new ImageData(
+        const imageData = new ImageData(
           resize.resize(
             buffer,
             input_width,
@@ -142,6 +136,8 @@ export const preprocessors = {
           width,
           height
         )
+        resize.cleanup()
+        return imageData
       }
     },
     defaultOptions: {
@@ -270,7 +266,13 @@ export const codecs = {
     detectors: [/^\x89PNG\x0D\x0A\x1A\x0A/],
     dec: async () => {
       await pngEncDecInit()
-      return { decode: pngEncDec.decode }
+      return {
+        decode: (buffer: Buffer | Uint8Array): Buffer => {
+          const imageData = pngEncDec.decode(buffer)
+          pngEncDec.cleanup()
+          return imageData
+        },
+      } as any
     },
     enc: async () => {
       await pngEncDecInit()
@@ -287,7 +289,9 @@ export const codecs = {
             width,
             height
           )
-          return oxipng.optimise(simplePng, opts.level)
+          const imageData = oxipng.optimise(simplePng, opts.level)
+          oxipng.cleanup()
+          return imageData
         },
       }
     },

packages/next/next-server/server/lib/squoosh/impl.ts

@@ -1,47 +1,9 @@
 import { codecs as supportedFormats, preprocessors } from './codecs'
 import ImageData from './image_data'
 
-type RotateOperation = {
-  type: 'rotate'
-  numRotations: number
-}
-type ResizeOperation = {
-  type: 'resize'
-  width: number
-}
-export type Operation = RotateOperation | ResizeOperation
-export type Encoding = 'jpeg' | 'png' | 'webp'
-
-export async function processBuffer(
-  buffer: Buffer | Uint8Array,
-  operations: Operation[],
-  encoding: Encoding,
-  quality: number
-): Promise<Buffer | Uint8Array> {
-  let imageData = await decodeBuffer(buffer)
-  for (const operation of operations) {
-    if (operation.type === 'rotate') {
-      imageData = await rotate(imageData, operation.numRotations)
-    } else if (operation.type === 'resize') {
-      if (imageData.width && imageData.width > operation.width) {
-        imageData = await resize(imageData, operation.width)
-      }
-    }
-  }
-
-  switch (encoding) {
-    case 'jpeg':
-      return encodeJpeg(imageData, { quality })
-    case 'webp':
-      return encodeWebp(imageData, { quality })
-    case 'png':
-      return encodePng(imageData)
-    default:
-      throw Error(`Unsupported encoding format`)
-  }
-}
-
-async function decodeBuffer(_buffer: Buffer | Uint8Array): Promise<ImageData> {
+export async function decodeBuffer(
+  _buffer: Buffer | Uint8Array
+): Promise<ImageData> {
   const buffer = Buffer.from(_buffer)
   const firstChunk = buffer.slice(0, 16)
   const firstChunkString = Array.from(firstChunk)
@@ -54,11 +16,10 @@ async function decodeBuffer(_buffer: Buffer | Uint8Array): Promise<ImageData> {
     throw Error(`Buffer has an unsupported format`)
   }
   const d = await supportedFormats[key].dec()
-  const rgba = d.decode(new Uint8Array(buffer))
-  return rgba
+  return d.decode(new Uint8Array(buffer))
 }
 
-async function rotate(
+export async function rotate(
   image: ImageData,
   numRotations: number
 ): Promise<ImageData> {
@@ -68,7 +29,7 @@ async function rotate(
   return await m(image.data, image.width, image.height, { numRotations })
 }
 
-async function resize(image: ImageData, width: number) {
+export async function resize(image: ImageData, width: number) {
   image = ImageData.from(image)
 
   const p = preprocessors['resize']
@@ -79,7 +40,7 @@ async function resize(image: ImageData, width: number) {
   })
 }
 
-async function encodeJpeg(
+export async function encodeJpeg(
   image: ImageData,
   { quality }: { quality: number }
 ): Promise<Buffer | Uint8Array> {
@@ -94,7 +55,7 @@ async function encodeJpeg(
   return Buffer.from(r)
 }
 
-async function encodeWebp(
+export async function encodeWebp(
   image: ImageData,
   { quality }: { quality: number }
 ): Promise<Buffer | Uint8Array> {
@@ -109,7 +70,9 @@ async function encodeWebp(
   return Buffer.from(r)
 }
 
-async function encodePng(image: ImageData): Promise<Buffer | Uint8Array> {
+export async function encodePng(
+  image: ImageData
+): Promise<Buffer | Uint8Array> {
   image = ImageData.from(image)
 
   const e = supportedFormats['oxipng']

packages/next/next-server/server/lib/squoosh/main.ts

@@ -1,25 +1,57 @@
 import { Worker } from 'jest-worker'
 import * as path from 'path'
 import { execOnce } from '../../../lib/utils'
-import { Operation, Encoding } from './impl'
+import { cpus } from 'os'
+
+type RotateOperation = {
+  type: 'rotate'
+  numRotations: number
+}
+type ResizeOperation = {
+  type: 'resize'
+  width: number
+}
+export type Operation = RotateOperation | ResizeOperation
+export type Encoding = 'jpeg' | 'png' | 'webp'
 
 const getWorker = execOnce(
   () =>
     new Worker(path.resolve(__dirname, 'impl'), {
       enableWorkerThreads: true,
+      // There will be at most 6 workers needed since each worker will take
+      // at least 1 operation type.
+      numWorkers: Math.max(1, Math.min(cpus().length - 1, 6)),
+      computeWorkerKey: (method) => method,
     })
 )
 
-export { Operation }
-
 export async function processBuffer(
   buffer: Buffer,
   operations: Operation[],
   encoding: Encoding,
   quality: number
 ): Promise<Buffer> {
   const worker: typeof import('./impl') = getWorker() as any
-  return Buffer.from(
-    await worker.processBuffer(buffer, operations, encoding, quality)
-  )
+
+  let imageData = await worker.decodeBuffer(buffer)
+  for (const operation of operations) {
+    if (operation.type === 'rotate') {
+      imageData = await worker.rotate(imageData, operation.numRotations)
+    } else if (operation.type === 'resize') {
+      if (imageData.width && imageData.width > operation.width) {
+        imageData = await worker.resize(imageData, operation.width)
+      }
+    }
+  }
+
+  switch (encoding) {
+    case 'jpeg':
+      return Buffer.from(await worker.encodeJpeg(imageData, { quality }))
+    case 'webp':
+      return Buffer.from(await worker.encodeWebp(imageData, { quality }))
+    case 'png':
+      return Buffer.from(await worker.encodePng(imageData))
+    default:
+      throw Error(`Unsupported encoding format`)
+  }
 }

packages/next/next-server/server/lib/squoosh/mozjpeg/mozjpeg_node_dec.js

@@ -39,12 +39,6 @@ var Module = (function () {
       a.buffer || v('Assertion failed: undefined')
       return a
     }
-    1 < process.argv.length && (ba = process.argv[1].replace(/\\/g, '/'))
-    process.argv.slice(2)
-    process.on('uncaughtException', function (a) {
-      if (!(a instanceof ja)) throw a
-    })
-    process.on('unhandledRejection', v)
     ca = function (a) {
       process.exit(a)
     }

packages/next/next-server/server/lib/squoosh/mozjpeg/mozjpeg_node_enc.js

@@ -39,12 +39,6 @@ var Module = (function () {
       a.buffer || u('Assertion failed: undefined')
       return a
     }
-    1 < process.argv.length && (da = process.argv[1].replace(/\\/g, '/'))
-    process.argv.slice(2)
-    process.on('uncaughtException', function (a) {
-      if (!(a instanceof la)) throw a
-    })
-    process.on('unhandledRejection', u)
     ea = function (a) {
       process.exit(a)
     }

packages/next/next-server/server/lib/squoosh/png/squoosh_oxipng.js

@@ -127,3 +127,10 @@ async function init(input) {
 }
 
 export default init
+
+// Manually remove the wasm and memory references to trigger GC
+export function cleanup() {
+  wasm = null
+  cachegetUint8Memory0 = null
+  cachegetInt32Memory0 = null
+}

packages/next/next-server/server/lib/squoosh/png/squoosh_png.js

@@ -179,3 +179,11 @@ async function init(input) {
 }
 
 export default init
+
+// Manually remove the wasm and memory references to trigger GC
+export function cleanup() {
+  wasm = null
+  cachegetUint8ClampedMemory0 = null
+  cachegetUint8Memory0 = null
+  cachegetInt32Memory0 = null
+}

packages/next/next-server/server/lib/squoosh/resize/squoosh_resize.js

@@ -130,3 +130,10 @@ async function init(input) {
 }
 
 export default init
+
+// Manually remove the wasm and memory references to trigger GC
+export function cleanup() {
+  wasm = null
+  cachegetUint8Memory0 = null
+  cachegetInt32Memory0 = null
+}

packages/next/next-server/server/lib/squoosh/webp/webp_node_dec.js

@@ -35,12 +35,6 @@ var Module = (function () {
       a.buffer || x('Assertion failed: undefined')
       return a
     }
-    1 < process.argv.length && process.argv[1].replace(/\\/g, '/')
-    process.argv.slice(2)
-    process.on('uncaughtException', function (a) {
-      throw a
-    })
-    process.on('unhandledRejection', x)
     e.inspect = function () {
       return '[Emscripten Module object]'
     }

packages/next/next-server/server/lib/squoosh/webp/webp_node_enc.js

@@ -35,12 +35,6 @@ var Module = (function () {
       a.buffer || u('Assertion failed: undefined')
       return a
     }
-    1 < process.argv.length && process.argv[1].replace(/\\/g, '/')
-    process.argv.slice(2)
-    process.on('uncaughtException', function (a) {
-      throw a
-    })
-    process.on('unhandledRejection', u)
     f.inspect = function () {
       return '[Emscripten Module object]'
     }