pytransitions / transitions

A lightweight, object-oriented finite state machine implementation in Python with many extensions
MIT License
5.49k stars 524 forks source link

Don't automatically cancel running_task in AsyncMachine #462

Closed broose-goose closed 3 years ago

broose-goose commented 3 years ago

I have a use case where a bunch of AsyncMachines are queuing tasks for a consumer. Whenever the consumer finishes processing the request, it triggers a state change on the AsyncMachine, letting the AsyncMachine know its good to continue.

What I found is that, when the AsyncMachine is processing the trigger, it runs a function called switch_model_context that updates the context that the AsyncMachine is running in. This function keeps killing my main loop task whenever the machine try's to switch context from the main loop to the consumer loop. Commenting out "running_task.cancel()" solves my issue.

https://github.com/pytransitions/transitions/blob/master/transitions/extensions/asyncio.py#L352-L365

Could you add a variable defined on the AsyncMachine that I could use to turn off this behavior? Something like this

# usage
machine = AsyncMachine(model, states, allow_context_switch=True)

# updated AsyncMachine.switch_model_context
def switch_model_context(self, model):
    running_task = self.async_tasks.get(model, None)
    if self.current_context.get() != running_task:
        if running_task is not None and running_task.done() is False and\
                self.allow_context_switch is False:
            _LOGGER.debug("Cancel running tasks...")
            running_task.cancel()
        self.async_tasks[model] = self.current_context.get()

I have updated my local copy to just comment the offending line out. I apologize very unconfident in my ability to actually make this change. Currently my tests always fail because tox always downloads a fresh copy of the library.

aleneum commented 3 years ago

Hi @broose-goose,

when AsyncMachine had been implemented, cancelling tasks of the exiting context was mentioned as required feature (full discussion here #259). The function AsyncMachine.switch_model_context was intentionally introduced to customize this behaviour.

I could think of two approaches for your use case:

First, override AsyncMachine.switch_model_context:

from transitions.extensions.asyncio import AsyncMachine

class KeepTasksAliveMachine(AsyncMachine):
    def switch_model_context(self, model):
        pass

Everything in that function is just there to keep track of running tasks. If you don't need it, you don't have to deal with it.

Second option: Shielding

If you have just a particular task you want to keep running, wrapping it in asyncio.shield might help:

async def model_callback(self):
    await async.shield(model.dont_cancel_this())
broose-goose commented 3 years ago

Thanks for sending that discussion! Totally get why that feature exists, I just trust myself not to step on my own feet.

Option one works for me. Wouldn't have thought of that as I am not a python developer; thank you so much!

aleneum commented 3 years ago

Thanks for sending that discussion! Totally get why that feature exists, I just trust myself not to step on my own feet.

It's a valid question and just shows me that there is room for improvement in the documentation.

Option one works for me. Wouldn't have thought of that as I am not a python developer;

Well, you managed to make sense of someone else's codebase and located an issue for your use case. I'd say you do pretty good on the Python front 👍. You are now one of us 🐍.