Phoenix transport: breaking early from async generator loop requires canceling receive_data_task

[pzingg]

I don't know if there is a Slack channel or other forum for the gql project, so I'm posing this here as an issue. I would like to see an example of a subscription that uses one of the async transports (such as PhoenixChannelWebsocketsTransport) for the use case below. It's something that I had working with gql 2.x, after adding my own (probably naive) implementation of Phoenix channels. The target server is Absinthe GraphQL.

The inputs for the process are:

• a mutation document and variables
• a subscription document and variables (subscription updates are triggered by the mutation)
• a coroutine (or asyncio task?) that can act on each subscription update received, as indicated below
• an optional timeout, after which the subscription will be ended

Then the process is:

1. Send a subscription request with the coroutine as a parameter.
2. If the subscription is successful, send a mutation request that will trigger subscription updates.
3. As each update arrives on the subscription (inside an async generator for loop), the coroutine acts on it (for example printing, or parsing the ExecutionResult and processing results). The coroutine may also, as it parses an update, decide to end the subscription if it decides that no more updates are expected, and return any data it wants.
4. If the timeout is exceeded, the process will unsubscribe, and return any results collected by the coroutine, as well as an indication that the process ended via a timeout.

It looks like the standard way a subscription is ended is when it receives a "complete" message on the websocket, but in my case I want the client to decide when to stop subscribing.

[pzingg]

I was successful in developing a coroutine that subscribes to an Absinthe subscription, after struggling to find the correct channel name, "__absinthe__:control", to use when constructing the PhoenixChannelWebsocketsTransport. Learning from some of the test example source code, I am almost there, except I don't know how to stop the subscription inside the asynchronous generator loop.

Here's the code

async def subscribe_to_updates(doc, vars, ...):
    absinthe_control_transport = PhoenixChannelWebsocketsTransport(...)
    async with Client(transport=absinthe_control_transport) as session:
        generator = session.subscribe(doc, variable_values=vars)
        async for result in generator:
            data = result['updated']
            print(data)
            if data['phase'] in ['completed', 'canceled']:
                # want to exit coroutine, setting data as the result
                print('break')
                await generator.aclose() # is this necessary?
                return data 
    return None

The code does receive results and prints the data, but the generator.aclose() and return statements do not end the execution. Instead, the awaitable sits until a timeout of 10 seconds happens (not sure which of the websockets timeouts this is). Same thing happens if I raise an exception inside the for loop. If I raise an exception or return before the for loop, the coroutine exits immediately, although there is an assertion failure because the websocket connection has not yet been made.

[pzingg]

I fixed the code by putting a break, not a return inside the loop, and making sure that the return is outside the async with ... as session block. The code is fairly close to the test_websocket_subscription_break in test_websocket_subscription.py and it does break immediately. However, I had to insert a call to cancel the transport's receive_data_task, followed by the await session._generator.aclose().

Works on Python 3.8 without failing assertions, throwing exceptions or leaving any unhandled exceptions.

async def subscribe_to_updates(doc, vars, ...):
    absinthe_control_transport = PhoenixChannelWebsocketsTransport(...)
    data = None
    async with Client(transport=absinthe_control_transport) as session:
        async for result in session.subscribe(doc, variable_values=vars):
            data = result['updated']
            if data['phase'] in ['completed', 'canceled']:
                # Here we want to exit coroutine, returning `data` as the result
                # Note: the following steps seem to be necessary
                # to avoid assertion failures and other exceptions
                print(f'cancelling receive_data_loop')
                session.transport.receive_data_task.cancel()
                print(f'closing generator')
                await session._generator.aclose()
                break

    return data

[leszekhanusz]

generator.aclose() is not needed if you have a recent version of python (at least 3.7)

In theory, in order to exit a subscription, you can:

• simply use break
• use return
• raise an exception
• use generator.aclose()
• or even cancel the current task

In all these cases, gql should close the connection gracefully.

The PhoenixChannelTransport has not been tested in great detail and you could help us test it further.

Could you please try the PR #227 to see if it is better for you ?

If you could then retry with logs set to DEBUG for the websockets code it would help debug the problems

[pzingg]

Thanks, it was able to use the logging output to track this down.

Here are my results on Python 3.8.10:

