How to create a pair of coroutines, that either all succeed or all fail?

Question

I would like to implement two coroutine functions, where one passes a value to the other. Return or exception is used to denote whether the transfer is successful. How to ensure that two corountines always agree on the result of the transfer, in any condition?

Below is an partially working implementation.

async def prevent_return_leak(f):
    assert isinstance(f, asyncio.Future)
    try:
        return await f
    except asyncio.CancelledError:
        if f.cancelled():
            raise
        asyncio.current_task().cancel() # delay to next yield
        return f.result()

class Exchange:
    '''
    Transfer a value. There can be at most one unfinished `give`
    call at any time. The same applies to `take`.
    '''
    def __init__(self):
        self.taker = None
        self.giver = None
        self.value = None

    async def take(self):
        '''
        Wait until `give(value)` is called, then return `value`.
        If this function returns, the `give` call from which `value`
        is taken is guaranteed to return.
        '''
        if self.giver is not None:
            if not self.giver.done():
                self.giver.set_result(None)
                return self.value
        f = asyncio.Future()
        assert self.taker is None
        self.taker = f
        try:
            return await prevent_return_leak(f)
        finally:
            assert self.taker is f
            self.taker = None

    async def give(self, value):
        '''
        Wait until `value` is taken by a `take` call. If this function
        returns, exactly one `take` call is guaranteed to return with
        `value`.
        '''
        if self.taker is not None:
            if not self.taker.done():
                self.taker.set_result(value)
                return
        f = asyncio.Future()
        assert self.giver is None
        self.giver = f
        try:
            return await prevent_return_leak(f)
        finally:
            assert self.giver is f
            self.giver = None

A simple test case:

async def take(e): 
    print(await e.take())
    print('take done')

async def test():
    e = Exchange()
    t1 = asyncio.create_task(take(e))
    await asyncio.sleep(0) # run event loop once
    t2 = asyncio.create_task(e.give(1))
    await asyncio.sleep(0) # run event loop once
    print(t1) 
    print(t2) 
    t1.cancel() 
    print(await t1) 
    print('test done')

loop.run_until_complete(test())                                                                                                              
# :2> wait_for=>
# :40> result=None>
# value
# take done
# ---------------------------------------------------------------------------
# CancelledError                            Traceback (most recent call last)

As you may see, take() finished normally, which is expected. However, test() is cancelled and the value get lost when take() returns to test(). This is due to asyncio.Task marking itself as cancelled if the wrapped coroutine requests self-cancellation just before return.

I event tried to patch asyncio.Task, but it only works for the pure python version, i.e. asyncio.tasks._PyTask. Is it even possible with asyncio?

How to create a pair of coroutines, that either all succeed or all fail?

0 Answers