fix(Worklets): RemoteFunction memory cycle

[tjzel]

Summary

Fixes #9661

In #9272 which was supposed to fix memory issues with Remote Function, I accidentally introduced a memory-leak due to reference retain cycle.

Before

Basically, RemoteFunctionRegistry would hold a strong reference to the (remote) function, which in turn would hold a strong reference to the serialized remote function, which would hold the C++ value.

This C++ value would erase the original function from the registry, but only in its destructor - which would never happen, because the registry itself would keep it alive forever.

flowchart TB
  subgraph RN[RN Runtime]
    R[RemoteFunctionRegistry]
    F[remote function]
    C[serialized remote function]
  end
  subgraph CPP[C++ heap]
    S[C++ value]
  end
  subgraph WK[Worker Runtime]
    H[holderFunction]
  end

  R --> F
  F -.-> C
  C --> S
  H --> S
  S ==>|"registry.delete(id) — never invoked"| R

  linkStyle 0,2,3 stroke:#2f9e44,stroke-width:2px
  linkStyle 1 stroke:#868e96,stroke-width:2px
  linkStyle 4 stroke:#e03131,stroke-width:3px

Loading

After

To fix this issue, I split the C++ Remote Function for RN callbacks into two subclasses - one would only be owned by the RN Runtime, and the other one from Worklet Runtimes.

If all Worklet Runtime references would get cleaned up, then they would schedule cleanup on the RN Runtime C++ side, to erase the original function from the registry. Then, this would in turn release the cache value and the remaining C++ value.

flowchart TB
  subgraph RN[RN Runtime]
    R[RemoteFunctionRegistry]
    F[remote function]
    C[serialized remote function]
  end
  subgraph CPP[C++ heap]
    O[RN-owned C++ value]
    P[Worklet-owned C++ value]
  end
  subgraph WK[Worker Runtime]
    H[holderFunction]
  end

  R --> F
  F -.-> C
  C --> O
  H --> P
  P --> O
  O -.-> P
  P ==>|"registry.delete(id)"| R

  linkStyle 0,2,3,4 stroke:#2f9e44,stroke-width:2px
  linkStyle 1,5 stroke:#868e96,stroke-width:2px
  linkStyle 6 stroke:#2f9e44,stroke-width:3px

Loading

Test Plan

I used the following code and was tracking constructor and destructor invocations of SerializableRemoteFunction in debugger.

Code

import { useEffect } from 'react';
import { Button, StyleSheet, View } from 'react-native';
import {
  scheduleOnRN,
  scheduleOnUI,
  createWorkletRuntime,
  scheduleOnRuntime,
  createSerializable,
} from 'react-native-worklets';

const workerRuntime = createWorkletRuntime();

export default function App() {
  useEffect(() => {
    scheduleOnUI(() => {
      function foo() {
        const _ = new Array(10000).fill(0);
        requestAnimationFrame(foo);
      }
      foo();
    });
  });

  return (
    <View style={styles.container}>
      <Button
        title="Init runtime"
        onPress={() => {
          scheduleOnUI(() => {
            // scheduleOnRuntime(workerRuntime, () => {
            'worklet';
            console.log('init');
          });
        }}
      />
      <Button
        title="Press me first"
        onPress={() => {
          function logTest() {
            console.log('Button pressed!');
          }

          // scheduleOnRuntime(workerRuntime, () => {
          scheduleOnUI(() => {
            'worklet';
            scheduleOnRN(logTest);
          });
        }}
      />
      <Button
        title="Trigger GC on RN"
        onPress={() => {
          console.log('Triggering GC on RN');
          gc();
        }}
      />
      <Button
        title="Trigger GC on Worker"
        onPress={() => {
          // scheduleOnRuntime(workerRuntime, () => {
          scheduleOnUI(() => {
            'worklet';
            console.log('Triggering GC on Worker');
            gc();
          });
        }}
      />
    </View>
  );
}

const styles = StyleSheet.create({
  container: {
    flex: 1,
    backgroundColor: '#fff',
    alignItems: 'center',
    justifyContent: 'center',
  },
});