From c3117223c5a552087ae5644662f849374005092c Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Wed, 10 Jan 2024 13:26:05 +0000
Subject: [PATCH] graphql-ws: optimise the "resolve" routine

* The resolve routine is implmented recursively in the graphql-ws
  library.
* Because the function is async this results in a large number of async
  tasks being created when the library is used with large, deeply nested
  schema.
* Async tasks have an overhead, above that of regular function calls.
* For the example in https://github.com/cylc/cylc-uiserver/issues/547
  this resulted in over 10 seconds of overheads.
---
 cylc/uiserver/websockets/resolve.py | 57 +++++++++++++++++++++++++++++
 cylc/uiserver/websockets/tornado.py |  4 +-
 2 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 cylc/uiserver/websockets/resolve.py

diff --git a/cylc/uiserver/websockets/resolve.py b/cylc/uiserver/websockets/resolve.py
new file mode 100644
index 00000000..b0f4c316
--- /dev/null
+++ b/cylc/uiserver/websockets/resolve.py
@@ -0,0 +1,57 @@
+# MIT License
+# 
+# Copyright (c) 2017, Syrus Akbary
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""
+This file contains an implementation of "resolve" derived from the one
+found in the graphql-ws library with the above license.
+"""
+
+
+async def resolve(
+    data: Any, _container: Union[List, Dict] = None, _key: Union[str, int] = None
+) -> None:
+    """
+    Recursively wait on any awaitable children of a data element and resolve any
+    Promises.
+    """
+    stack = [(data, _container, _key)]
+
+    while stack:
+        data, _container, _key = stack.pop()
+
+        if is_awaitable(data):
+            data = await data
+            if isinstance(data, Promise):
+                data = data.value  # type: Any
+            if _container is not None:
+                _container[_key] = data
+        if isinstance(data, dict):
+            items = data.items()
+        elif isinstance(data, list):
+            items = enumerate(data)
+        else:
+            items = None
+        if items is not None:
+            stack.extend([
+                (child, data, key)
+                for key, child in items
+            ])
diff --git a/cylc/uiserver/websockets/tornado.py b/cylc/uiserver/websockets/tornado.py
index fcabac34..1bb1cbb5 100644
--- a/cylc/uiserver/websockets/tornado.py
+++ b/cylc/uiserver/websockets/tornado.py
@@ -15,7 +15,6 @@
 from graphql.execution.middleware import MiddlewareManager
 from graphql_ws.base import ConnectionClosedException
 from graphql_ws.base_async import (
-    resolve,
     BaseAsyncConnectionContext,
     BaseAsyncSubscriptionServer
 )
@@ -29,6 +28,9 @@
 from typing import Union, Awaitable, Any, List, Tuple, Dict, Optional
 
 from cylc.uiserver.authorise import AuthorizationMiddleware
+from cylc.uiserver.websockets.resolve import resolve
+
+
 setup_observable_extension()
 
 NO_MSG_DELAY = 1.0