rx.background and StateManager.modify_state provides safe exclusive access to state

[masenf]

App.modify_state provides exclusive access to state AND sends a state update after exiting the async contextmanager.

Background EventHandler example

@rx.background enables an EventHandler to not block the event loop!

A background task can either yield events that will update the state or enter async with self context block to directly munge the state.

import asyncio

import reflex as rx


class State(rx.State):
    counter: int = 0
    _task_id: int = 0

    @rx.background
    async def handle_event(self):
        async with self:
            task_id = self._task_id
            self._task_id += 1
        print(f"[{task_id}] background event starting; counter = {self.counter}")
        for ix in range(5):
            await asyncio.sleep(2)
            print(f"[{task_id}] [{ix}] background event stale counter = {self.counter}")
            async with self:
                print(f"[{task_id}] [{ix}] background event bumping counter = {self.counter}")
                self.counter += 1
            yield rx.console_log(f"Background task {task_id} is on iteration {ix}")

    @rx.background
    async def handle_event_yield_only(self):
        async with self:
            task_id = self._task_id
            self._task_id += 1
        print(f"[{task_id}] background event starting; counter = {self.counter}")
        for ix in range(5):
            await asyncio.sleep(2)
            yield State.increment_arbitrary(5)
            yield rx.console_log(f"Background task {task_id} is on iteration {ix}")

    def increment(self):
        self.counter += 1

    def increment_arbitrary(self, amount: int):
        self.counter += int(amount)

    def reset_counter(self):
        self.counter = 0

    async def chain_background_task(self, direct: bool):
        if direct:
            await self.non_blocking_pause()
        else:
            # must yield / return, no direct calling of background tasks
            yield State.non_blocking_pause()

    async def blocking_pause(self):
        await asyncio.sleep(5)
        return rx.console_log("Blocking pause done")

    @rx.event.background
    async def non_blocking_pause(self):
        await asyncio.sleep(5)
        return rx.console_log("Non-blocking pause done")


def index() -> rx.Component:
    return rx.vstack(
        rx.heading(State.counter),
        rx.button("Delayed Increment", on_click=State.handle_event),
        rx.button("Yield Increment", on_click=State.handle_event_yield_only),
        rx.button("Increment", on_click=State.increment),
        rx.button("Blocking Pause", on_click=State.blocking_pause),
        rx.button("Non-Blocking Pause", on_click=State.non_blocking_pause),
        rx.button("Chain background", on_click=lambda: State.chain_background_task(False)),
        rx.button("Chain background direct", on_click=lambda: State.chain_background_task(True)),
        rx.button("Reset", on_click=State.reset_counter),
    )


app = rx.App()
app.add_page(index)
app.compile()

Background task example app

import asyncio
import uuid

import reflex as rx


TASKS = set()


def background_task(token) -> asyncio.Task:
    task_id = uuid.uuid4()
    async def _task():
        for ix in range(10):
            await asyncio.sleep(1)
            async with app.modify_state(token) as state:
                state.messages.append(f"[{task_id}] [{ix}] Hello to {token}")
        print(f"Task {task_id} finished")
    return asyncio.create_task(_task())


class State(rx.State):
    messages: list[str] = []
    counter: int = 0

    async def start_background_task(self):
        """Start a background task."""
        t = background_task(self.get_token())
        t.add_done_callback(TASKS.discard)
        TASKS.add(t)

    async def increment_counter(self):
        await asyncio.sleep(0.25)
        self.counter += 1

    @rx.var
    def reversed_messages(self) -> list[str]:
        return self.messages[::-1]


def index() -> rx.Component:
    return rx.fragment(
        rx.color_mode_button(rx.color_mode_icon(), float="right"),
        rx.vstack(
            rx.heading(f"Counter: {State.counter}"),
            rx.button("Start background task", on_click=State.start_background_task),
            rx.button("Increment counter", on_click=State.increment_counter),
            rx.foreach(
                State.reversed_messages,
                lambda m: rx.text(m),
            ),
        ),
    )


app = rx.App()
app.add_page(index)
app.compile()

Run it with Redis?

docker run -d -p 6379:6379 redis
REDIS_URL=localhost:6379 reflex run

Bugs

Click Increment counter about 10 times quickly, then click Start Background Task... then it deadlocks 😔 Fixed via b471403

Remaining Work

• redis lock expiry
• unit tests for StateProxy and background events
• end-to-end integration tests w/ redis

[picklelo]

Awesome, working really well for me! I think the async with self syntax is nice. I didn't review the state.py file yet, will add comments to that

[masenf]

@picklelo thanks for the feedback. I addressed almost everything here: https://github.com/reflex-dev/reflex/pull/1676/files/f6ba2ab4eb855837ae517e05ffc187765bbaaa93..1093229068b6d7f5aec2917c5d08b0ad43fb4211

I also went ahead and fixed all of the AppHarness tests to work with redis, so now we're running with both StateManagerMemory and StateManagerRedis for better coverage. This will also ensure that future tests and features are at least considering how they interop with redis backing store. The refactoring of the StateManager as an abstract base class also opens the door to simpler implementation of future backing stores if we ever need or want that.

As for async with self I've been pondering over potentially better syntax, but i haven't found anything i particularly like. I don't really want to expose an underscore method as part of the API, but I also am reluctant to restrict downstream users from another name. And I'm a sucker for the conciseness of async with self. I think we should put that one to the community and see what people think.

As an aside, I had to kick the CI on reflex-web a few times; I wouldn't think it's my changes here, but the error is weird and I don't really understand it. Edit: same error on a different PR that's already merged.

[masenf]

Further testing has revealed that this sort of breaks yielding updates in normal event handlers. Need to write a test case and investigate.

Fixed and test case written ✔️

[picklelo]

Awesome!!

[picklelo]

Why do we set this if we're overriding it again in the __init__?

[masenf]

I didn't want to make it Optional and default it to None because then we have to put is not None checks all over the place.

But if we do something like StateManager.create(state=DefaultState) here, and redis is used, then we create an unserializable default, which pydantic doesn't like.

So the default StateManager is the StateManagerMemory, then inside __init__, we call StateManager.create(...) to update the reference to either memory or redis.

[picklelo]

Awesome, this is much cleaner we can extend to other backends in the future too