1. cancel receive_data_task and await generator.aclose: quick exit, no exceptions
2. only cancel receive_data_task (without await generator.aclose): waits for timeout in _close_coro: starting clean_close
3. only await generator.aclose (without cancel receive_data_task): Task exception was never retrieved. assertion failed in "websockets.py", line 179, in _receive, assert self.websocket is not None
4. break without either: waits for timeout as in case 2, also fails assertion in case 3.

[leszekhanusz]

About the 10 seconds delay, we have 3 different timeout defined by default to 10 seconds:

• connect_timeout
• close_timeout
• ack_timeout

In your case, the problematic timeout seems to be the close_timeout
For confirmation, you can change the timer value in the transport init

[leszekhanusz]

But please test the PR #227 with only a break, the assertion bug should disappear.

[pzingg]

With PR #227 I could remove the receive_data_task.cancel() but still need to do the await session._generator.aclose() to avoid waiting. If I omit that, there is a 10 second timeout between these two log entries:

DEBUG:gql.transport.websockets:close: starting
DEBUG:gql.transport.websockets:_fail: starting with exception: TransportClosed('Websocket GraphQL transport closed by user')
DEBUG:gql.transport.websockets:wait_close: starting
DEBUG:gql.transport.websockets:_close_coro: starting
DEBUG:gql.transport.websockets:_close_coro: starting clean_close
...10 second pause here...
INFO:gql.transport.websockets:>>> {"topic": "__absinthe__:control", "event": "phx_leave", "payload": {}, "ref": 3}

I added some debugging to _clean_close, and can verify that in this code:

        # Send 'stop' message for all current queries
        for query_id, listener in self.listeners.items():

            if listener.send_stop:
                await self._send_stop_message(query_id)
                listener.send_stop = False

        # Wait that there is no more listeners (we received 'complete' for all queries)
        try:
            await asyncio.wait_for(self._no_more_listeners.wait(), self.close_timeout)
        except asyncio.TimeoutError:  # pragma: no cover
            pass

1. The _send_stop_message IS being awaited (there is only one listener).
2. The wait_for(_no_more_listeners.wait()) IS timing out.

[leszekhanusz]

In the websockets transport, in the subscriptions protocol you need to send a "stop" message to close a subscription on the backend.
When the user wants to close the connection (not an exception or connection closed), we want to send a stop for each open subscription and wait for the complete message for all of them before sending the final message and closing the connection.

The _no_more_listener event is used to know when all the subscriptions generators have been been properly closed.

For the phoenix transport, there is apparently no "stop" message in the protocol. So the _send_stop_message will simulate receiving a "complete" message when we need to close the connection cleanly.

In theory, in that case, all the subscriptions generators still open will receive the "complete" answer or be already closed.
So in theory at some point, len(self.listeners) will be 0 and the _no_more_listeners event will be set.

Apparently this is not the case for you...
Could you reproduce your behavior in a test case ?

[pzingg]

In the code that I had written to support Absinthe subscriptions with earlier versions of gql, I tore down the client subscription by sending an "unsubscribe" Phoenix Channel message over the websocket. This is part of the Absinthe protocol, and corresponds I guess to the "complete" type you use with vanilla Websockets.

Here is some pseudocode in the style of _send_stop_message. Do you want me to try and implement it in a fork of PR #227?

async def _send_stop_message(self, listener_query_id: int) -> None:
        """Send an unsubscribe message to Absinthe for a listener's query_id.
        """
        query_id = self.next_query_id
        self.next_query_id += 1

        subscription_id = None
        for sub_id, q_id in self.subscription_ids_to_query_ids.items():
            if q_id == listener_query_id:
                subscription_id = sub_id
                break
                
        if subscription_id is None:
            raise ValueError(f"No subscription for {listener_query_id}")

        stop_message = json.dumps(
            {
                "topic": self.channel_name, # This was "__absinthe__:control"
                "event": "unsubscribe",
                "payload": {"subscriptionId": subscription_id},
                "ref": query_id,
            }
        )

        await self._send(stop_message)

After the unsubscribe request is sent, the socket will receive a "phx_reply" from Absinthe (see below). The ref (query_id) of the request must be saved in a dict or array, so that the "phx_reply" that comes back can be mapped into a "complete" answer_type in _parse_answer.

