Cache worklets

[piaskowyk]

Summary

Currently, every time we want to execute a shareable worklet, we need to call toJSValue() to convert Reanimated's Shareable into a runnable JavaScript function. This operation can be quite expensive for larger methods that have many dependencies in their closure (such as objects and other worklets). Previously, the result of toJSValue() wasn't cached, which meant we had to convert the same shareable multiple times, especially on every call to runOnUI() and in response to events - potentially on every frame.

This happens because the part of the code - runGuarded - is called frequently. You can see this code here.

This PR introduces the retention of all worklets and caches the result of toJSValue() per runtime.

⚠️ This change is potentially risky, and it's challenging to predict if there are any edge cases where caching everything might not be appropriate. However, at this moment, we haven't found any regressions related to memory issues or crashes.

⚠️ This PR changes the default behavior of worklets. Previously, worklets were stateless and destroyed their closure after every invocation, but now the closure will persist as long as the worklet lives.

stateless vs stateful example

export default function Example() {
  let counter = {value: 0};
  const workletFunction = () => {
    'worklet';
    counter.value++;
    console.log(counter);
  };
  return <Button title="click" onPress={runOnUI(workletFunction)}/>;
}

Previous output

{"value": 1}
{"value": 1}
{"value": 1}

Current output

{"value": 1}
{"value": 2}
{"value": 3}

However, after the render, the worklet and their closure will be created again.

Issue reproduction example

code

import { Text, StyleSheet, View, Button } from 'react-native';

import React from 'react';
import Animated, { withSpring, useAnimatedScrollHandler, withClamp, withDecay, withDelay, withTiming, useAnimatedRef, scrollTo, runOnUI } from 'react-native-reanimated';

function mleko() {
  'worklet';
  console.log('mleko');
}

export default function EmptyExample() {
  const aref = useAnimatedRef<Animated.ScrollView>();
  const onScroll = useAnimatedScrollHandler({
    onBeginDrag: () => {
      'worklet'
      withSpring;
      withClamp;
      withDecay;
      withDelay;
      withTiming;
    },
  });
  return (
    <View style={styles.container}>
      <Button onPress={() => {
        runOnUI(() => {
          mleko();
          withTiming;
          console.log('runOnUI');
          scrollTo(aref, 0, 100, true);
        })();
      }} title='click' />
      <Animated.ScrollView onScroll={onScroll} ref={aref}>
        {Array.from({ length: 1000 }, (_, i) => (
          <Text key={i}>Item_____________ {i}</Text>
        ))}
      </Animated.ScrollView>
    </View>
  );
}

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

Memory test example

I have tested whether this change leads to memory leaks, and according to my tests, the behavior remains exactly the same as before. Here is my test example:

code

import { Text, StyleSheet, View, Button } from 'react-native';

import {useRef} from 'react';
import Animated, { withSpring, useAnimatedScrollHandler, withClamp, withDecay, withDelay, withTiming, useAnimatedRef, scrollTo, runOnUI } from 'react-native-reanimated';

export default function EmptyExample() {
  const aref = useAnimatedRef<Animated.ScrollView>();
  const onScroll = useAnimatedScrollHandler({
    onBeginDrag: () => {
      'worklet'
      withSpring;
      withClamp;
      withDecay;
      withDelay;
      withTiming;
    },
  });

  const ref = useRef(0);

  function test1() {
    const obj = {a: 5};
    for(let i = 0; i < 10000; i++) {
      runOnUI(() => {
        const a = 5 + obj.a;
        if (a < 5) {
          console.log('a', a);
        }
      })();
    }
  }

  return (
    <View style={styles.container}>
      <Button onPress={test1} title='test1' />
      <Button onPress={() => {
        global.gc();
        runOnUI(() => {
          global.gc();
        })();
      }} title='gc' />
      <Button onPress={() => {
        runOnUI(() => {
          withTiming;
          scrollTo(aref, 0, 100, true);
        })();
      }} title='scroll' />
      <Animated.ScrollView onScroll={onScroll} ref={aref}>
        {Array.from({ length: 1000 }, (_, i) => (
          <Text key={i}>Item_____________ {i}</Text>
        ))}
      </Animated.ScrollView>
    </View>
  );
}

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

To test for memory leaks, you can follow these steps:

1. Add a counter of shareable worklets

