Skip to content

Commit

Permalink
Maintain a strong reference to rpc_tasks.
Browse files Browse the repository at this point in the history
Only a weakref is kept in the loop. See python docs on creating tasks.
https://docs.python.org/3/library/asyncio-task.html#creating-tasks

This fixes an issue with request hangs when tasks were being garbage collected.

PiperOrigin-RevId: 634035959
  • Loading branch information
OskarBun authored and Copybara-Service committed May 15, 2024
1 parent bc5570b commit 05ae7fb
Showing 1 changed file with 18 additions and 1 deletion.
19 changes: 18 additions & 1 deletion src/python/grpcio/grpc/_cython/_cygrpc/aio/server.pyx.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -780,6 +780,16 @@ async def _schedule_rpc_coro(object rpc_coro,
), name="HandleExceptions[%s]" % _decode(rpc_state.method()))
_add_callback_handler(rpc_task, rpc_state)
await _handle_cancellation_from_core(rpc_task, rpc_state, loop)
try:
# Propagate any errors not handled by _handle_exceptions. If not awaited
# there will be logs of the form "Task exception was never retrieved".
# Catching it here we can provide traceback and debugging logs.
await rpc_task
except:
_LOGGER.exception('Exception not handled by _handle_exceptions in servicer method [%s]' % (
_decode(rpc_state.method()),
))
traceback.print_exc()


async def _handle_rpc(list generic_handlers, tuple interceptors,
Expand Down Expand Up @@ -954,6 +964,7 @@ cdef class AioServer:
self._server.start(backup_queue=False)
cdef RPCState rpc_state
server_started.set_result(True)
rpc_tasks = set()

while True:
# When shutdown begins, no more new connections.
Expand Down Expand Up @@ -985,9 +996,15 @@ cdef class AioServer:
rpc_coro,
rpc_state,
self._loop
)
),
name="rpc_task",
)

# loop.create_task only holds a weakref to the task.
# Maintain reference to tasks to avoid garbage collection.
rpc_tasks.add(rpc_task)
rpc_task.add_done_callback(rpc_tasks.discard)

if self._limiter is not None:
self._limiter.decrease_once_finished(rpc_task)

Expand Down

0 comments on commit 05ae7fb

Please sign in to comment.