3

Writing Redux Reducers in Rust

 2 years ago
source link: https://fiberplane.dev/blog/writing-redux-reducers-in-rust/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

Writing Redux Reducers in Rust

writing-redux-reducers-in-rust.png
arend.jpg

Arend van Beelen

Principal Software Engineer

Wed Apr 06 2022 - 9 min. read

Introduction

At Fiberplane, we are building a collaborative notebook editor for incident response and infrastructure debugging. As part of the challenges we faced implementing such an editor, we had several complex functions written in Rust that we had to use on the frontend through WebAssembly.

This post will explore how we integrated this WASM code into a React/Redux application, as well as why we ended up writing our own bindings generator for it.

Initial solution: Wasm-bindgen to the rescue

The way we initially integrated our Rust code was relatively straightforward: We simply used wasm-bindgen, the go-to tool for integrating Rust WASM code in the browser. The functions that implemented our business logic were pure functions that we invoked from our Redux thunks and reducers as appropriate. All the code that dealt with Redux directly was still being written in TypeScript.

Everything worked but some pain points quickly emerged. While wasm-bindgen allows for passing complex data structures back and forth through JSON serialization, we ran into some practical limitations:

  • Rust structs were exposed as classes to TypeScript. This didn’t match our usage of plain-old objects that we wished to store in our Redux state. Worse, type-safe enums were not supported at all. This led us to use ts-rs for generating our TypeScript types, with manually written glue code where we had to pass objects with types generated from ts-rs into functions generated by wasm-bindgen. Again, things worked, but the solution was becoming increasingly brittle.
  • Serialization overhead also became an issue. Exposing pure functions was great for testability but some operations were large and passing them back and forth across the WASM bridge repeatedly became a bottleneck for us. This was especially problematic when a client had to perform conflict resolution. This is a complex procedure that involves several round-trips in and out of WASM. This eventually made us rethink how we did our state management...

Iteration: Moving the Reducers into Rust

When your core logic is written in Rust, while your state is managed in TypeScript, you have an impedance mismatch. Every time you want to invoke your logic, you pay the price of serializing the relevant state back and forth. We could not escape this problem entirely (we still needed to have state in TypeScript, so that React could render it), but we minimized its impact by moving our state into Rust.

How we did this, you might ask?

1. Write Reducers in Rust

First, we ported our reducers to Rust. This resulted in reducers with signatures such as this one:

Rust
fn reducer(state: &State, action: Action) -> (State, Vec<SideEffectDescriptor>)
Copy

These are pure functions, just as reducers should be. They take an existing state and an action, and return a new state. They also return a list of side effect descriptors, which we’ll talk about later.

2. Expose Reducers to TypeScript

Next, we expose these reducers to TypeScript, while avoiding having to serialize the state across the bridge:

