Feature request: Run a cleanup function when the molecule gets unmounted

[asfktz]

First, thank you for your work on this library! Decoupling Bunshi from Jotai was a great decision. I believe Bunshi has the potential to benefit not only Jotai but many other libraries as well.

One thing I've noticed that's missing for integration with other libraries is the ability to define a cleanup function that is called when the molecule gets unmounted. This isn’t an issue with Jotai since it already has a cleanup mechanism, but when working with xstate, for example, there's no way to stop the machine, which leads to a memory leak.

I made a small demo the demonstrate the problem: https://codesandbox.io/s/molecule-cleanup-mvt4cg?file=/App.tsx

1. Click "Mount Molecule" - the molecule will initialize the machine
2. Open the console – You'll see 2 logs every second*
3. Click "Unmount Molecole" – the machine keeps logging

* Becouse of React's Stricy mode "mount, unmount & remount" behavior. the machine is initialized twice. so we actually have 2 machine instances running at the same time.

I imagine something like this:

export const TimerMolecule = molecule((_, scope) => {
  scope(ComponentScope);

  const timerMachine = createMachine({
    id: "counter",
    context: { count: 0 },
    initial: "ticking",
    states: {
      ticking: {
        after: {
          1000: {
            target: "ticking",
            actions: [console.log, assign({ count: (ctx) => ctx.count + 1 })]
          }
        }
      }
    }
  });

  onCleanup(() => {
    timerMachine.stop();
  });

  return interpret(timerMachine).start();
});

[loganvolkers]

First, thank you for your work on this library! Decoupling Bunshi from Jotai was a great decision. I believe Bunshi has the potential to benefit not only Jotai but many other libraries as well.

Thanks for the feedback @asfktz! Next up looking at recipes for things like tanstack-query and floating-ui.

One thing I've noticed that's missing for integration with other libraries is the ability to define a cleanup function that is called when the molecule gets unmounted. This isn’t an issue with Jotai since it already has a cleanup mechanism, but when working with xstate, for example, there's no way to stop the machine, which leads to a memory leak.

A few parts to this question.

1. I agree -- this is a feature of tsyringe and some other DI frameworks that would be useful

2. This is challenging -- I foresee some trickiness in getting this to work properly across React, Vue and Vanilla JS

3. There may be workarounds -- as you described, this is not a problem for jotai or nanostores, so there may be a workaround for XState too. Potentially useActor is the wrong hook to use. useMolecule is basically a replacement for XState's internal useIdleInterpreter, but the rest of the code in useActor and useMachien for starting a machine only once isn't encapsulated or exported. Maybe @Andarist or someone from the XState team could chime in here, maybe we need to create a mirror issue in the XState repo, or maybe we should just open a PR.

4. React strict mode might be a bug -- that sounds like it could be something to fix imminently with use-sync-external-store

[asfktz]

This is challenging -- I foresee some trickiness in getting this to work properly across React, Vue and Vanilla JS

Do you mean it's challenging to know when to call the cleanup function? I imagine having some sort of listeners counter per molecule, like Jotai and [nanostores](https://github.com/nanostores/nanostores#:~:text=import%20%7B%20onMount%20%7D%20from%20%27nanostores%27%0A%0AonMount(%24profile%2C%20()%20%3D%3E%20%7B%0A%20%20//%20Mount%20mode%0A%20%20return%20()%20%3D%3E%20%7B%0A%20%20%20%20//%20Disabled%20mode%0A%20%20%7D%0A%7D)) have.

Maybe it will be simpler to trigger the cleanup when the scope is unmounted?

export const MoleculeA = molecule((_, scope) => {
  scope(MyScope);

  onCleanup(MyScope, () => {
     // cleanup logic
  });

  return {};
});

The disadvantage of this approach is that you have to declare a scope for molecules without an explicit scope if you want to let them have a cleanup function:

export const MyScope = createScope();

export const MoleculeA = molecule((_, scope) => {
  scope(MyScope);

  onCleanup(MyScope, () => {
     // cleanup logic
  });

  return {};
});

export const MoleculeB = molecule((getMol) => {
  getMol(MoleculeA);

  // Need to be aware of the scope on order /:
  onCleanup(MyScope, () => {
     // cleanup logic
  });

  return {};
});

There may be workarounds -- as you described, this is not a problem for jotai or nanostores, so there may be a workaround for XState too. Potentially useActor is the wrong hook to use. useMolecule is basically a replacement for XState's internal useIdleInterpreter, but the rest of the code in useActor and useMachien for starting a machine only once isn't encapsulated or exported. Maybe @Andarist or someone from the XState team could chime in here, maybe we need to create a mirror issue in the XState repo, or maybe we should just open a PR.

