| # | |
| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| # | |
| # | |
| import asyncio | |
| import json | |
| import logging | |
| from typing import Optional | |
| import httpx | |
| import motor.motor_asyncio | |
| import websockets | |
| from pyee import AsyncIOEventEmitter | |
| from pymongo import UpdateOne | |
| from .config import MongoConfig, ResilientDBConfig | |
| from .exceptions import ResilientPythonCacheError | |
| logging.basicConfig(level=logging.INFO, format='%(message)s') | |
| logger = logging.getLogger(__name__) | |
| logging.getLogger("httpx").setLevel(logging.WARNING) | |
| logging.getLogger("httpcore").setLevel(logging.WARNING) | |
| class ResilientPythonCache(AsyncIOEventEmitter): | |
| def __init__(self, mongo_config: MongoConfig, resilient_db_config: ResilientDBConfig): | |
| super().__init__() | |
| self.mongo_config = mongo_config | |
| self.resilient_db_config = resilient_db_config | |
| self.http_endpoint: str | |
| self.ws_endpoint: str | |
| self.reconnect_interval: int | |
| self.fetch_interval: int | |
| self.current_block_number: int = 0 | |
| self.fetch_task: Optional[asyncio.Task] = None | |
| self.ws_task: Optional[asyncio.Task] = None | |
| self.is_closing: bool = False | |
| self.reconnect_attempts: int = 0 | |
| self.initialize_endpoints() | |
| def initialize_endpoints(self): | |
| config = self.resilient_db_config | |
| self.reconnect_interval = config.reconnect_interval | |
| self.fetch_interval = config.fetch_interval | |
| if not config.base_url.startswith("resilientdb://"): | |
| raise ResilientPythonCacheError("Invalid protocol in base_url. Expected 'resilientdb://'") | |
| url = config.base_url[len("resilientdb://"):] | |
| hostname_port = url.split("/")[0] | |
| hostname, port = (hostname_port.split(":") + [None])[:2] | |
| port = port or "" | |
| http_protocol = "https" if config.http_secure else "http" | |
| ws_protocol = "wss" if config.ws_secure else "ws" | |
| if port: | |
| host_port = f"{hostname}:{port}" | |
| else: | |
| host_port = hostname | |
| self.http_endpoint = config.http_endpoint or f"{http_protocol}://{host_port}/v1/blocks" | |
| self.ws_endpoint = config.ws_endpoint or f"{ws_protocol}://{host_port}/blockupdatelistener" | |
| logger.info(f"HTTP Endpoint: {self.http_endpoint}") | |
| logger.info(f"WebSocket Endpoint: {self.ws_endpoint}") | |
| async def initialize(self): | |
| try: | |
| self.mongo_client = motor.motor_asyncio.AsyncIOMotorClient(self.mongo_config.uri) | |
| self.db = self.mongo_client[self.mongo_config.db_name] | |
| self.collection = self.db[self.mongo_config.collection_name] | |
| await self.mongo_client.admin.command('ping') | |
| logger.info(f"Connected to MongoDB database: {self.mongo_config.db_name}, " | |
| f"collection: {self.mongo_config.collection_name}") | |
| await self.fetch_and_sync_initial_blocks() | |
| self.fetch_task = asyncio.create_task(self.start_periodic_fetch()) | |
| self.ws_task = asyncio.create_task(self.connect_websocket()) | |
| except Exception as e: | |
| logger.error("Error initializing connections:") | |
| logger.error(e) | |
| raise ResilientPythonCacheError(str(e)) from e | |
| async def fetch_and_sync_initial_blocks(self): | |
| try: | |
| last_block = await self.collection.find_one({}, sort=[("id", -1)]) | |
| self.current_block_number = last_block['id'] if last_block and 'id' in last_block else 0 | |
| batch_size = 100 | |
| concurrency_limit = 5 | |
| batch_ranges = [] | |
| blocks_fetched = True | |
| async with httpx.AsyncClient() as client: | |
| while blocks_fetched: | |
| min_seq = self.current_block_number + 1 | |
| max_seq = min_seq + batch_size - 1 | |
| url = f"{self.http_endpoint}/{min_seq}/{max_seq}" | |
| logger.info(f"Checking for blocks from {min_seq} to {max_seq}") | |
| response = await client.get(url) | |
| blocks = response.json() | |
| if isinstance(blocks, list) and blocks: | |
| self.current_block_number = max(block.get('id', 0) for block in blocks) | |
| batch_ranges.append({'min_seq': min_seq, 'max_seq': max_seq}) | |
| else: | |
| blocks_fetched = False | |
| if not batch_ranges: | |
| logger.info("No new blocks to sync.") | |
| return | |
| logger.info(f"Total batches to fetch: {len(batch_ranges)}") | |
| semaphore = asyncio.Semaphore(concurrency_limit) | |
| tasks = [self.fetch_and_sync_batch(batch['min_seq'], batch['max_seq'], semaphore) | |
| for batch in batch_ranges] | |
| await asyncio.gather(*tasks) | |
| except Exception as e: | |
| logger.error("Error fetching initial blocks:") | |
| logger.error(e) | |
| raise ResilientPythonCacheError(str(e)) from e | |
| async def fetch_and_sync_batch(self, min_seq: int, max_seq: int, semaphore: asyncio.Semaphore): | |
| async with semaphore: | |
| try: | |
| async with httpx.AsyncClient() as client: | |
| url = f"{self.http_endpoint}/{min_seq}/{max_seq}" | |
| logger.info(f"Fetching blocks from {min_seq} to {max_seq}") | |
| response = await client.get(url) | |
| blocks = response.json() | |
| if isinstance(blocks, list) and blocks: | |
| processed_blocks = self.process_blocks(blocks) | |
| bulk_ops = [ | |
| UpdateOne( | |
| {'id': block['id']}, | |
| {'$set': block}, | |
| upsert=True | |
| ) for block in processed_blocks | |
| ] | |
| if bulk_ops: | |
| result = await self.collection.bulk_write(bulk_ops) | |
| logger.info(f"Blocks {min_seq} to {max_seq} synced: " | |
| f"Inserted {result.upserted_count}, Modified {result.modified_count}") | |
| self.emit('data', processed_blocks) # Emit 'data' event with new blocks | |
| else: | |
| logger.info(f"No blocks fetched for range {min_seq} to {max_seq}") | |
| except Exception as e: | |
| logger.error(f"Error fetching blocks from {min_seq} to {max_seq}:") | |
| logger.error(e) | |
| def process_blocks(self, blocks: list) -> list: | |
| for block in blocks: | |
| transactions = block.get('transactions', []) | |
| if isinstance(transactions, list): | |
| for transaction in transactions: | |
| if 'value' in transaction and isinstance(transaction['value'], str): | |
| value = transaction['value'].strip() | |
| if (value.startswith('{') and value.endswith('}')) or \ | |
| (value.startswith('[') and value.endswith(']')): | |
| try: | |
| transaction['value'] = json.loads(value) | |
| except json.JSONDecodeError: | |
| logger.warning(f"Could not parse transaction value for key {transaction.get('key')}. " | |
| f"Keeping original value.") | |
| return blocks | |
| async def start_periodic_fetch(self): | |
| logger.info(f"Started periodic fetch every {self.fetch_interval / 1000} seconds.") | |
| try: | |
| while not self.is_closing: | |
| await self.fetch_and_sync_new_blocks() | |
| await asyncio.sleep(self.fetch_interval / 1000) | |
| except asyncio.CancelledError: | |
| pass # Allow task to be cancelled gracefully | |
| async def fetch_and_sync_new_blocks(self): | |
| try: | |
| batch_size = 100 | |
| concurrency_limit = 5 | |
| batch_ranges = [] | |
| blocks_fetched = True | |
| async with httpx.AsyncClient() as client: | |
| while blocks_fetched: | |
| min_seq = self.current_block_number + 1 | |
| max_seq = min_seq + batch_size - 1 | |
| url = f"{self.http_endpoint}/{min_seq}/{max_seq}" | |
| logger.info(f"Checking for new blocks from {min_seq} to {max_seq}") | |
| response = await client.get(url) | |
| blocks = response.json() | |
| if isinstance(blocks, list) and blocks: | |
| self.current_block_number = max(block.get('id', 0) for block in blocks) | |
| batch_ranges.append({'min_seq': min_seq, 'max_seq': max_seq}) | |
| else: | |
| blocks_fetched = False | |
| if not batch_ranges: | |
| logger.info("No new blocks to sync.") | |
| return | |
| logger.info(f"Total new batches to fetch: {len(batch_ranges)}") | |
| semaphore = asyncio.Semaphore(concurrency_limit) | |
| tasks = [self.fetch_and_sync_batch(batch['min_seq'], batch['max_seq'], semaphore) | |
| for batch in batch_ranges] | |
| await asyncio.gather(*tasks) | |
| except Exception as e: | |
| logger.error("Error fetching new blocks:") | |
| logger.error(e) | |
| self.emit('error', e) | |
| async def connect_websocket(self): | |
| try: | |
| async with websockets.connect(self.ws_endpoint) as websocket: | |
| logger.info(f"Connected to WebSocket: {self.ws_endpoint}") | |
| self.reconnect_attempts = 0 | |
| self.emit('connected') | |
| # Removed the duplicate log statement here | |
| # logger.info("WebSocket connected.") | |
| async for message in websocket: | |
| logger.info(f"Received message: {message}") | |
| try: | |
| parsed_message = json.loads(message) | |
| except json.JSONDecodeError: | |
| parsed_message = {"type": message} | |
| if parsed_message.get("type") == "Update blocks": | |
| await self.fetch_and_sync_new_blocks() | |
| else: | |
| logger.warning(f"Received unrecognized message: {message}") | |
| except (websockets.exceptions.ConnectionClosedError, | |
| websockets.exceptions.InvalidURI, | |
| websockets.exceptions.InvalidHandshake) as e: | |
| logger.warning(f"WebSocket closed or failed to connect: {e}") | |
| if not self.is_closing: | |
| await self.attempt_reconnection() | |
| except asyncio.CancelledError: | |
| pass # Allow task to be cancelled gracefully | |
| except Exception as e: | |
| logger.error("WebSocket error:") | |
| logger.error(e) | |
| if not self.is_closing: | |
| await self.attempt_reconnection() | |
| async def attempt_reconnection(self): | |
| delay = self.reconnect_interval * (2 ** self.reconnect_attempts) / 1000 # Convert ms to seconds | |
| logger.info(f"Attempting to reconnect in {delay} seconds...") | |
| await asyncio.sleep(delay) | |
| self.reconnect_attempts += 1 | |
| logger.info(f"Reconnection attempt #{self.reconnect_attempts}") | |
| self.ws_task = asyncio.create_task(self.connect_websocket()) | |
| async def close(self): | |
| try: | |
| self.is_closing = True | |
| if self.fetch_task: | |
| self.fetch_task.cancel() | |
| try: | |
| await self.fetch_task | |
| except asyncio.CancelledError: | |
| pass | |
| if self.ws_task: | |
| self.ws_task.cancel() | |
| try: | |
| await self.ws_task | |
| except asyncio.CancelledError: | |
| pass | |
| await self.mongo_client.close() | |
| logger.info("Closed MongoDB and WebSocket connections.") | |
| self.emit('closed') | |
| except Exception as e: | |
| logger.error("Error closing connections:") | |
| logger.error(e) | |
| raise ResilientPythonCacheError(str(e)) from e |