Rust
static mut STATE: Lazy<RefCell<State>> = Lazy::new(|| RefCell::new(State::default()));
#[fp_export_impl(protocol_bindings)]
fn reducer_bridge(action: Action) -> ReducerResult {
// This is safe as long as we only call this from a single-threaded WASM context:
unsafe {
let old_state = STATE.get_mut();
let (new_state, side_effects) = reducer(old_state, action);
let state_update = StateUpdate::from_states(old_state, &new_state);
STATE.replace(new_state);
ReducerResult {
state_update,
side_effects,
Copy

There’s a couple of things going on here:

  • On the first line of code we see the state instance that lives in Rust. It’s a global mutable variable. This is generally frowned upon because this is unsafe in a multi-threaded environment. For us it’s okay, since we only use this in a single-threaded WASM environment.
  • Next, we see the reducer function itself, which has a few noteworthy features:
    • No state is passed in. This way only actions still need to be serialized on reducer invocations, while the global state is injected in the body. And after the call to the original reducer, the global state is replaced with the new state.
    • A ReducerResult type is returned which contains the side effect descriptors and a state_update field. This state_update is effectively a diff between the old state and the new state, so that we only need to serialize parts of the new state that have actually changed. If your state is small, you probably can get away with simply returning the full new state. Unfortunately, our state is pretty large at this point, so that was not an option for us.
    • An #[fp_export_impl] annotation that provides a little spoiler for our new bindings generator, which we will discuss below.

3. Call the Reducer from TypeScript

We still have regular Redux reducers in TypeScript, as not all of our state slices live in Rust (let’s face it, not all reducers need to be in Rust, and it’s simpler not to move them). For those that do live in Rust, the respective TypeScript reducers invoke their Rust counterparts:

TypeScript
export default function reducer(
state = initialState,
action: Action
): StateWithSideEffects<State> {
const result = reducerBridge(action); // This calls the Rust reducer.
state = stateUpdateReducer(state, result.stateUpdate);
const { sideEffects } = result;
return { state, sideEffects };
Copy

The main takeaway here is that we invoke the Rust reducer without passing along the state but instead we update the Redux state using the stateUpdate that we received as part of the ReducerResult (our state_update field from the previous section, as our bindings automatically convert between casing conventions). We use a separate stateUpdateReducer() that applies the returned update, which I omitted for brevity.

At this point, you can see that we do indeed have two copies of our state: One in TypeScript and one in Rust. The TypeScript one is used for rendering in React, while the one in Rust is the real source of truth that our core logic operates on. Having two copies of the state does consume some more memory but, overall, this is negligible for us. One downside though, the definition of initialState in TypeScript, is still something we need to manually keep in sync with the State::default() implementation we have in Rust. From there, any updates to the state are kept in sync by the reducers.

4. Handle Side Effects

We would be done by now, if it weren’t for those side effect descriptors we’ve been carrying around. Remember I said our initial solution used to call Rust functions from Redux thunks? The reason we called those (pure) functions from thunks rather than a reducer was because their result would trigger new side effects. But now that all the Rust code is called strictly inside the reducer, it can’t trigger side effects anymore. At least, not directly.

So, instead, we let the reducer return side effect descriptors, simple objects not entirely unlike Redux actions. By letting our reducer return these objects, it gains the ability to decide which side effects to trigger, without sacrificing its functional purity.

Implementation-wise, side effect descriptors are stored in the Redux store like any other state slice. And then we use a custom middleware to trigger the actual effects:

TypeScript
export const sideEffectDispatcher =
(store: Store<RootState, any>) => (next: any) => (action: any) => {
const oldSideEffects = store.getState()?.sideEffects;
// Call the reducers:
const result = next(action);
const { sideEffects } = store.getState();
if (sideEffects && sideEffects !== oldSideEffects) {
// Trigger the side effects:
for (const descriptor of sideEffects) {
store.dispatch(thunkForDescriptor(descriptor));
return result;
Copy

Here, thunkForDescriptor() is a function that looks a bit like a reducer. It contains a large switch (descriptor.type) { ... } statement that returns the appropriate thunk for every type of side effect.

Along came fp-bindgen...

You thought we were almost done, didn’t you? Still, there is one more thing to discuss that is central to how we expose our Rust reducers today.

Remember how I mentioned things were starting to become brittle when we had to manually wire up wasm-bindgen and ts-rs? We still had to write some glue code that often wasn’t type-safe. Now, imagine how that worked out when we had most of our state management and our most important reducers all moved into Rust... Yeah, indeed.

Meanwhile, our business also had a need for full-stack WASM plugins. wasm-bindgen could not support that use case because it assumes a JavaScript host environment, while our servers are all running Rust. So, we created our own bindings generator called fp-bindgen.

As we already had quite some experience in integrating these various tools, we made sure fp-bindgen could generate the bindings for our Rust reducers as well. By unifying the tooling, we improved type safety and avoided maintaining hand-written glue code. At the same time, we also made it possible to efficiently pass data between our plugins and the code behind our Redux reducers.

fp-bindgen is now open-source and we invite you to try it out. If you are interested in full-stack WASM plugins, or you have a need for advanced Rust ↔ TypeScript interoperability, this might be for you.

Wrapping up

Writing Redux Reducers in Rust is not for the faint of heart. There are lots of challenges you may face. And unless you have a clear business case, I would not recommend taking this road lightly. You will have to deal with quite some complexity to get this up and running.

That said, if you do take this route, you may find Rust to be a remarkably pleasant language to write reducers in:

  • You do not need to worry about accidental mutations, as the Rust compiler will enforce immutability requirements for you.
  • Convenience similar to using Immer — an amazing tool — but without even needing a library.
  • General niceties such as if-else expressions and pattern matching.

Plus you get the benefit of sharing code with your backend, without being tied to Node.js. For us, it’s worth it.

Should you find yourself in a similar situation to us, we hope you found this helpful. fp-bindgen can be found on GitHub: fp-bindgen and feel free to join our Discord with any questions you might have!

And if working on challenges such as these sounds exciting to you, we’re hiring!


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK