Maintenance and bugfixes (#13)

* Fix bug that double ran place_stateful_orders.py each restart

* Try terminate and kill in run_all_scripts.py + add more logging

* Fix asyncio event-loop blocking in place_stateful_orders.py

* Add more detailed logging for BigQuery insert errors

* Update clob pair id + add TODOs for code cleanup

* Refactor precompute_order to not need to query the chain

* Update place_orders.py to only use asyncio-friendly code

* Handle place_stateful_orders.py account sequence bug

* Bump clob pair id to 1000 to cover future market listings

* Update batch writer to fall back to GCS insert automatically

Author: matthewdowney

bq_helpers.py

@@ -1,8 +1,9 @@
 import asyncio
 import logging
 from datetime import datetime
-from typing import List, Callable
+from typing import List, Callable, Type, Awaitable
 
+from google.api_core.exceptions import GoogleAPICallError
 from google.cloud import bigquery, storage
 from google.cloud.bigquery import SchemaField
 from google.cloud.exceptions import NotFound
@@ -57,21 +58,40 @@ def __init__(
         self.batch_timeout = batch_timeout
         self.last_flush_time = datetime.utcnow()
 
+        async def default_error_413_handler(data_buffer, e):
+            logging.error(f"Error inserting {len(data_buffer)} rows (code = {e.code}): {e}")
+
+        self.error_413_handler = default_error_413_handler
+
     async def enqueue_data(self, data):
         await self.queue.put(data)
 
     async def flush_data(self, data_buffer):
         try:
             errors = await asyncio.to_thread(
-                self.bq_client.insert_rows_json, self.table_ref, data_buffer
+                self.bq_client.insert_rows_json,
+                self.table_ref,
+                data_buffer
             )
             if errors:
                 logging.error(f"Errors occurred: {errors}")
+        except GoogleAPICallError as e:
+            if e.code == 413:
+                await self.error_413_handler(data_buffer, e)
+            else:
+                logging.error(f"Error inserting {len(data_buffer)} rows (code = {e.code}): {e}")
         except Exception as e:
-            logging.error(f"Error inserting rows: {e}")
+            logging.error(f"Error inserting {len(data_buffer)} rows: {e}")
         finally:
             self.last_flush_time = datetime.utcnow()
 
+    def set_error_413_handler(self, handler: Callable[[List, Type[Exception]], Awaitable[None]]):
+        """
+        Set an async function handler that will be called when a 413 (payload
+        too large) error occurs.
+        """
+        self.error_413_handler = handler
+
     async def batch_writer_loop(self):
         workers = [asyncio.create_task(self.worker()) for _ in range(self.worker_count)]
         await asyncio.gather(*workers)
@@ -140,9 +160,9 @@ async def gcs_writer_loop(self):
                 logging.error(f"Error inserting data: {e}")
 
     def _process_and_insert(self, data: dict) -> None:
-        data = self.middleware_fn(data)
+        data = [self.middleware_fn(data) for data in data]
         return gcs_insert.insert_via_gcs(
-            self.client, self.bucket, self.table, [data]
+            self.client, self.bucket, self.table, data
         )
 
     def set_middleware(self, middleware_fn: Callable[[dict], dict]):

client_helpers.py

@@ -43,6 +43,16 @@ def setup_clients(grpc_endpoint):
     return client, ledger_client
 
 
+def get_current_block_with_retries(client: CompositeClient):
+    for i in range(5):
+        try:
+            return client.get_current_block()
+        except Exception as error:
+            if i == 4:
+                raise error
+    raise Exception("Failed to get current block")
+
+
 def get_markets_data(client, market):
     markets_response = client.indexer_client.markets.get_perpetual_markets(market)
     return markets_response.data["markets"][market]
@@ -62,19 +72,20 @@ def place_order(client, block, msg_and_order, batch_writer):
 
 
 def precompute_order(
-    client,
-    ledger_client,
-    market,
-    subaccount,
-    side,
-    price,
-    client_id,
-    good_til_block,
-    good_til_block_time,
-    size,
-    order_flags,
-    time_in_force,
-    sequence_number,
+        client,
+        ledger_client,
+        market,
+        subaccount,
+        side,
+        price,
+        client_id,
+        good_til_block,
+        good_til_block_time,
+        size,
+        order_flags,
+        time_in_force,
+        sequence_number,
+        account_number,
 ):
     # precompute order and sign the order
     clob_pair_id = market["clobPairId"]
@@ -114,17 +125,20 @@ def precompute_order(
     memo = None
 
     fee = ledger_client.estimate_fee_from_gas(gas_limit)
-    account = ledger_client.query_account(sender.address())
 
     tx.seal(
         SigningCfg.direct(sender.public_key(), sequence_number),
         fee=fee,
         gas_limit=gas_limit,
         memo=memo,
     )
-    tx.sign(sender.signer(), ledger_client.network_config.chain_id, account.number)
+    tx.sign(
+        sender.signer(),
+        ledger_client.network_config.chain_id,
+        account_number
+    )
     tx.complete()
-    return (tx, msg)
+    return tx, msg
 
 
 async def place_orders(client, block, msg_and_orders, batch_writer):

listen_to_grpc_stream.py

@@ -22,7 +22,7 @@
 # Dataset configuration
 DATASET_ID = "full_node_stream"
 TABLE_ID = "responses"
-CLOB_PAIR_IDS = range(133)
+CLOB_PAIR_IDS = range(1000)  # Have to update if too many more pairs are added
 
 # If data to too large for direct insert, use this GCS bucket to sideload
 GCS_BUCKET = "full_node_stream_sideload"
@@ -97,6 +97,9 @@ async def listen_to_stream_and_write_to_bq(
     retry_count = 0
     start_time = datetime.utcnow()
 
+    msg_count = 0
+    msg_count_time = datetime.utcnow()
+
     while retry_count < MAX_RETRIES_PER_DAY:
         try:
             stub = QueryStub(channel)
@@ -107,12 +110,19 @@ async def listen_to_stream_and_write_to_bq(
                 row = process_message(response, server_address)
 
                 # If the row is too large, sideload into BQ via GCS
-                too_large_for_direct_insert = len(row['response']) > 5_000_000
+                too_large_for_direct_insert = len(row['response']) > 1_000_000
                 if too_large_for_direct_insert:
-                    await gcs_writer.enqueue_data(row)
+                    await gcs_writer.enqueue_data([row])
                 else:
                     await batch_writer.enqueue_data(row)
 
+                # Log message counts every 15 mins
+                msg_count += 1
+                if datetime.utcnow() - msg_count_time > timedelta(minutes=15):
+                    logging.info(f"Saw {msg_count} msgs in the last 15 minutes")
+                    msg_count_time = datetime.utcnow()
+                    msg_count = 0
+
             await batch_writer.enqueue_data(
                 process_error("Stream ended", server_address)
             )
@@ -166,24 +176,32 @@ def preprocess_row_for_gcs(row: dict) -> dict:
 
 
 async def main(server_address):
+    # Writer for exceptionally large rows (book snapshots)
+    gcs_writer = GCSWriter(DATASET_ID, TABLE_ID, SCHEMA, GCS_BUCKET)
+    gcs_writer.set_middleware(preprocess_row_for_gcs)
+    gcs_writer_task = asyncio.create_task(gcs_writer.gcs_writer_loop())
+
     # Writer for direct BigQuery inserts
     batch_writer = BatchWriter(
         DATASET_ID, TABLE_ID, WORKER_COUNT, BATCH_SIZE, BATCH_TIMEOUT
     )
     batch_writer_task = asyncio.create_task(batch_writer.batch_writer_loop())
 
-    # Writer for exceptionally large rows (book snapshots)
-    gcs_writer = GCSWriter(DATASET_ID, TABLE_ID, SCHEMA, GCS_BUCKET)
-    gcs_writer.set_middleware(preprocess_row_for_gcs)
-    gcs_writer_task = asyncio.create_task(gcs_writer.gcs_writer_loop())
+    # If the direct writer encounters a 413 error, fall back to sideloading
+    async def error_413_handler(data_buffer, e):
+        logging.error(f"Error 413 occurred, sideloading {len(data_buffer)} rows.")
+        await gcs_writer.enqueue_data(data_buffer)
+    batch_writer.set_error_413_handler(error_413_handler)
 
     # Adjust to use secure channel if needed
+    logging.info(f"Connecting to server at {server_address}...")
     async with grpc.aio.insecure_channel(server_address, GRPC_OPTIONS) as channel:
         await listen_to_stream_and_write_to_bq(
             channel, batch_writer, gcs_writer, server_address
         )
 
     await asyncio.gather(batch_writer_task, gcs_writer_task)
+    logging.error("Listening task finished (unexpected)")
 
 
 if __name__ == "__main__":