enhance discovery function and fix bugs (#46)

- Make discovery function more robust to work with different types of responses, and catch timeout and json value exceptions
- add usage of ThreadpoolExecutor to speed up discovery function
- fix issue with avalanche endpoints for discovery tool

Author: erwin-wee

chainbench/main.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import logging
 import os
 import shlex
@@ -10,6 +11,7 @@
 from click import Context, Parameter
 from locust import runners
 
+from chainbench.tools.discovery.rpc import DiscoveryResult
 from chainbench.user.evm import EVMMethods
 from chainbench.util.cli import (
     ContextData,
@@ -388,11 +390,13 @@ def validate_clients(ctx: Context, param: Parameter, value: str) -> list[str]:
 
     if value is not None:
         input_client_list = value.split(",")
-        client_list = RPCDiscovery.get_client_names()
-        for client in input_client_list:
-            if client not in client_list:
+        client_list: list[str] = []
+        for client in RPCDiscovery.get_clients():
+            client_list.extend(client.get_cli_argument_names())
+        for client_name in input_client_list:
+            if client_name not in client_list:
                 raise click.BadParameter(
-                    f"Client {client} is not supported. "
+                    f"Client {client_name} is not supported. "
                     f"Use 'chainbench list clients' to list all available clients."
                 )
     else:
@@ -421,17 +425,15 @@ def discover(endpoint: str | None, clients: list[str]) -> None:
 
     from chainbench.tools.discovery.rpc import RPCDiscovery
 
+    rpc_discovery = RPCDiscovery(endpoint, clients)
     click.echo(f"Please wait, discovering methods available on {endpoint}...")
-    results = RPCDiscovery.discover_methods(endpoint, clients)
-
-    # Print the results
-    for result in results:
-        if result.supported is True:
-            click.echo(f"{result.method} ✔")
-        elif result.supported is False:
-            click.echo(f"{result.method} ✖")
-        else:
-            click.echo(f"{result.method}: {result.error_message}")
+
+    def get_discovery_result(method: str) -> None:
+        result: DiscoveryResult = rpc_discovery.discover_method(method)
+        click.echo(result.to_string())
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+        executor.map(get_discovery_result, rpc_discovery.methods)
 
 
 @cli.group(name="list", help="Lists values of the given type.")

chainbench/tools/discovery/rpc.py

@@ -9,10 +9,6 @@
 __location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
 
 
-class JSONRPCError(Exception):
-    pass
-
-
 @dataclass
 class RPCMethod:
     name: str
@@ -43,11 +39,26 @@ class DiscoveryResult:
     supported: bool | None = None
     error_message: str | None = None
 
+    def to_string(self) -> str:
+        if self.supported is True:
+            return f"{self.method}, ✔"
+        elif self.supported is False:
+            return f"{self.method}, ✖"
+        else:
+            return f"{self.method}, {self.error_message}"
+
 
 class RPCDiscovery:
     METHODS_FILE = Path(os.path.join(__location__, "methods.json"))
     CLIENTS_FILE = Path(os.path.join(__location__, "clients.json"))
 
+    def __init__(self, endpoint: str, clients: list[str]):
+        self.endpoint = endpoint
+        self.clients = clients
+
+        self.methods = self.get_methods_list(clients)
+        self.http = httpx.Client()
+
     @staticmethod
     def _parse_methods() -> list[RPCMethod]:
         methods = []
@@ -69,12 +80,13 @@ def _parse_clients() -> list[RPCClient]:
         return clients
 
     @classmethod
-    def _get_methods_list(cls, clients: list[str]) -> list[str]:
+    def get_methods_list(cls, clients: list[str]) -> list[str]:
         methods = []
         for method in cls._parse_methods():
             for client in clients:
                 if client in method.supported_clients:
-                    methods.append(method.name)
+                    if method.name not in methods:
+                        methods.append(method.name)
         return methods
 
     @classmethod
@@ -86,38 +98,47 @@ def get_clients(cls) -> list[RPCClient]:
         return cls._parse_clients()
 
     @classmethod
-    def discover_methods(cls, endpoint: str, clients: list[str]) -> list[DiscoveryResult]:
-        result: list[DiscoveryResult] = []
-        http = httpx.Client(base_url=endpoint)
-
-        method_list = cls._get_methods_list(clients)
-
-        for method in method_list:
-            data = {
-                "id": 1,
-                "jsonrpc": "2.0",
-                "method": method,
-                "params": [],
+    def check_response(cls, method: str, response: httpx.Response) -> DiscoveryResult:
+        keywords = ["not supported", "unsupported"]
+
+        if response.status_code not in [200, 400]:
+            return DiscoveryResult(method, None, f"HTTP error {response.status_code}")
+        try:
+            response_json = response.json()
+            if "error" in response_json:
+                error_code = response_json["error"]["code"]
+                if error_code in [-32600, -32604]:
+                    for keyword in keywords:
+                        if keyword in response_json["error"]["message"].lower():
+                            return DiscoveryResult(method, False)
+                if error_code == -32601:
+                    return DiscoveryResult(method, False)
+                if error_code == -32602:
+                    return DiscoveryResult(method, True)
+                return DiscoveryResult(
+                    method,
+                    None,
+                    f"Unable to determine. Unknown error {response_json['error']['code']}: "
+                    f"{response_json['error']['message']}",
+                )
+            return DiscoveryResult(method, True)
+        except ValueError:
+            return DiscoveryResult(method, None, f"Value error {response.json()}")
+
+    def discover_method(self, method: str) -> DiscoveryResult:
+        data = {
+            "id": 1,
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": [],
+        }
+        try:
+            headers = {
+                "User-Agent": "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) "
+                "AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148"
             }
-            response = http.post(endpoint, json=data)
-            if response.status_code != 200:
-                print(f"HTTP Error {response.status_code} with method {method}")
-            else:
-                response_json = response.json()
-                if "error" in response_json:
-                    if response_json["error"]["code"] == -32602:
-                        result.append(DiscoveryResult(method, True))
-                    elif response_json["error"]["code"] == -32601:
-                        result.append(DiscoveryResult(method, False))
-                    else:
-                        result.append(
-                            DiscoveryResult(
-                                method,
-                                None,
-                                f"Unable to determine. Unknown error {response_json['error']['code']}: "
-                                f"{response_json['error']['message']}",
-                            )
-                        )
-                else:
-                    result.append(DiscoveryResult(method, True))
-        return result
+            response = self.http.post(self.endpoint, json=data, headers=headers)
+            return self.check_response(method, response)
+
+        except httpx.TimeoutException as e:
+            return DiscoveryResult(method, None, f"HTTP Timeout Exception: {e}")