I tried to find such workaround, I ended up wrapping my actors with Jotai 😅

[loganvolkers]

Do you mean it's challenging to know when to call the cleanup function?

Exactly. The MoleculeInjector has both a get and use function, and we want to be clear and backwards compatible about how those work.

Maybe it will be simpler to trigger the cleanup when the scope is unmounted?

This idea might work! Since scopes already have a lifecycle, it could be tapped in to.

The disadvantage of this approach is that you have to declare a scope for molecules without explicit scopes if you want to let them have a cleanup function.

Molecules without scopes are never cleaned up. They will hang out in the internal WeakMap until the Molecule itself is destroyed. Changing this implicit contract is also not backwards-compatible.

[asfktz]

This idea might work! Since scopes already have a lifecycle, it could be tapped in to.

Cleaning up by scope will be a great improvement in itself! And if, by any chance, any Molecule could have its own cleanup without break up backwards compatibility — that would be just perfect.

[loganvolkers]

There's also the possibility of introducing another special scope, like ComponentScope, that is less global than the default scope.

Something like UsageScope or LazyScope or GloballyButWithCleanupScope that helps support cleanup for global molecules.

[asfktz]

How about something like DefaultScope?

export const MyScope = createScope();

export const MoleculeA = molecule((_, scope) => {
  scope(MyScope);

  onCleanup(MyScope, () => {
     // cleanup logic
  });

  return {};
});

export const MoleculeB = molecule((getMol) => {
  getMol(MoleculeA);

  onCleanup(DefaultScope, () => {
     // cleanup logic
  });

  return {};
});

As a user, I don't know what the default scope is, but I know I want one in order to cleanup my logic.

[Andarist]

If the lifecycle of a molecule should be tied to some component's lifecycle then the ability to run cleanup is crucial. I don't know much about this project to really understand what's going on or to advise how to fix this issue. I imagine that you can do a way better job at that than me ;p

When it comes to React's Strict Mode... yeeah, that's annoying. I was just working on improvement in this area in XState, you can see this work here. Note though that this PR targets the upcoming XState v5 which is better suited for the state rehydration etc so if u'd like to use the same strategy in v4... YMMV. I'm also abusing some internal-ish APIs there so I wouldn't recommend you to just copy this code over to this project here.

[loganvolkers]

If the lifecycle of a molecule should be tied to some component's lifecycle then the ability to run cleanup is crucial.

Agreed.

I don't know much about this project to really understand what's going on or to advise how to fix this issue. I imagine that you can do a way better job at that than me ;p

Thanks for responding! Very happy that you responded despite this being a small and rather unknown library.

For what we're trying to achieve, check out the XState recipe for React and Vue: https://www.bunshi.org/recipes/xstate/

When it comes to React's Strict Mode... yeeah, that's annoying. I was just working on improvement in this area in XState, you can see this work https://github.com/statelyai/xstate/pull/4436. Note though that this PR targets the upcoming XState v5 which is better suited for the state rehydration etc so if u'd like to use the same strategy in v4... YMMV. I'm also abusing some internal-ish APIs there so I wouldn't recommend you to just copy this code over to this project here.

I was just looking at that. The workaround we'd need to avoid cleanup in the molecules is having a hook that can activate an idle actor by calling start() when it's used and stop() when it is unmount, like below.

import isDevelopment from '#is-development';
import { useCallback, useEffect } from 'react';
import { useSyncExternalStore } from 'use-sync-external-store/shim';
import {
  ActorRefFrom,
  AnyActorLogic,
  ActorOptions,
  SnapshotFrom
} from 'xstate';
import { useIdleActor } from './useActorRef.ts';
import { stopRootWithRehydration } from './stopRootWithRehydration.ts';

export function useActor<TLogic extends AnyActorLogic>(
  logic: TLogic,
  options: ActorOptions<TLogic> = {}
): [SnapshotFrom<TLogic>, ActorRefFrom<TLogic>['send'], ActorRefFrom<TLogic>] {
  if (
    isDevelopment &&
    !!logic &&
    'send' in logic &&
    typeof logic.send === 'function'
  ) {
    throw new Error(
      `useActor() expects actor logic (e.g. a machine), but received an ActorRef. Use the useSelector(actorRef, ...) hook instead to read the ActorRef's snapshot.`
    );
  }

  const actorRef = useIdleActor(logic, options as any);

  return useActivateActor(actorRef)
}

