Fix memory leak in image optimization (vercel#23565)

This RP fixes the problem that the image optimization API uses a large amount of memory, and is not correctly freed afterwards. There're multiple causes of this problem:

### 1. Too many WebAssembly instances are created

We used to do all the image processing operations (decode, resize, rotate, encodeJpeg, encodePng, encodeWebp) inside each worker thread, where each operation creates at least one WASM instance, and we create `os.cpus().length - 1` workers by default. That means in the worst case, there will be `N*6` WASM instances created (N is the number of CPU cores minus one).

This PR changes it to a pipeline-like architecture: there will be at most 6 workers, and the same type of operations will always be assigned to the same worker. With this change, 6 WASM instances will be created in the worst case.

### 2. WebAssembly memory can't be deallocated

It's known that [WebAssembly can't simply deallocate its memory as of today](https://stackoverflow.com/a/51544868/2424786). And due to the implementation/design of the WASM modules that we are using, they're not very suitable for long-running cases and it's more like a one-off use. For each operation like resize, it will allocate **new memory** to store that data. So the memory will increase quickly as more images are processed.

The fix is to get rid of `execOnce` for WASM module initializations, so each time a new WASM module will be created and the old module will be GC'd entirely as there's no reference to it. That's the only and easiest way to free the memory use of a WASM module AFAIK.

### 3. WebAssembly memory isn't correctly freed after finishing the operation

`wasm-bindgen` generates code with global variables like `cachegetUint8Memory0` and `wasm` that always hold the WASM memory as a reference. We need to manually clean them up after finishing each operation. 

This PR ensures that these variables will be deleted so the memory overhead can go back to 0 when an operation is finished.

### 4. Memory leak inside event listeners

`emscripten` generates code with global error listener registration (without cleaning them up): https://github.com/vercel/next.js/blob/99a4ea6/packages/next/next-server/server/lib/squoosh/webp/webp_node_dec.js#L39-L43

And the listener has references to the WASM instance directly or indirectly: https://github.com/vercel/next.js/blob/99a4ea6/packages/next/next-server/server/lib/squoosh/webp/webp_node_dec.js#L183-L192 (`e`, `y`, `r`).

That means whenever a WASM module is created (emscripten), its memory will be kept by the global scope. And when we replace the WASM module with a new one, the newer will be added again and the old will still be referenced, which causes a leak.

Since we're running them inside worker threads (which will retry on fail), this PR simply removes these listeners.

### Test

Here're some statistics showing that these changes have improved the memory usage a lot (the app I'm using to test has one page of 20 high-res PNGs):

Before this PR (`next@10.1.0`):

<img src="https://user-images.githubusercontent.com/3676859/113058480-c3496100-91e0-11eb-9e5a-b325e484adac.png" width="500">

Memory went from ~250MB to 3.2GB (peak: 3.5GB) and never decreased again.

With fix 1 applied:

<img src="https://user-images.githubusercontent.com/3676859/113059060-921d6080-91e1-11eb-8ac6-83c70c1f2f75.png" width="500">

Memory went from ~280MB to 1.5GB (peak: 2GB).

With fix 1+2 applied:

<img src="https://user-images.githubusercontent.com/3676859/113059207-bf6a0e80-91e1-11eb-845a-870944f9e116.png" width="500">

Memory went from ~280MB to 1.1GB (peak: 1.6GB).

With fix 1+2+3+4 applied:

<img src="https://user-images.githubusercontent.com/3676859/113059362-ec1e2600-91e1-11eb-8d9a-8fbce8808802.png" width="500">

It's back to normal; memory changed from ~300MB to ~480MB, peaked at 1.2GB. You can clearly see that GC is working correctly here.

---

## Bug

- [x] Related issues vercel#23189, vercel#23436
- [ ] Integration tests added

## Feature

- [ ] Implements an existing feature request or RFC. Make sure the feature request has been accepted for implementation before opening a PR.
- [ ] Related issues linked using `fixes #number`
- [ ] Integration tests added
- [ ] Documentation added
- [ ] Telemetry added. In case of a feature if it's used or not.

## Documentation / Examples

- [ ] Make sure the linting passes

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]'
     }