fix(🐛): memory leak with CPU surfaces on Node/Web (#2327)

docs/docs/getting-started/headless.md

@@ -15,13 +15,16 @@ You will notice in the example below that the import URL looks different than th
 
 ```tsx
 import { LoadSkiaWeb } from "@shopify/react-native-skia/lib/commonjs/web/LoadSkiaWeb";
-import { Fill, makeOffscreenSurface, drawOffscreen } from "@shopify/react-native-skia/lib/commonjs/headless";
+import { Fill, makeOffscreenSurface, drawOffscreen, getSkiaExports } from "@shopify/react-native-skia/lib/commonjs/headless";
 
 (async () => {
   const width = 256;
   const height = 256;
   const r = size * 0.33;
   await LoadSkiaWeb();
+  // Once that CanvasKit is loaded, you can access Skia via getSkiaExports()
+  // Alternatively you can do const {Skia} = require("@shopify/react-native-skia")
+  const {Skia} = getSkiaExports();
   const surface = makeOffscreenSurface(width, height);
   const image = drawOffscreen(surface,
     <Group blendMode="multiply">

package/src/__tests__/setup.ts

@@ -25,7 +25,9 @@ export const processResult = (
   surface.flush();
   const image = surface.makeImageSnapshot();
   surface.getCanvas().clear(Float32Array.of(0, 0, 0, 0));
-  return checkImage(image, relPath, { overwrite });
+  const result = checkImage(image, relPath, { overwrite });
+  image.dispose();
+  return result;
 };
 
 interface CheckImageOptions {

package/src/headless/index.ts

@@ -7,11 +7,11 @@ import { JsiSkApi } from "../skia/web";
 import { SkiaRoot } from "../renderer/Reconciler";
 import { JsiDrawingContext } from "../dom/types";
 import type { SkSurface } from "../skia";
+import { Skia } from "../skia/types";
 
 export * from "../renderer/components";
 
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-let Skia: any;
+let Skia: Skia;
 
 export const makeOffscreenSurface = (width: number, height: number) => {
   if (!Skia) {
@@ -24,12 +24,20 @@ export const makeOffscreenSurface = (width: number, height: number) => {
   return surface;
 };
 
+export const getSkiaExports = () => {
+  if (!Skia) {
+    Skia = JsiSkApi(CanvasKit);
+  }
+  return { Skia };
+};
+
 export const drawOffscreen = (surface: SkSurface, element: ReactNode) => {
   const root = new SkiaRoot(Skia, false);
   root.render(element);
   const canvas = surface.getCanvas();
   const ctx = new JsiDrawingContext(Skia, canvas);
   root.dom.render(ctx);
+  root.unmount();
   surface.flush();
   return surface.makeImageSnapshot();
 };

package/src/renderer/__tests__/SkiaDOM.spec.tsx

@@ -15,8 +15,8 @@ describe("Test introductionary examples from our documentation", () => {
         <Circle cx={width / 2} cy={height - r} r={r} color="yellow" />
       </Group>
     );
-    expect(root.children().length).toBe(1);
-    const child = root.children()[0]!;
+    expect(root.dom.children().length).toBe(1);
+    const child = root.dom.children()[0]!;
     expect(child).toBeDefined();
     expect(child.type).toBe(NodeType.Group);
     expect(child.children().length).toBe(3);
@@ -31,8 +31,8 @@ describe("Test introductionary examples from our documentation", () => {
         <Paint color="#ade6d8" style="stroke" strokeWidth={strokeWidth / 2} />
       </Circle>
     );
-    expect(root.children().length).toBe(1);
-    const child = root.children()[0]!;
+    expect(root.dom.children().length).toBe(1);
+    const child = root.dom.children()[0]!;
     expect(child).toBeDefined();
     expect(child.type).toBe(NodeType.Circle);
     const circle = child;

package/src/renderer/__tests__/Surfaces.spec.tsx

@@ -0,0 +1,43 @@
+import React from "react";
+
+import { Circle, Group } from "../components";
+import { processResult } from "../../__tests__/setup";
+import { setupSkia } from "../../skia/__tests__/setup";
+
+import { drawOnNode } from "./setup";
+
+describe("Surface", () => {
+  it("A raster surface shouldn't leak (1)", () => {
+    const { Skia } = setupSkia();
+    // When leaking, the WASM memory limit will be reached quite quickly
+    // causing the test to fail
+    for (let i = 0; i < 500; i++) {
+      const surface = Skia.Surface.Make(1920, 1080)!;
+      const canvas = surface.getCanvas();
+      canvas.clear(Skia.Color("cyan"));
+      surface.flush();
+      const image = surface.makeImageSnapshot();
+      image.dispose();
+      surface.dispose();
+    }
+  });
+  it("A raster surface shouldn't leak (2)", () => {
+    for (let i = 0; i < 1000; i++) {
+      //const t = performance.now();
+      const r = 128;
+      const surface = drawOnNode(
+        <>
+          <Group blendMode="multiply">
+            <Circle cx={r} cy={r} r={r} color="cyan" />
+            <Circle cx={r} cy={r} r={r} color="magenta" />
+            <Circle cx={r} cy={r} r={r} color="yellow" />
+          </Group>
+        </>
+      );
+      surface.flush();
+      //console.log(`Iteration ${i} took ${Math.floor(performance.now() - t)}ms`);
+      processResult(surface, "snapshots/leak.png");
+      surface.dispose();
+    }
+  });
+});

package/src/renderer/__tests__/setup.tsx

@@ -185,8 +185,9 @@ export const height = 256 * PIXEL_RATIO;
 export const center = { x: width / 2, y: height / 2 };
 
 export const drawOnNode = (element: ReactNode) => {
-  const { surface: ckSurface, draw } = mountCanvas(element);
+  const { surface: ckSurface, draw, root } = mountCanvas(element);
   draw();
+  root.unmount();
   return ckSurface;
 };
 
@@ -201,7 +202,7 @@ export const mountCanvas = (element: ReactNode) => {
   root.render(element);
   return {
     surface: ckSurface,
-    root: root.dom,
+    root,
     draw: () => {
       const ctx = new JsiDrawingContext(Skia, canvas);
       root.dom.render(ctx);
@@ -211,7 +212,7 @@ export const mountCanvas = (element: ReactNode) => {
 
 export const serialize = (element: ReactNode) => {
   const { root } = mountCanvas(element);
-  const serialized = serializeNode(root);
+  const serialized = serializeNode(root.dom);
   return JSON.stringify(serialized);
 };
 

package/src/skia/web/JsiSkImage.ts

@@ -45,7 +45,11 @@ export const toBase64String = (bytes: Uint8Array) => {
 };
 
 export class JsiSkImage extends HostObject<Image, "Image"> implements SkImage {
-  constructor(CanvasKit: CanvasKit, ref: Image) {
+  constructor(
+    CanvasKit: CanvasKit,
+    ref: Image,
+    private releaseCtx?: () => void
+  ) {
     super(CanvasKit, ref, "Image");
   }
 
@@ -128,6 +132,8 @@ export class JsiSkImage extends HostObject<Image, "Image"> implements SkImage {
     return toBase64String(bytes);
   }
 
+  // TODO: this is leaking on Web
+  // Add signature with allocated buffer
   readPixels(srcX?: number, srcY?: number, imageInfo?: ImageInfo) {
     const info = this.getImageInfo();
     const pxInfo: CKImageInfo = {
@@ -148,6 +154,9 @@ export class JsiSkImage extends HostObject<Image, "Image"> implements SkImage {
 
   dispose = () => {
     this.ref.delete();
+    if (this.releaseCtx) {
+      this.releaseCtx();
+    }
   };
 
   makeNonTextureImage(): SkImage {
@@ -157,14 +166,25 @@ export class JsiSkImage extends HostObject<Image, "Image"> implements SkImage {
       ...partialInfo,
       colorSpace,
     };
-    const pixels = this.ref.readPixels(0, 0, info) as Uint8Array | null;
+
+    var pixelLen = info.width * info.height * 4;
+    const pixelPtr = this.CanvasKit.Malloc(Uint8Array, pixelLen);
+    const pixels = this.ref.readPixels(
+      0,
+      0,
+      info,
+      pixelPtr,
+      info.width * 4
+    ) as Uint8Array | null;
     if (!pixels) {
       throw new Error("Could not create image from bytes");
     }
     const img = this.CanvasKit.MakeImage(info, pixels, info.width * 4);
     if (!img) {
       throw new Error("Could not create image from bytes");
     }
-    return new JsiSkImage(this.CanvasKit, img);
+    return new JsiSkImage(this.CanvasKit, img, () => {
+      this.CanvasKit.Free(pixelPtr);
+    });
   }
 }

package/src/skia/web/JsiSkSurface.ts

@@ -11,12 +11,19 @@ export class JsiSkSurface
   extends HostObject<Surface, "Surface">
   implements SkSurface
 {
-  constructor(CanvasKit: CanvasKit, ref: Surface) {
+  constructor(
+    CanvasKit: CanvasKit,
+    ref: Surface,
+    private releaseCtx?: () => void
+  ) {
     super(CanvasKit, ref, "Surface");
   }
 
   dispose = () => {
     this.ref.delete();
+    if (this.releaseCtx) {
+      this.releaseCtx();
+    }
   };
 
   flush() {

package/src/skia/web/JsiSkSurfaceFactory.ts

@@ -11,11 +11,26 @@ export class JsiSkSurfaceFactory extends Host implements SurfaceFactory {
   }
 
   Make(width: number, height: number) {
-    const surface = this.CanvasKit.MakeSurface(width, height);
+    var pixelLen = width * height * 4;
+    const pixelPtr = this.CanvasKit.Malloc(Uint8Array, pixelLen);
+    const surface = this.CanvasKit.MakeRasterDirectSurface(
+      {
+        width: width,
+        height: height,
+        colorType: this.CanvasKit.ColorType.RGBA_8888,
+        alphaType: this.CanvasKit.AlphaType.Unpremul,
+        colorSpace: this.CanvasKit.ColorSpace.SRGB,
+      },
+      pixelPtr,
+      width * 4
+    );
     if (!surface) {
       return null;
     }
-    return new JsiSkSurface(this.CanvasKit, surface);
+    surface.getCanvas().clear(this.CanvasKit.TRANSPARENT);
+    return new JsiSkSurface(this.CanvasKit, surface, () => {
+      this.CanvasKit.Free(pixelPtr);
+    });
   }
 
   MakeOffscreen(width: number, height: number) {