export function useActivateActor(actorRef: ReturnType<typeof useIdleActor): [SnapshotFrom<TLogic>, ActorRefFrom<TLogic>['send'], ActorRefFrom<TLogic>]{
  const getSnapshot = useCallback(() => {
    return actorRef.getSnapshot();
  }, [actorRef]);

  const subscribe = useCallback(
    (handleStoreChange) => {
      const { unsubscribe } = actorRef.subscribe(handleStoreChange);
      return unsubscribe;
    },
    [actorRef]
  );

  const actorSnapshot = useSyncExternalStore(
    subscribe,
    getSnapshot,
    getSnapshot
  );

  useEffect(() => {
    actorRef.start();

    return () => {
      stopRootWithRehydration(actorRef);
    };
  }, [actorRef]);

  return [actorSnapshot, actorRef.send, actorRef] as any;
}

And then the counter component would use useActivateActor so that nothing starts until the counter is mounted, and everything is cleaned up when the counter is removed.

import { useActivateActor } from "@xstate/react";
import { useMolecule } from "bunshi/react";
import React from "react";
import { CountMolecule } from "./molecules";

function Counter() {
  const countActor = useMolecule(CountMolecule);
  const [state, send] = useActivateActor(countActor);
  const increment = () => send({ type: 'INCREMENT' });
  return (
    <div>
      <p>Times clicked: {state.context.count}</p>
      <button onClick={increment}>Increment count</button>
    </div>
  );
}

Unfortunately this solution is incomplete, because each component that uses useActivateActor could start or stop the shared countActor.

[loganvolkers]

Put together a proof of concept for cleanup with a new mounted function. This needs to exists because the molecules should be stateless functions that can be re-run.

Here's an example of the syntax:

export const TimerMolecule = molecule((_, scope) => {
  scope(ComponentScope);

  const timerMachine = createMachine({
    id: "counter",
    context: { count: 0 },
    initial: "ticking",
    states: {
      ticking: {
        after: {
          1000: {
            target: "ticking",
            actions: [console.log, assign({ count: (ctx) => ctx.count + 1 })]
          }
        }
      }
    }
  });

   const actor = interpret(timerMachine);
   mounted(() => {
      actor.start();
      return ()=> actor.stop();
   });

  return actor;
});

Thoughts or opinions on the syntax?

[asfktz]

Love it!

Here's my thoughts:

1. I suggest following the convention of calling it onMount (just for the sake of familiarity):

   • Jotai: https://jotai.org/docs/core/atom#:~:text=onMount
   • NanoStores: https://github.com/nanostores/nanostores#store-events:~:text=onMount
   • Solidjs: https://www.solidjs.com/docs/latest/api#onmount

2. From a user perspective (who is not familiar with the implementation details) I wonder why scope is an argument and mounted is not.

This needs to exists because the molecules should be stateless functions that can be re-run.

3. In a sense it's similar to useEffect, no? That is why Strict Mode forces us to define a cleanup function, to make sure we can re-run our effects without causing memory leaks. if so, isn't it safe to call actor.start() at the root of of the molecule and only define onCleanup?

It could be that I miss something important about the way molecules work.

[loganvolkers]

I suggest following the convention of calling it onMount (just for the sake of familiarity):

Great references. I agree with onMount. onXXXX seems to be a common way of indicating lifecycle callbacks, vue has onMounted too.

From a user perspective (who is not familiar with the implementation details) I wonder why scope is an argument and mounted is not.

There's the possibility of making scope and mol both global, too. I have even considered combining them both into a simple use function that accepts either molecules or scopes.

In a sense it's similar to useEffect, no? That is why Strict Mode forces us to define a cleanup function, to make sure we can re-run our effects without causing memory leaks. if so, isn't it safe to call actor.start() at the root of of the molecule and only define onCleanup?

On closer inspection, this might be a defect in the current implementation. Molecules are being created multiple times, but they need not be.

[asfktz]

There's the possibility of making scope and mol both global, too. I have even considered combining them both into a simple use function that accepts either molecules or scopes.

Interesting, I like this direction!

[loganvolkers]

@asfktz this has been fixed in 2.1.0-rc.2 and there's a working demo here

Can you please give it a try and make sure it works as expected for you?

[asfktz]

Hi @loganvolkers!

It seems that the machine instance continues to run even after the molecule is unmounted. Is this the intended behavior? Additionally, it mounts twice, which aligns with the expected behavior in React with Strict Mode. However, without a cleanup function, this behavior is undesirable.

https://github.com/saasquatch/bunshi/assets/199747/99a3855a-0150-4827-af31-f1f62a62b251

[loganvolkers]

Sorry I should have been more clear. There is a new onMount callback function for doing lazy starting and stopping (i.e. "cleanup")

Are you using that in your code?

I added it to the working demo: https://codesandbox.io/p/sandbox/molecule-cleanup-mvt4cg