Call immutable web API only once when previous call blocks (#3085)

Author: wjsi

mars/services/cluster/api/web.py

@@ -44,7 +44,7 @@ def _convert_node_dict(node_info_list: Dict[str, Dict]):
             res[node_addr] = res_dict
         return res
 
-    @web_api("nodes", method=["get", "post"])
+    @web_api("nodes", method=["get", "post"], cache_blocking=True)
     async def get_nodes_info(self):
         watch = bool(int(self.get_argument("watch", "0")))
         env = bool(int(self.get_argument("env", "0")))
@@ -104,7 +104,7 @@ async def get_nodes_info(self):
             result["nodes"] = self._convert_node_dict(nodes)
         self.write(json.dumps(result))
 
-    @web_api("bands", method="get")
+    @web_api("bands", method="get", cache_blocking=True)
     async def get_all_bands(self):
         role_arg = self.get_argument("role", None)
         role = NodeRole(int(role_arg)) if role_arg is not None else None
@@ -135,19 +135,19 @@ async def get_all_bands(self):
                 )
             )
 
-    @web_api("versions", method="get")
+    @web_api("versions", method="get", cache_blocking=True)
     async def get_mars_versions(self):
         cluster_api = await self._get_cluster_api()
         self.write(json.dumps(list(await cluster_api.get_mars_versions())))
 
-    @web_api("pools", method="get")
+    @web_api("pools", method="get", cache_blocking=True)
     async def get_node_pool_configs(self):
         cluster_api = await self._get_cluster_api()
         address = self.get_argument("address", "") or None
         pools = list(await cluster_api.get_node_pool_configs(address))
         self.write(json.dumps({"pools": pools}))
 
-    @web_api("stacks", method="get")
+    @web_api("stacks", method="get", cache_blocking=True)
     async def get_node_thread_stacks(self):
         cluster_api = await self._get_cluster_api()
         address = self.get_argument("address", "") or None

mars/services/task/api/web.py

@@ -91,22 +91,25 @@ async def submit_tileable_graph(self, session_id: str):
         )
         self.write(task_id)
 
-    @web_api("", method="get")
+    @web_api("", method="get", cache_blocking=True)
     async def get_task_results(self, session_id: str):
         progress = bool(int(self.get_argument("progress", "0")))
         oscar_api = await self._get_oscar_task_api(session_id)
         res = await oscar_api.get_task_results(progress=progress)
         self.write(json.dumps({"tasks": [_json_serial_task_result(r) for r in res]}))
 
     @web_api(
-        "(?P<task_id>[^/]+)", method="get", arg_filter={"action": "fetch_tileables"}
+        "(?P<task_id>[^/]+)",
+        method="get",
+        arg_filter={"action": "fetch_tileables"},
+        cache_blocking=True,
     )
     async def get_fetch_tileables(self, session_id: str, task_id: str):
         oscar_api = await self._get_oscar_task_api(session_id)
         res = await oscar_api.get_fetch_tileables(task_id)
         self.write(serialize_serializable(res))
 
-    @web_api("(?P<task_id>[^/]+)", method="get")
+    @web_api("(?P<task_id>[^/]+)", method="get", cache_blocking=True)
     async def get_task_result(self, session_id: str, task_id: str):
         oscar_api = await self._get_oscar_task_api(session_id)
         res = await oscar_api.get_task_result(task_id)
@@ -116,19 +119,24 @@ async def get_task_result(self, session_id: str, task_id: str):
         "(?P<task_id>[^/]+)/tileable_graph",
         method="get",
         arg_filter={"action": "get_tileable_graph_as_json"},
+        cache_blocking=True,
     )
     async def get_tileable_graph_as_json(self, session_id: str, task_id: str):
         oscar_api = await self._get_oscar_task_api(session_id)
         res = await oscar_api.get_tileable_graph_as_json(task_id)
         self.write(json.dumps(res))
 
-    @web_api("(?P<task_id>[^/]+)/tileable_detail", method="get")
+    @web_api("(?P<task_id>[^/]+)/tileable_detail", method="get", cache_blocking=True)
     async def get_tileable_details(self, session_id: str, task_id: str):
         oscar_api = await self._get_oscar_task_api(session_id)
         res = await oscar_api.get_tileable_details(task_id)
         self.write(json.dumps(res))
 
-    @web_api("(?P<task_id>[^/]+)/(?P<tileable_id>[^/]+)/subtask", method="get")
+    @web_api(
+        "(?P<task_id>[^/]+)/(?P<tileable_id>[^/]+)/subtask",
+        method="get",
+        cache_blocking=True,
+    )
     async def get_tileable_subtasks(
         self, session_id: str, task_id: str, tileable_id: str
     ):
@@ -139,7 +147,12 @@ async def get_tileable_subtasks(
         )
         self.write(json.dumps(res))
 