{
 "event":"phx_reply",
 "payload":{"response":{"subscriptionId":"__absinthe__:doc:-576460..."},
 "status":"ok"},
 "ref":3, 
 "topic":"__absinthe__:control"
}

I have this unsubscribe logic working in a fork of PR #227, but it still requires that I call await generator.aclose to avoid a timeout. When _send_stop_message is called after a GeneratorExit exception, I don't see where or how the transport waits and parses the reply, which then puts the "complete" message on the queue, after which the listener is removed, if that is the sequence that should happen. Instead, in the logs, it looks like the unsubscribe reply is only picked up after the transport itself is being shut down (after "close websocket connection" below):

--- generator.aclose() called explicitly 
DEBUG:gql.transport.websockets:Exception in subscribe: GeneratorExit()
INFO:gql.transport.websockets:>>> {"topic": "__absinthe__:control", "event": "unsubscribe", "payload": {"subscriptionId": "__absinthe__:doc:-57646..."}, "ref": 3}
DEBUG:gql.transport.websockets:close: starting
DEBUG:gql.transport.websockets:_fail: starting with exception: TransportClosed('Websocket GraphQL transport closed by user')
DEBUG:gql.transport.websockets:wait_close: starting
DEBUG:gql.transport.websockets:_close_coro: starting
DEBUG:gql.transport.websockets:_close_coro: starting clean_close
-- leaving the channel but the unsubscribe ack has not been processed!
INFO:gql.transport.websockets:>>> {"topic": "__absinthe__:control", "event": "phx_leave", "payload": {}, "ref": 4}
DEBUG:gql.transport.websockets:_close_coro: sending exception to listeners
DEBUG:gql.transport.websockets:_close_coro: close websocket connection
-- messages that were still in the socket are finally received!
INFO:gql.transport.websockets:<<< {"event":"phx_reply","payload":{"response":{"subscriptionId":"__absinthe__:doc:-576460752303421754:A0E9DB4A02934BAF259C928DDCE706242F67FAB0D01923BEFA9857AC883E3507"},"status":"ok"},"ref":3,"topic":"__absinthe__:control"}
INFO:gql.transport.websockets:<<< {"event":"phx_reply","payload":{"response":{},"status":"ok"},"ref":4,"topic":"__absinthe__:control"}
INFO:gql.transport.websockets:<<< {"event":"phx_close","payload":{},"ref":null,"topic":"__absinthe__:control"}
DEBUG:gql.transport.websockets:close: starting
DEBUG:gql.transport.websockets:_fail: starting with exception: TransportClosed('Websocket GraphQL transport closed by user')
DEBUG:gql.transport.websockets:close_task is not None in _fail. Previous exception is: TransportClosed('Websocket GraphQL transport closed by user') New exception is: TransportClosed('Websocket GraphQL transport closed by user')
DEBUG:gql.transport.websockets:wait_close: starting
DEBUG:gql.transport.websockets:_close_coro: websocket connection closed
DEBUG:gql.transport.websockets:_close_coro: start cleanup
DEBUG:gql.transport.websockets:_close_coro: exiting
DEBUG:gql.transport.websockets:wait_close: done
DEBUG:gql.transport.websockets:close: done
DEBUG:gql.transport.websockets:wait_close: done
DEBUG:gql.transport.websockets:close: done
DEBUG:gql.transport.websockets:Exiting _receive_data_loop()

[leszekhanusz]

Feel free to do a PR with your fork so that we could work together on the same branch.

If you call generator.aclose(), then you send the unsubscribe directly. It does not mean that the connection should close at this point as there could be other subscriptions still in progress. If this is your only subscription, then _no_more_listeners will be set directly. We don't wait for the reply here and delete the corresponding listener.

Then you close the transport, it does a clean_close, and because there is no other listeners, it does not need to send stop messages and will directly send the phx_leave message.

So your log seems perfectly fine, this is the same behavior with the websockets transport, we wait for the stop reply only if the subscription is still open in order to put a "complete" in the listener queue and end the subscription generator cleanly.

This is not normal that you need to call generator.aclose() manually to avoid a delay. If you could reproduce that in a test case I would be able to help debug it.

[pzingg]

I will run the test suite on my changes, and if all the Phoenix Channel tests pass, I will try to build a smallest possible Absinthe server and test case to see what happens. Then I'll put in a PR for this repo.