feat: Support multiple indexer-service endpoints (#49)

* feat: combine query count from multiple service endpoints Signed-off-by: Marc-André Dumas <[email protected]> * fix: remove unused code Signed-off-by: Marc-André Dumas <[email protected]> * fix: move pcrpy dependency to dev Signed-off-by: Marc-André Dumas <[email protected]> * fix: lint Signed-off-by: Marc-André Dumas <[email protected]> * fix: moved argparse Signed-off-by: Marc-André Dumas <[email protected]> * fix: add mandatory arguments Signed-off-by: Marc-André Dumas <[email protected]> * docs: args for multiple indexer-service endpoints Signed-off-by: Marc-André Dumas <[email protected]> Signed-off-by: Marc-André Dumas <[email protected]>
semiotic-ai · Jan 24, 2023 · b1db19c · b1db19c
1 parent 24324e7
commit b1db19c
Show file tree

Hide file tree

Showing 8 changed files with 12,880 additions and 924 deletions.
diff --git a/README.md b/README.md
@@ -66,8 +66,9 @@ optional arguments:
                         URL to the indexer-agent management GraphQL endpoint. [env var:
                         INDEXER_AGENT_MGMT_ENDPOINT] (default: None)
   --indexer-service-metrics-endpoint INDEXER_SERVICE_METRICS_ENDPOINT
-                        HTTP endpoint for the indexer-service metrics. [env var:
-                        INDEXER_SERVICE_METRICS_ENDPOINT] (default: None)
+                        HTTP endpoint for the indexer-service metrics. Can be a comma-separated 
+                        for multiple endpoints. [env var: INDEXER_SERVICE_METRICS_ENDPOINT]
+                        (default: None)
   --qps-observation-duration QPS_OBSERVATION_DURATION
                         Duration of the measurement period of the query-per-second after a price
                         multiplier update. [env var: QPS_OBSERVATION_DURATION] (default: 60)

diff --git a/autoagora/config.py b/autoagora/config.py
@@ -110,7 +110,7 @@ def init_config(argv: Optional[Sequence[str]] = None):
         "--indexer-service-metrics-endpoint",
         env_var="INDEXER_SERVICE_METRICS_ENDPOINT",
         required=True,
-        help="HTTP endpoint for the indexer-service metrics.",
+        help="HTTP endpoint for the indexer-service metrics. Can be a comma-separated for multiple endpoints.",
     )
 
     #

diff --git a/autoagora/query_metrics.py b/autoagora/query_metrics.py
@@ -4,7 +4,7 @@
 import asyncio
 import logging
 import re
-from typing import Dict
+from functools import reduce
 
 import aiohttp
 import backoff
@@ -16,51 +16,26 @@ class HTTPError(Exception):
     """Catch-all for HTTP errors"""
 
 
-@backoff.on_exception(
-    backoff.expo, aiohttp.ClientError, max_time=30, logger=logging.root
-)
-async def indexer_service_metrics() -> str:
-    async with aiohttp.ClientSession() as session:
-        async with session.get(args.indexer_service_metrics_endpoint) as response:
-            assert response.status == 200, f"{response.status=}"
-            result = await response.text()
-    return result
-
-
-@backoff.on_exception(
-    backoff.expo, aiohttp.ClientError, max_time=30, logger=logging.root
-)
-async def query_counts() -> Dict[str, int]:
-    async with aiohttp.ClientSession() as session:
-        async with session.get(args.indexer_service_metrics_endpoint) as response:
-            assert response.status == 200
-
-            results = re.findall(
-                r'indexer_service_queries_ok{deployment="(Qm[a-zA-Z0-9]{44})"} ([0-9]*)',
-                await indexer_service_metrics(),
-            )
-
-    results = {subgraph: int(value) for subgraph, value in results}
-    logging.debug("Query counts: %s", results)
-
-    return results
-
-
 @backoff.on_exception(
     backoff.expo, (aiohttp.ClientError, HTTPError), max_time=30, logger=logging.root
 )
 async def subgraph_query_count(subgraph: str) -> int:
+    endpoints = args.indexer_service_metrics_endpoint.split(",")
+    results = []
     async with aiohttp.ClientSession() as session:
-        async with session.get(args.indexer_service_metrics_endpoint) as response:
-            if response.status != 200:
-                raise HTTPError(response.status)
-
-            results = re.findall(
-                r'indexer_service_queries_ok{{deployment="{subgraph}"}} ([0-9]*)'.format(
-                    subgraph=subgraph
-                ),
-                await indexer_service_metrics(),
-            )
+        for endpoint in endpoints:
+            async with session.get(endpoint) as response:
+                if response.status != 200:
+                    raise HTTPError(response.status)
+
+                results.extend(
+                    re.findall(
+                        r'indexer_service_queries_ok{{deployment="{subgraph}"}} ([0-9]*)'.format(
+                            subgraph=subgraph
+                        ),
+                        await response.text(),
+                    )
+                )
 
     if len(results) == 0:
         # The subgraph query count will not be in the metric if it hasn't received any
@@ -69,7 +44,7 @@ async def subgraph_query_count(subgraph: str) -> int:
     if len(results) == 1:
         return int(results[0])
     else:
-        raise RuntimeError(f"More than one matching metric entry for {subgraph}.")
+        return reduce(lambda x, y: int(x) + int(y), results)
 
 
 if __name__ == "__main__":

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -47,6 +47,7 @@ sphinxcontrib-napoleon = "^0.7"
 myst-parser = "^0.17.2"
 pylint = "^2.14.5"
 pytest-cov = "^4.0.0"
+vcrpy = "^4.2.1"
 
 [tool.semantic_release]
 upload_to_repository = false

diff --git a/tests/test_query_metrics.py b/tests/test_query_metrics.py
@@ -0,0 +1,56 @@
+import asyncio
+
+import vcr
+
+import autoagora.query_metrics
+from autoagora.config import init_config
+
+
+class TestQueryMetrics:
+    def test_subgraph_query_count(self):
+        init_config(
+            [
+                "--indexer-agent-mgmt-endpoint",
+                "http://nowhere",
+                "--postgres-host",
+                "nowhere",
+                "--postgres-username",
+                "nowhere",
+                "--postgres-password",
+                "nowhere",
+                "--indexer-service-metrics-endpoint",
+                "http://indexer-service.default.svc.cluster.local:7300/metrics",
+            ]
+        )
+        with vcr.use_cassette("vcr_cassettes/test_subgraph_query_count.yaml"):
+            res = asyncio.run(
+                autoagora.query_metrics.subgraph_query_count(
+                    "Qmadj8x9km1YEyKmRnJ6EkC2zpJZFCfTyTZpuqC3j6e1QH"
+                )
+            )
+        assert res == 938
+
+    def test_subgraph_query_count_multiple_endpoints(self):
+        init_config(
+            [
+                "--indexer-agent-mgmt-endpoint",
+                "http://nowhere",
+                "--postgres-host",
+                "nowhere",
+                "--postgres-username",
+                "nowhere",
+                "--postgres-password",
+                "nowhere",
+                "--indexer-service-metrics-endpoint",
+                "http://indexer-service-0:7300/metrics,http://indexer-service-1:7300/metrics",
+            ]
+        )
+        with vcr.use_cassette(
+            "vcr_cassettes/test_subgraph_query_count_multiple_endpoints.yaml"
+        ):
+            res = asyncio.run(
+                autoagora.query_metrics.subgraph_query_count(
+                    "Qmadj8x9km1YEyKmRnJ6EkC2zpJZFCfTyTZpuqC3j6e1QH"
+                )
+            )
+        assert res == 2607