+int ShareableWorklet::objCounter = 0;
ShareableWorklet::ShareableWorklet(jsi::Runtime &rt, const jsi::Object &worklet)
    : ShareableObject(rt, worklet) {
  valueType_ = WorkletType;
+  objCounter++;
}
ShareableWorklet::~ShareableWorklet() {
+  objCounter--;
}

2. Add breakpoints to the constructor and destructor.
3. Press the test1 button a few times.
4. Press the gc (garbage collection) button a few times.
5. Check if the counter returns to the value it had at the beginning.
6. Note: The counter should never reach 0 because there are some internal worklets that exist in the global scope and should never be destructed during the lifetime of the React Context.

Additional tests cases

I've also checked for any regressions in our Example app, but everything seems to be functioning normally.

code

import { Text, StyleSheet, View, Button } from 'react-native';

import React, {useEffect, useRef} from 'react';
import Animated, { withSpring, useAnimatedScrollHandler, withClamp, withDecay, withDelay, withTiming, useAnimatedRef, scrollTo, runOnUI, runOnJS, useSharedValue } from 'react-native-reanimated';

function mleko() {
  'worklet';
  console.log('mleko');
}

export default function EmptyExample() {
  const aref = useAnimatedRef<Animated.ScrollView>();
  const onScroll = useAnimatedScrollHandler({
    onBeginDrag: () => {
      'worklet'
      withSpring;
      withClamp;
      withDecay;
      withDelay;
      withTiming;
    },
  });

  const ref = useRef(0);
  const sv = useSharedValue(0);

  function test1() {
    sv.value++;
    function tmp(arg: number) {
      sv.value++;
      console.log('tmp', ref.current, sv.value);
    }
    const obj = {a: 5};
    for(let i = 0; i < 10000; i++) {
      runOnUI(() => {
        const a = 5 + obj.a;
        if (a < 5) {
          console.log('a', a);
        }
      })();
    }
  }

  const remoteObject = {a: 5};
  function test2() {
    sv.value++;
    remoteObject.a++;
    function a({a, b}: {a: number, b: number} = {a: 5, b: 10}) {
      sv.value++;
      console.log('a', a, b, ref, remoteObject, sv.value);
    }
    function schedule(obj: any) {
      sv.value++;
      console.log('schedule', obj, remoteObject, sv.value);
      runOnUI((tmp) => {
        sv.value++;
        console.log('runOnUI', tmp, sv.value);
        runOnJS(a)(tmp as any);
      })({a: 3, b: 4});
    }
    runOnUI(() => {
      sv.value++;
      runOnJS(schedule)({a: 1, b: 2});
    })();
  }

  useEffect(() => {
    // setInterval(() => {
    //   ref.current++;
    //   console.log('ref', ref.current);
    // }, 1000);
  }, []);

  return (
    <View style={styles.container}>
      <Button onPress={test1} title='test1' />
      <Button onPress={test2} title='test2' />
      <Button onPress={() => {
        global.gc();
        runOnUI(() => {
          global.gc();
        })();
      }} title='gc' />
      <Button onPress={() => {
        runOnUI(() => {
          mleko();
          withTiming;
          console.log('runOnUI');
          scrollTo(aref, 0, 100, true);
        })();
      }} title='scroll' />
      <Animated.ScrollView onScroll={onScroll} ref={aref}>
        {Array.from({ length: 1000 }, (_, i) => (
          <Text key={i}>Item_____________ {i}</Text>
        ))}
      </Animated.ScrollView>
    </View>
  );
}

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

[tjzel]

Let's use name Static Worklets for those. Let's also handle it with 'static worklet'; directive in our Babel Plugin.

[gaearon]

Just curious — what determines cache-ability?

[tjzel]

@gaearon Immutable closure

[christianbaroni]

FWIW, I was playing around with these changes in the Rainbow app and was able to enable caching for nearly all worklets, with no visible issues. Could be I’m missing certain cases or problems it creates, but runOnJS was the only case I noticed break with caching.

As far as I can tell, it resolves a considerable number of frame drops on both iOS and Android, which without these changes happen consistently. There was a much more pronounced improvement with the full caching. In my testing (limited), 99% of worklets were being cached. We use Reanimated pretty heavily. Gestures in particular were noticeably smoother.

[piaskowyk]

