Generalise the @cancellable annotation so it can be used on functions other than just servlet methods. (#13662)

This commit is contained in:
reivilibre 2022-08-31 11:16:05 +00:00 committed by GitHub
parent a160406d24
commit 7bc110a19e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 75 additions and 75 deletions

View file

@ -33,7 +33,6 @@ from typing import (
Optional,
Pattern,
Tuple,
TypeVar,
Union,
)
@ -64,6 +63,7 @@ from synapse.logging.context import defer_to_thread, preserve_fn, run_in_backgro
from synapse.logging.opentracing import active_span, start_active_span, trace_servlet
from synapse.util import json_encoder
from synapse.util.caches import intern_dict
from synapse.util.cancellation import is_function_cancellable
from synapse.util.iterutils import chunk_seq
if TYPE_CHECKING:
@ -94,68 +94,6 @@ HTML_ERROR_TEMPLATE = """<!DOCTYPE html>
HTTP_STATUS_REQUEST_CANCELLED = 499
F = TypeVar("F", bound=Callable[..., Any])
_cancellable_method_names = frozenset(
{
# `RestServlet`, `BaseFederationServlet` and `BaseFederationServerServlet`
# methods
"on_GET",
"on_PUT",
"on_POST",
"on_DELETE",
# `_AsyncResource`, `DirectServeHtmlResource` and `DirectServeJsonResource`
# methods
"_async_render_GET",
"_async_render_PUT",
"_async_render_POST",
"_async_render_DELETE",
"_async_render_OPTIONS",
# `ReplicationEndpoint` methods
"_handle_request",
}
)
def cancellable(method: F) -> F:
"""Marks a servlet method as cancellable.
Methods with this decorator will be cancelled if the client disconnects before we
finish processing the request.
During cancellation, `Deferred.cancel()` will be invoked on the `Deferred` wrapping
the method. The `cancel()` call will propagate down to the `Deferred` that is
currently being waited on. That `Deferred` will raise a `CancelledError`, which will
propagate up, as per normal exception handling.
Before applying this decorator to a new endpoint, you MUST recursively check
that all `await`s in the function are on `async` functions or `Deferred`s that
handle cancellation cleanly, otherwise a variety of bugs may occur, ranging from
premature logging context closure, to stuck requests, to database corruption.
Usage:
class SomeServlet(RestServlet):
@cancellable
async def on_GET(self, request: SynapseRequest) -> ...:
...
"""
if method.__name__ not in _cancellable_method_names and not any(
method.__name__.startswith(prefix) for prefix in _cancellable_method_names
):
raise ValueError(
"@cancellable decorator can only be applied to servlet methods."
)
method.cancellable = True # type: ignore[attr-defined]
return method
def is_method_cancellable(method: Callable[..., Any]) -> bool:
"""Checks whether a servlet method has the `@cancellable` flag."""
return getattr(method, "cancellable", False)
def return_json_error(
f: failure.Failure, request: SynapseRequest, config: Optional[HomeServerConfig]
) -> None:
@ -389,7 +327,7 @@ class _AsyncResource(resource.Resource, metaclass=abc.ABCMeta):
method_handler = getattr(self, "_async_render_%s" % (request_method,), None)
if method_handler:
request.is_render_cancellable = is_method_cancellable(method_handler)
request.is_render_cancellable = is_function_cancellable(method_handler)
raw_callback_return = method_handler(request)
@ -551,7 +489,7 @@ class JsonResource(DirectServeJsonResource):
async def _async_render(self, request: SynapseRequest) -> Tuple[int, Any]:
callback, servlet_classname, group_dict = self._get_handler_for_request(request)
request.is_render_cancellable = is_method_cancellable(callback)
request.is_render_cancellable = is_function_cancellable(callback)
# Make sure we have an appropriate name for this handler in prometheus
# (rather than the default of JsonResource).