Support tracking latency from multiple indexers (#14)

* Add server_address field and new table to indexer websocket listener

* Increase wait time for place_stateful_orders

* Update indexer ws subscriber to also listen to stateful order address

* Fix bug where batch writer could get stuck with timeout = 0 and drop data

* Remove excessive sleeping in SIGTERM/SIGINT handler for run_all_scripts.py

Author: matthewdowney

bigquery_to_datadog.py

@@ -32,9 +32,9 @@
                 UNION ALL
                 SELECT p.received_at
                     , TIMESTAMP_DIFF(p.received_at, s.sent_at, millisecond) AS latency
-                    , "indexer" AS server_address
+                    , p.server_address
                 FROM `{project_id}.latency_experiments.long_running_two_sided_orders` s
-                JOIN `{project_id}.indexer_stream.received_orders_and_cancels` p
+                JOIN `{project_id}.indexer_stream_new.received_orders_and_cancels` p
                     ON p.client_id = CAST(s.client_id AS STRING)
                    AND p.address = s.address
                    AND p.received_at > TIMESTAMP("{start_timestamp}")
@@ -68,9 +68,9 @@
             UNION ALL
             SELECT p.received_at
                 , TIMESTAMP_DIFF(p.received_at, s.sent_at, millisecond) AS latency
-                , "indexer" AS server_address
+                , p.server_address
             FROM `{project_id}.latency_experiments.long_running_stateful_orders` s
-            JOIN `{project_id}.indexer_stream.received_orders_and_cancels` p
+            JOIN `{project_id}.indexer_stream_new.received_orders_and_cancels` p
                 ON p.client_id = CAST(s.client_id AS STRING)
             AND p.address = s.address
             AND p.received_at > TIMESTAMP("{start_timestamp}")

bq_helpers.py

@@ -103,7 +103,7 @@ async def worker(self):
                 elapsed_time = (
                         datetime.utcnow() - self.last_flush_time
                 ).total_seconds()
-                dynamic_timeout = max(0, self.batch_timeout - elapsed_time)
+                dynamic_timeout = max(0.1, self.batch_timeout - elapsed_time)
                 data = await asyncio.wait_for(self.queue.get(), timeout=dynamic_timeout)
                 data_buffer.append(data)
 

listen_to_websocket.py

@@ -3,7 +3,7 @@
 
 Usage: python listen_to_websocket.py 
 """
-
+import argparse
 import asyncio
 import json
 import logging
@@ -14,7 +14,6 @@
 import websockets
 from google.cloud import bigquery
 from google.cloud.bigquery import SchemaField
-from v4_client_py.clients.constants import Network
 
 # Import the BigQuery helpers
 from bq_helpers import create_table, BatchWriter
@@ -23,41 +22,45 @@
 with open("config.json", "r") as config_file:
     config_json = json.load(config_file)
 
-DATASET_ID = "indexer_stream"
+DATASET_ID = "indexer_stream_new"
 TABLE_ID = "responses"
 
 SCHEMA = [
     SchemaField("received_at", "TIMESTAMP", mode="REQUIRED"),
     SchemaField("uuid", "STRING", mode="REQUIRED"),
     SchemaField("response", "JSON", mode="NULLABLE"),
+    SchemaField("server_address", "STRING", mode="REQUIRED"),
 ]
 
 TIME_PARTITIONING = bigquery.TimePartitioning(field="received_at")
+CLUSTERING_FIELDS = ["server_address"]
+
 # Batch settings
 BATCH_SIZE = 9
 BATCH_TIMEOUT = 10
 WORKER_COUNT = 1
 
 
-def process_message(message):
+def process_message(message, url):
     return {
         "received_at": datetime.utcnow().isoformat("T") + "Z",
         "uuid": str(uuid.uuid4()),
         "response": message,
+        "server_address": url
     }
 
 
 class AsyncSocketClient:
-    def __init__(self, config, subaccount_ids, batch_writer):
-        self.url = config.websocket_endpoint
+    def __init__(self, indexer_url: str, subaccount_ids, batch_writer):
+        self.url = indexer_url
         self.subaccount_ids = subaccount_ids
         self.batch_writer = batch_writer
 
     async def connect(self):
         retries = 0
         while True:
             try:
-                async with websockets.connect(self.url) as websocket:
+                async with websockets.connect(f"wss://{self.url}/v4/ws") as websocket:
                     if self.subaccount_ids:
                         for subaccount_id in self.subaccount_ids:
                             await self.subscribe(
@@ -81,7 +84,7 @@ async def connect(self):
     async def consumer_handler(self, websocket):
         async for message in websocket:
             await self.batch_writer.enqueue_data(
-                process_message(message)
+                process_message(message, self.url)
             )  # Enqueue data for batch writing
 
     async def send(self, websocket, message):
@@ -97,17 +100,19 @@ async def subscribe(self, websocket, channel, params=None):
         await self.send(websocket, message)
 
 
-async def main():
+async def main(indexer_url: str):
     batch_writer = BatchWriter(
         DATASET_ID, TABLE_ID, WORKER_COUNT, BATCH_SIZE, BATCH_TIMEOUT
     )
-    config = Network.config_network().indexer_config
     subaccount_ids = [
         "/".join([config_json["maker_address"], str(0)]),
         "/".join([config_json["taker_address"], str(0)]),
+        "/".join([config_json["stateful_address"], str(0)]),
     ]
     client = AsyncSocketClient(
-        config, subaccount_ids=subaccount_ids, batch_writer=batch_writer
+        indexer_url,
+        subaccount_ids=subaccount_ids,
+        batch_writer=batch_writer,
     )
 
     batch_writer_task = asyncio.create_task(batch_writer.batch_writer_loop())
@@ -116,8 +121,19 @@ async def main():
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run the Indexer Websocket client for a given URL, e.g. indexer.v4testnet.dydx.exchange."
+    )
+    parser.add_argument(
+        "--indexer_url",
+        type=str,
+        help="The indexer API to read from.",
+    )
+    args = parser.parse_args()
+
+    log_id = args.indexer_url.replace(":", "_").replace("/", "_")
     handler = RotatingFileHandler(
-        "listen_to_websocket.log",
+        f"listen_to_websocket_{log_id}.log",
         maxBytes=5 * 1024 * 1024,  # 5 MB
         backupCount=5
     )
@@ -127,5 +143,5 @@ async def main():
         format="%(asctime)s - %(levelname)s - %(message)s",
     )
 
-    create_table(DATASET_ID, TABLE_ID, SCHEMA, TIME_PARTITIONING)
-    asyncio.run(main())
+    create_table(DATASET_ID, TABLE_ID, SCHEMA, TIME_PARTITIONING, CLUSTERING_FIELDS)
+    asyncio.run(main(args.indexer_url))

place_stateful_orders.py

@@ -56,7 +56,7 @@
 NUM_BLOCKS = 10_000
 DYDX_MNEMONIC = config["stateful_mnemonic"]
 GTBT_DELTA = 5
-PLACE_INTERVAL = 12
+PLACE_INTERVAL = 13
 
 
 async def listen_to_block_stream_and_place_orders(batch_writer):

run_all_scripts.py

@@ -28,14 +28,14 @@
 CHECK_INTERVAL = 250  # Check every 250 seconds
 
 SCRIPT_CONFIGS = {
-    "websocket": {
-        "script_name": "listen_to_websocket.py",
-        "table_id": "indexer_stream.responses",
-        "timestamp_column": "received_at",
-        "filter": "",
-        "args": [],
-        "time_threshold": timedelta(seconds=90),
-    },
+    # "websocket": {
+    #     "script_name": "listen_to_websocket.py",
+    #     "table_id": "indexer_stream.responses",
+    #     "timestamp_column": "received_at",
+    #     "filter": "",
+    #     "args": [],
+    #     "time_threshold": timedelta(seconds=90),
+    # },
     "place_orders": {
         "script_name": "place_orders.py",
         "table_id": "latency_experiments.long_running_two_sided_orders",
@@ -74,6 +74,16 @@
         "time_threshold": timedelta(seconds=90),
     }
 
+for addr in config["indexer_addresses"]:
+    SCRIPT_CONFIGS[f"indexer {addr}"] = {
+        "script_name": "listen_to_websocket.py",
+        "table_id": "indexer_stream_new.responses",
+        "timestamp_column": "received_at",
+        "filter": f'server_address = "{addr}"',
+        "args": ["--indexer_url", addr],
+        "time_threshold": timedelta(seconds=90),
+    }
+
 
 def get_latest_timestamp(table_id, timestamp_column, filter_condition):
     filter_clause = f"WHERE TIMESTAMP_TRUNC({timestamp_column}, DAY) = TIMESTAMP(CURRENT_DATE()) AND {filter_condition}" if filter_condition else ""
@@ -95,14 +105,13 @@ def start_script(script_name, args):
     return subprocess.Popen(["python", script_name] + args)
 
 
-def force_kill_process(process: subprocess.Popen, pname: str):
+def terminate_process(process: subprocess.Popen, pname: str):
     # Try to terminate the process
     logging.info(f"Terminating process {pname}...")
     process.terminate()
 
-    # Wait for a few seconds to give the process time to terminate
-    time.sleep(3)
 
+def force_kill_process(process: subprocess.Popen, pname: str):
     # Check if the process has terminated
     if process.poll() is None:
         # Process is still alive, so forcefully kill it
@@ -147,6 +156,8 @@ def check_and_restart_script(
         should_restart = True
 
     if should_restart:
+        terminate_process(process, script_name)
+        time.sleep(1)  # Wait for the process to terminate
         force_kill_process(process, script_name)
         return start_script(script_name, args)
     else:
@@ -162,8 +173,14 @@ def main():
     # Gracefully handle Ctrl+C
     def signal_handler(sig, frame):
         logging.info("Received termination signal, shutting down...")
+        for pname, p in processes.items():
+            terminate_process(p, pname)
+
+        time.sleep(3)  # Wait for the processes to terminate
+
         for pname, p in processes.items():
             force_kill_process(p, pname)
+
         sys.exit(0)
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)