Thanks for testing and the insights! 🙌 We're thinking about enabling caching for all worklets, but we want to make sure it won't cause any issues. So, we plan to look into it more thoroughly first.

[piaskowyk]

@christianbaroni You mentioned that you encountered an issue where caching breaks runOnJS. Could you tell me more about how I can reproduce this issue?

[christianbaroni]

Here's the patch I'm using. I haven't observed any issues with this implementation:

react-native-reanimated+3.16.3.patch

diff --git a/node_modules/react-native-reanimated/Common/cpp/worklets/SharedItems/Shareables.cpp b/node_modules/react-native-reanimated/Common/cpp/worklets/SharedItems/Shareables.cpp
index 4087065..f22bd06 100644
--- a/node_modules/react-native-reanimated/Common/cpp/worklets/SharedItems/Shareables.cpp
+++ b/node_modules/react-native-reanimated/Common/cpp/worklets/SharedItems/Shareables.cpp
@@ -47,7 +47,12 @@ jsi::Value makeShareableClone(
   if (value.isObject()) {
     auto object = value.asObject(rt);
     if (!object.getProperty(rt, "__workletHash").isUndefined()) {
-      shareable = std::make_shared<ShareableWorklet>(rt, object);
+      if (shouldRetainRemote.isBool() && shouldRetainRemote.getBool()) {
+        shareable =
+            std::make_shared<RetainingShareable<ShareableWorklet>>(rt, object);
+      } else {
+        shareable = std::make_shared<ShareableWorklet>(rt, object);
+      }
     } else if (!object.getProperty(rt, "__init").isUndefined()) {
       shareable = std::make_shared<ShareableHandle>(rt, object);
     } else if (object.isFunction(rt)) {
diff --git a/node_modules/react-native-reanimated/src/shareables.ts b/node_modules/react-native-reanimated/src/shareables.ts
index 44e961f..1b57980 100644
--- a/node_modules/react-native-reanimated/src/shareables.ts
+++ b/node_modules/react-native-reanimated/src/shareables.ts
@@ -128,6 +128,7 @@ export function makeShareableCloneRecursive<T>(
   const type = typeof value;
   const isTypeObject = type === 'object';
   const isTypeFunction = type === 'function';
+  let isCacheableWorklet = false;
   if ((isTypeObject || isTypeFunction) && value !== null) {
     const cached = shareableMappingCache.get(value);
     if (cached === shareableMappingFlag) {
@@ -166,6 +167,7 @@ export function makeShareableCloneRecursive<T>(
       } else if (isPlainJSObject(value) || isTypeFunction) {
         toAdapt = {};
         if (isWorkletFunction(value)) {
+          isCacheableWorklet = true;
           if (__DEV__) {
             const babelVersion = value.__initData.version;
             if (babelVersion !== undefined && babelVersion !== jsVersion) {
@@ -275,7 +277,7 @@ Offending code was: \`${getWorkletCode(value)}\``);
       }
       const adapted = NativeReanimatedModule.makeShareableClone(
         toAdapt,
-        shouldPersistRemote,
+        shouldPersistRemote || isCacheableWorklet,
         value
       );
       shareableMappingCache.set(value, adapted);

Quite a bit of the shareables.ts file has since been restructured, but at a glance this seems equivalent to your latest changes @piaskowyk.

My earlier comment and the 99% figure were a little misleading. I only noticed issues when I set isCacheableWorklet above to true in all cases (for all shareables, not just all worklets).

Here's what I noticed break when caching all shareables (RN 0.74.3, Paper):

• useSyncSharedValue.ts
  • This hook no longer synced properly (at least with the configuration used here): DappBrowser.tsx
  • Could have been wrong in assuming runOnJS was to blame — didn't dig into where it was failing, may have had more to do with useAnimatedReaction not seeing state changes
• dispatchCommand
  • Used here to scroll an Animated.ScrollView: useBrowserScrollView.ts
    • Probably due to stale args — believe it worked if I defined the scrollTo point outside of dispatchCommand, but within the onTouchesMove worklet

From what I remember, most non-worklet shareables were cacheable without apparent issue. But seemed far more susceptible to edge cases and bugs, and not sure it's even desirable or whether it's redundant.

Hope that helps — lmk if you have any other questions.

[piaskowyk]

Thanks for all provided informations! That makes sense now - caching all shareables (including shared values) isn't correct. I've tested it intensively recently, and it seems like we can cache all worklets. This should result in a noticeable performance improvement.

[tjzel]

@piaskowyk Just to be clear, previously they destroyed their closure not after the invocation, but after the given worklet instance was disposed. For instance, a worklet from useAnimatedStyle would destroy it's closure only after a re-render (restarting the mapper).

[piaskowyk]

Okay, to clarify further 😅: the closure lives as long as the UI copy of the worklet method exists on the UI runtime. Whenever a worklet needs to be copied from one runtime to another, the closure is recreated. useAnimatedStyle() is an exception because it keeps a reference to a worklet in the mappers registry 👍

[gaearon]

Would it be a good idea at this point to patch this into our app (tentatively in prod as well)? I don't have a good mental model about what exactly is being cached here and what makes it safe (and why it wasn't being cached in the past). I.e. what is the exact tradeoff. I still don't fully understand what's implied by a closure being "immutable" earlier (are we talking about bindings or objects they point to? deep or shallow? etc)

[gaearon]

Oh wait, you did add details to the PR description, let me read through those.

[tjzel]

@gaearon I think it would work. I guess you adhere to good practices of React in the Bluesky app so you wouldn't be impacted by slight change of (never defined) behavior.

[gaearon]

The "stateless vs stateful example" in the PR description is very illuminating. I agree this does seem to be borderline undefined behavior. I guess the previous behavior is a little less fragile but the optimization may very well be worth it. Especially considering the compiler linter should catch these cases.

[christianbaroni]

One thing I noticed experimenting with more aggressive caching is that the vast majority of the values being repeatedly processed were static JS constants defined at the module scope and used in worklets, e.g.:

const ITEM_HEIGHT = 40;

I was thinking perhaps the babel plugin could identify these, which could then allow caching and sharing them across worklets. Our worklets use significantly more static JS values than dynamic ones (I’d guess 100:1), and the current closure mechanism seems to make no distinction.

Is this a correct understanding of the current behavior?

When I enabled caching for worklets and all JS values, I saw, I’d say, ~double the performance gain vs. worklet caching alone — I’m thinking due to all static values being cached.

[tjzel]

@christianbaroni In my opinion expanding the Reanimated Babel Plugin is not the way to go. The environments in which the developers build React/React Native can be very different. That brings a lot of issues when we decide to 'assume' something in the Plugin.

Common const elements located at module level are cached in shareableCache (and even frozen!) already.

The problem is not about the React Runtime at all. All values from the React Runtime that are going to be serialized are serialized only once (which brings its own problems at times). This pull request addresses de-serialization, which used to happen on each worklet instantiation on the UI Runtime.

---

Good thing about the change is that it's reducing differences in behavior between Worklet Runtimes and the React Runtime. While Worklet Runtime are actual instances of runtimes which follow ECMAScript etc., the API we provide to operate on them makes them behave different than you'd expect. This is the course we've been trying to steer towards for a long time and I hope we can stay on it.

[tjzel]

you only cache once

[gaearon]

Is there an alpha release we can try?

[efstathiosntonas]

@gaearon you can try nightly build tomorrow.

[gaearon]

Tried updating to the nightly but the app instacrashes on an Android device on every start in Release mode. In Development it's fine.

Not sure if related to this PR.

  E  FATAL EXCEPTION: mqt_native_modules
Process: xyz.blueskyweb.app, PID: 12643
com.facebook.react.common.JavascriptException: Error: Exception in HostFunction: expected 0 arguments, got 1, js engine: hermes, stack:
callNativeSyncHook@1:161022
nonPromiseMethodWrapper@1:168634
NativeReanimated@1:705613
anonymous@1:699443
loadModuleImplementation@1:146748
guardedLoadModule@1:146277
metroRequire@1:145907
metroImportDefault@1:145971
makeShareableCloneRecursive@1:696452
makeShareableCloneRecursive@1:697238
makeShareableCloneRecursive@1:697127
makeShareableCloneRecursive@1:697238
makeMutableNative@1:691095
anonymous@1:690542
loadModuleImplementation@1:146748
guardedLoadModule@1:146277

[gaearon]

b2b58a6 is good
6fad03e is bad

Something in between

[piaskowyk]

Thanks for the info! I'll check to see if it's related and figure out how to fix it 🫡

[gaearon]

Confirmed this fixes the perf issue we were seeing