-    @web_api("(?P<task_id>[^/]+)", method="get", arg_filter={"action": "progress"})
+    @web_api(
+        "(?P<task_id>[^/]+)",
+        method="get",
+        arg_filter={"action": "progress"},
+        cache_blocking=True,
+    )
     async def get_task_progress(self, session_id: str, task_id: str):
         oscar_api = await self._get_oscar_task_api(session_id)
         res = await oscar_api.get_task_progress(task_id)

mars/services/web/core.py

@@ -49,20 +49,30 @@ class _WebApiDef(NamedTuple):
 
 
 def web_api(
-    sub_pattern: str, method: Union[str, List[str]], arg_filter: Optional[Dict] = None
+    sub_pattern: str,
+    method: Union[str, List[str]],
+    arg_filter: Optional[Dict] = None,
+    cache_blocking: bool = False,
 ):
     if not sub_pattern.endswith("$"):  # pragma: no branch
         sub_pattern += "$"
     methods = method if isinstance(method, list) else [method]
 
     def wrapper(func):
         @functools.wraps(func)
-        async def wrapped(self, *args, **kwargs):
+        async def wrapped(self: "MarsServiceWebAPIHandler", *args, **kwargs):
             try:
-                res = func(self, *args, **kwargs)
-                if inspect.isawaitable(res):
-                    res = await res
+                if not inspect.iscoroutinefunction(func):
+                    return func(self, *args, **kwargs)
+                elif not cache_blocking or self.request.method.lower() != "get":
+                    res = await func(self, *args, **kwargs)
+                else:
+                    res = await self._create_or_get_url_future(
+                        func, self, *args, **kwargs
+                    )
                 return res
+            except GeneratorExit:
+                raise
             except:  # noqa: E722  # nosec  # pylint: disable=bare-except
                 exc_type, exc, tb = sys.exc_info()
                 err_msg = (
@@ -102,8 +112,9 @@ async def _get_api_by_key(
 
 
 class MarsServiceWebAPIHandler(MarsRequestHandler):
-    _root_pattern = None
-    _method_to_handlers = None
+    _root_pattern: str = None
+    _method_to_handlers: Dict[str, Dict[Callable, _WebApiDef]] = None
+    _uri_to_futures: Dict[str, asyncio.Task] = None
 
     def __init__(self, *args, **kwargs):
         self._collect_services()
@@ -119,6 +130,21 @@ def _get_api_by_key(
             with_key_arg=with_key_arg,
         )
 
+    def _create_or_get_url_future(self, func, *args, **kw):
+        if self._uri_to_futures is None:
+            type(self)._uri_to_futures = dict()
+
+        uri = self.request.uri
+        if uri in self._uri_to_futures:
+            return self._uri_to_futures[uri]
+
+        def _future_remover(_fut):
+            self._uri_to_futures.pop(uri, None)
+
+        task = self._uri_to_futures[uri] = asyncio.create_task(func(*args, **kw))
+        task.add_done_callback(_future_remover)
+        return task
+
     @classmethod
     def _collect_services(cls):
         if cls._method_to_handlers is not None:
@@ -144,9 +170,7 @@ def get_root_pattern(cls):
 
     @functools.lru_cache(100)
     def _route_sub_path(self, http_method: str, sub_path: str):
-        handlers = self._method_to_handlers[
-            http_method.lower()
-        ]  # type: Dict[Callable, _WebApiDef]
+        handlers = self._method_to_handlers[http_method.lower()]
         method, kwargs = None, None
         for handler_method, web_api_def in handlers.items():
             match = web_api_def.sub_pattern_compiled.match(sub_path)

mars/services/web/tests/test_core.py

@@ -28,6 +28,7 @@
 class TestAPIHandler(MarsServiceWebAPIHandler):
     __test__ = False
     _root_pattern = "/api/test/(?P<test_id>[^/]+)"
+    _call_counter = 0
 
     @web_api("", method="get")
     def get_method_root(self, test_id):
@@ -58,6 +59,12 @@ async def get_with_timeout(self, test_id):
         await asyncio.sleep(100)
         raise ValueError(test_id)
 
+    @web_api("subtest_delay_cache", method="get", cache_blocking=True)
+    async def get_with_blocking_cache(self, test_id):
+        await asyncio.sleep(1)
+        type(self)._call_counter += 1
+        self.write(test_id)
+
 
 @pytest.fixture
 async def actor_pool():
@@ -138,6 +145,17 @@ def url_recorder(request):
             f"http://localhost:{web_port}/api/test/test_id/subtest_error"
         )
 
+    # test multiple request into long immutable requests
+    req_uri = f"http://localhost:{web_port}/api/test/test_id/subtest_delay_cache"
+    tasks = [asyncio.create_task(client.fetch(req_uri)) for _ in range(2)]
+    await asyncio.sleep(0.5)
+    assert TestAPIHandler._call_counter == 0
+    assert len(TestAPIHandler._uri_to_futures) == 1
+
+    await asyncio.gather(*tasks)
+    assert TestAPIHandler._call_counter == 1
+    assert len(TestAPIHandler._uri_to_futures) == 0
+
     with pytest.raises(TimeoutError):
         await client.fetch(
             f"http://localhost:{web_port}/api/test/test_id/subtest_delay",