walrasquant-lib 0.4.20__py3-none-any.whl

walrasquant/core/cache.py

@@ -0,0 +1,688 @@
+import msgspec
+import asyncio
+import threading
+import re
+import redis
+import nexuslog as logging
+
+from typing import Dict, Set, List, Optional, Any, cast
+from collections import defaultdict
+from returns.maybe import maybe
+from pathlib import Path
+
+from walrasquant.schema import (
+    Order,
+    Position,
+    ExchangeType,
+    Kline,
+    BookL1,
+    Trade,
+    AccountBalance,
+    Balance,
+    FundingRate,
+    IndexPrice,
+    MarkPrice,
+    BookL2,
+)
+from walrasquant.constants import STATUS_TRANSITIONS, AccountType, KlineInterval
+from walrasquant.core.entity import TaskManager, get_redis_client_if_available
+from walrasquant.core.nautilius_core import LiveClock, MessageBus
+from walrasquant.constants import StorageType, ParamBackend
+from walrasquant.backends import SQLiteBackend, PostgreSQLBackend, MemoryBackend
+
+
+class AsyncCache:
+    _backend: SQLiteBackend | PostgreSQLBackend | MemoryBackend
+
+    def __init__(
+        self,
+        strategy_id: str,
+        user_id: str,
+        msgbus: MessageBus,
+        clock: LiveClock,
+        task_manager: TaskManager,
+        storage_backend: StorageType = StorageType.SQLITE,
+        db_path: str = ".keys/cache.db",
+        sync_interval: int = 60,  # seconds
+        expired_time: int = 3600,  # seconds
+    ):
+        parent_dir = Path(db_path).parent
+        if not parent_dir.exists():
+            parent_dir.mkdir(parents=True, exist_ok=True)
+
+        self.strategy_id = strategy_id
+        self.user_id = user_id
+        self._storage_backend = storage_backend
+        self._db_path = db_path
+
+        self._log = logging.getLogger(name=type(self).__name__)
+        self._clock = clock
+
+        # in-memory save
+        self._mem_orders: Dict[str, Order] = {}  # oid -> Order
+        self._mem_open_orders: Dict[ExchangeType, Set[str]] = defaultdict(
+            set
+        )  # exchange_id -> set(oid)
+        self._mem_symbol_open_orders: Dict[str, Set[str]] = defaultdict(
+            set
+        )  # symbol -> set(oid)
+        self._mem_symbol_orders: Dict[str, Set[str]] = defaultdict(
+            set
+        )  # symbol -> set(oid)
+        self._mem_positions: Dict[str, Position] = {}  # symbol -> Position
+        self._mem_account_balance: Dict[AccountType, AccountBalance] = defaultdict(
+            AccountBalance
+        )
+        self._mem_params: Dict[str, Any] = {}  # params cache
+        self._cancel_intent_oids: Set[str] = (
+            set()
+        )  # oids currently pending cancel intent
+        self._inflight_orders: Dict[str, Set[str]] = defaultdict(
+            set
+        )  # symbol -> set(oid) orders submitted but not yet acknowledged by exchange
+
+        # set params
+        self._sync_interval = sync_interval  # sync interval
+        self._expired_time = expired_time  # expire time
+        self._task_manager = task_manager
+
+        self._kline_cache: Dict[str, Kline] = {}
+        self._bookl1_cache: Dict[str, BookL1] = {}
+        self._trade_cache: Dict[str, Trade] = {}
+        self._bookl2_cache: Dict[str, BookL2] = {}
+        self._funding_rate_cache: Dict[str, FundingRate] = {}
+        self._index_price_cache: Dict[str, IndexPrice] = {}
+        self._mark_price_cache: Dict[str, MarkPrice] = {}
+
+        self._msgbus = msgbus
+        self._msgbus.subscribe(topic="kline", handler=self._update_kline_cache)
+        self._msgbus.subscribe(topic="bookl1", handler=self._update_bookl1_cache)
+        self._msgbus.subscribe(topic="trade", handler=self._update_trade_cache)
+        self._msgbus.subscribe(topic="bookl2", handler=self._update_bookl2_cache)
+        self._msgbus.subscribe(
+            topic="funding_rate", handler=self._update_funding_rate_cache
+        )
+        self._msgbus.subscribe(
+            topic="index_price", handler=self._update_index_price_cache
+        )
+        self._msgbus.subscribe(
+            topic="mark_price", handler=self._update_mark_price_cache
+        )
+
+        self._storage_initialized = False
+        self._table_prefix = self.safe_table_name(f"{self.strategy_id}_{self.user_id}")
+
+        self._position_lock = threading.RLock()  # Lock for position updates
+        self._order_lock = threading.RLock()  # Lock for order updates
+        self._balance_lock = threading.RLock()  # Lock for balance updates
+
+        # Redis client for parameter storage (lazy initialization)
+        self._redis_client = None
+        self._redis_client_initialized = False
+
+    ################# # base functions ####################
+
+    @staticmethod
+    def safe_table_name(name: str) -> str:
+        name = re.sub(r"[^a-zA-Z0-9_]", "_", name)
+        return name.lower()
+
+    async def _init_storage(self):
+        """Initialize the storage backend"""
+        if self._storage_backend == StorageType.SQLITE:
+            self._backend = SQLiteBackend(
+                strategy_id=self.strategy_id,
+                user_id=self.user_id,
+                table_prefix=self._table_prefix,
+                log=self._log,
+                db_path=self._db_path,
+            )
+        elif self._storage_backend == StorageType.POSTGRESQL:
+            self._backend = PostgreSQLBackend(
+                strategy_id=self.strategy_id,
+                user_id=self.user_id,
+                table_prefix=self._table_prefix,
+                log=self._log,
+            )
+        elif self._storage_backend == StorageType.MEMORY:
+            self._backend = MemoryBackend()
+
+        assert self._backend is not None
+
+        await self._backend.start()
+        self._storage_initialized = True
+
+    async def _load_params_from_db(self):
+        """Load existing parameters from database"""
+        try:
+            existing_params = self._backend.get_all_params()
+            self._mem_params.update(existing_params)
+            if existing_params:
+                self._log.debug(
+                    f"Loaded {len(existing_params)} parameters from database"
+                )
+        except Exception as e:
+            self._log.error(f"Error loading parameters from database: {e}")
+
+    async def start(self):
+        """Start the cache"""
+        await self._init_storage()
+        if self._storage_backend != StorageType.MEMORY:
+            # Load existing parameters from database
+            await self._load_params_from_db()
+            self._task_manager.create_task(self._periodic_sync())
+
+    async def _periodic_sync(self):
+        """Periodically sync the cache"""
+        while True:
+            await self._backend.sync_orders(self._mem_orders)
+            await self._backend.sync_positions(self._mem_positions)
+            await self._backend.sync_open_orders(
+                self._mem_open_orders, self._mem_orders
+            )
+            await self._backend.sync_balances(self._mem_account_balance)
+            await self._backend.sync_params(self._mem_params)
+            self._cleanup_expired_data()
+            await asyncio.sleep(self._sync_interval)
+
+    async def sync_orders(self):
+        with self._order_lock:
+            await self._backend.sync_orders(self._mem_orders)
+
+    async def sync_positions(self):
+        with self._position_lock:
+            await self._backend.sync_positions(self._mem_positions)
+
+    async def sync_open_orders(self):
+        with self._order_lock:
+            await self._backend.sync_open_orders(
+                self._mem_open_orders, self._mem_orders
+            )
+
+    async def sync_balances(self):
+        with self._balance_lock:
+            await self._backend.sync_balances(self._mem_account_balance)
+
+    async def sync_params(self):
+        await self._backend.sync_params(self._mem_params)
+
+    def _cleanup_expired_data(self):
+        """Cleanup expired data"""
+        current_time = self._clock.timestamp_ms()
+        expire_before = current_time - self._expired_time * 1000
+
+        with self._order_lock:
+            expired_orders = []
+            for oid, order in self._mem_orders.copy().items():
+                if order.timestamp is not None and order.timestamp < expire_before:
+                    expired_orders.append(oid)
+
+                    if not order.is_closed:
+                        self._log.warning(f"order {oid} is not closed, but expired")
+
+            for oid in expired_orders:
+                del self._mem_orders[oid]
+                self._log.debug(f"removing order {oid} from memory")
+                for symbol, order_set in self._mem_symbol_orders.copy().items():
+                    self._log.debug(f"removing order {oid} from symbol {symbol}")
+                    order_set.discard(oid)
+
+    async def close(self):
+        """关闭缓存"""
+        if self._storage_initialized and self._backend:
+            try:
+                await self._backend.sync_orders(self._mem_orders)
+                await self._backend.sync_positions(self._mem_positions)
+                await self._backend.sync_open_orders(
+                    self._mem_open_orders, self._mem_orders
+                )
+                await self._backend.sync_balances(self._mem_account_balance)
+                await self._backend.sync_params(self._mem_params)
+            except Exception as e:
+                # Never let cache close crash shutdown; log and continue
+                self._log.error(f"Error closing cache (sync phase): {e}")
+            finally:
+                try:
+                    await self._backend.close()
+                except Exception as e:
+                    self._log.error(f"Error closing storage backend: {e}")
+
+    ################ # cache public data  ###################
+
+    def _update_kline_cache(self, kline: Kline):
+        key = f"{kline.symbol}-{kline.interval.value}"
+        self._kline_cache[key] = kline
+
+    def _update_bookl1_cache(self, bookl1: BookL1):
+        self._bookl1_cache[bookl1.symbol] = bookl1
+
+    def _update_trade_cache(self, trade: Trade):
+        self._trade_cache[trade.symbol] = trade
+
+    def _update_bookl2_cache(self, bookl2: BookL2):
+        self._bookl2_cache[bookl2.symbol] = bookl2
+
+    def _update_funding_rate_cache(self, funding_rate: FundingRate):
+        self._funding_rate_cache[funding_rate.symbol] = funding_rate
+
+    def _update_index_price_cache(self, index_price: IndexPrice):
+        self._index_price_cache[index_price.symbol] = index_price
+
+    def _update_mark_price_cache(self, mark_price: MarkPrice):
+        self._mark_price_cache[mark_price.symbol] = mark_price
+
+    def kline(self, symbol: str, interval: KlineInterval) -> Optional[Kline]:
+        """
+        Retrieve a Kline object from the cache by symbol.
+
+        :param symbol: The symbol of the Kline to retrieve.
+        :return: The Kline object if found, otherwise None.
+        """
+        key = f"{symbol}-{interval.value}"
+        return self._kline_cache.get(key, None)
+
+    def bookl1(self, symbol: str) -> Optional[BookL1]:
+        """
+        Retrieve a BookL1 object from the cache by symbol.
+
+        :param symbol: The symbol of the BookL1 to retrieve.
+        :return: The BookL1 object if found, otherwise None.
+        """
+        return self._bookl1_cache.get(symbol, None)
+
+    def bookl2(self, symbol: str) -> Optional[BookL2]:
+        """
+        Retrieve a BookL2 object from the cache by symbol.
+        """
+        return self._bookl2_cache.get(symbol, None)
+
+    def trade(self, symbol: str) -> Optional[Trade]:
+        """
+        Retrieve a Trade object from the cache by symbol.
+
+        :param symbol: The symbol of the Trade to retrieve.
+        :return: The Trade object if found, otherwise None.
+        """
+        return self._trade_cache.get(symbol, None)
+
+    def funding_rate(self, symbol: str) -> Optional[FundingRate]:
+        """
+        Retrieve a FundingRate object from the cache by symbol.
+        """
+        return self._funding_rate_cache.get(symbol, None)
+
+    def index_price(self, symbol: str) -> Optional[IndexPrice]:
+        """
+        Retrieve an IndexPrice object from the cache by symbol.
+        """
+        return self._index_price_cache.get(symbol, None)
+
+    def mark_price(self, symbol: str) -> Optional[MarkPrice]:
+        """
+        Retrieve a MarkPrice object from the cache by symbol.
+        """
+        return self._mark_price_cache.get(symbol, None)
+
+    ################ # cache private data  ###################
+
+    def _check_status_transition(self, order: Order):
+        if order.oid is None:
+            return True
+        previous_order = self._mem_orders.get(order.oid)
+        if not previous_order:
+            return True
+
+        if order.status not in STATUS_TRANSITIONS[previous_order.status]:
+            self._log.warning(
+                f"Order id: {order.oid} Invalid status transition: {previous_order.status} -> {order.status}"
+            )
+            return False
+
+        return True
+
+    def _apply_position(self, position: Position):
+        with self._position_lock:
+            if position.is_closed:
+                self._mem_positions.pop(position.symbol, None)
+            else:
+                self._mem_positions[position.symbol] = position
+
+    def _apply_balance(self, account_type: AccountType, balances: List[Balance]):
+        with self._balance_lock:
+            self._mem_account_balance[account_type]._apply(balances)
+
+    def get_balance(self, account_type: AccountType) -> AccountBalance:
+        with self._balance_lock:
+            return self._mem_account_balance[account_type]
+
+    @maybe
+    def get_position(self, symbol: str) -> Optional[Position]:
+        with self._position_lock:
+            if position := self._mem_positions.get(symbol, None):
+                return position
+
+    def get_all_positions(
+        self, exchange: Optional[ExchangeType] = None
+    ) -> Dict[str, Position]:
+        with self._position_lock:
+            positions = {
+                symbol: position
+                for symbol, position in self._mem_positions.copy().items()
+                if (
+                    (exchange is None or position.exchange == exchange)
+                    and position.is_opened
+                )
+            }
+            return positions
+
+    def _order_status_update(self, order: Order) -> bool:
+        if order.oid is None:
+            return False
+        with self._order_lock:
+            if not self._check_status_transition(order):
+                return False
+            self._mem_orders[order.oid] = order
+
+            # Remove from inflight tracking once exchange acknowledges the order
+            self._inflight_orders[order.symbol].discard(order.oid)
+
+            # Ensure order is tracked in all sets if it's open (handles WebSocket arriving before REST API)
+            if not order.is_closed:
+                self._mem_open_orders[order.exchange].add(order.oid)
+                self._mem_symbol_orders[order.symbol].add(order.oid)
+                self._mem_symbol_open_orders[order.symbol].add(order.oid)
+                if order.is_cancel_failed:
+                    self._cancel_intent_oids.discard(order.oid)
+            else:
+                self._mem_open_orders[order.exchange].discard(order.oid)
+                self._mem_symbol_open_orders[order.symbol].discard(order.oid)
+                self._cancel_intent_oids.discard(order.oid)
+            return True
+
+    def mark_all_cancel_intent(self, symbol: str) -> None:
+        with self._order_lock:
+            oids = self._mem_symbol_open_orders.get(symbol, set())
+            self._cancel_intent_oids.update(oids)
+            self._log.debug(
+                f"Marked {len(oids)} orders as cancel intent for symbol {symbol}"
+            )
+
+    def mark_cancel_intent(self, oid: str) -> None:
+        with self._order_lock:
+            self._cancel_intent_oids.add(oid)
+            self._log.debug(f"Marked order {oid} as cancel intent")
+
+    def add_inflight_order(self, symbol: str, oid: str) -> None:
+        with self._order_lock:
+            self._inflight_orders[symbol].add(oid)
+
+    def get_inflight_orders(self, symbol: str) -> Set[str]:
+        with self._order_lock:
+            return self._inflight_orders.get(symbol, set()).copy()
+
+    async def wait_for_inflight_orders(
+        self, symbol: str, timeout: float = 5.0, interval: float = 0.1
+    ) -> None:
+        waited = 0.0
+        while self.get_inflight_orders(symbol) and waited < timeout:
+            await asyncio.sleep(interval)
+            waited += interval
+
+        remaining = self.get_inflight_orders(symbol)
+        if remaining:
+            self._log.warning(
+                f"Timeout waiting for inflight orders to resolve for {symbol}, "
+                f"remaining inflight: {remaining}"
+            )
+
+    # NOTE: this function is not for user to call, it is for internal use
+    def _get_all_balances_from_db(self, account_type: AccountType) -> List[Balance]:
+        with self._balance_lock:
+            return self._backend.get_all_balances(account_type)
+
+    # NOTE: this function is not for user to call, it is for internal use
+    def _get_all_positions_from_db(
+        self, exchange_id: ExchangeType
+    ) -> Dict[str, Position]:
+        return self._backend.get_all_positions(exchange_id)
+
+    @maybe
+    def get_order(self, oid: str) -> Optional[Order]:
+        with self._order_lock:
+            return self._backend.get_order(oid, self._mem_orders)
+
+    def get_symbol_orders(self, symbol: str, in_mem: bool = True) -> Set[str]:
+        """Get all orders for a symbol from memory and storage"""
+        with self._order_lock:
+            memory_orders = self._mem_symbol_orders.get(symbol, set())
+            if not in_mem:
+                storage_orders = self._backend.get_symbol_orders(symbol)
+                return memory_orders.union(storage_orders)
+            return memory_orders
+
+    def get_open_orders(
+        self,
+        symbol: str | None = None,
+        exchange: ExchangeType | None = None,
+        *,
+        include_canceling: bool = False,
+    ) -> Set[str]:
+        with self._order_lock:
+            if symbol is not None:
+                orders = self._mem_symbol_open_orders[symbol].copy()
+            elif exchange is not None:
+                orders = self._mem_open_orders[exchange].copy()
+            else:
+                raise ValueError("Either `symbol` or `exchange` must be specified")
+
+            if include_canceling:
+                return orders
+
+            return orders.difference(self._cancel_intent_oids)
+
+    ################ # parameter cache  ###################
+
+    def _ensure_redis_client(self) -> None:
+        """
+        Lazy initialization of Redis client.
+        Raises RuntimeError if Redis is not available.
+        """
+        if not self._redis_client_initialized:
+            self._redis_client = get_redis_client_if_available()
+            self._redis_client_initialized = True
+            if self._redis_client is None:
+                raise RuntimeError(
+                    "Redis client is not available. Please ensure Redis is running and configured properly."
+                )
+
+    def _get_redis_param_key(self, key: str) -> str:
+        """Generate Redis key for parameter storage"""
+        return f"nexus:param:{self.strategy_id}:{self.user_id}:{key}"
+
+    def get_param(
+        self,
+        key: str,
+        default: Any = None,
+        param_backend: ParamBackend = ParamBackend.MEMORY,
+    ) -> Any:
+        """
+        Get a parameter from the cache
+
+        Args:
+            key: Parameter key
+            default: Default value if key not found
+            param_backend: Storage backend (MEMORY or REDIS)
+
+        Returns:
+            Parameter value or default
+
+        Raises:
+            RuntimeError: If param_backend is REDIS but Redis is not available
+        """
+        if param_backend == ParamBackend.REDIS:
+            self._ensure_redis_client()
+            _redis = self._redis_client
+            assert _redis is not None
+            redis_key = self._get_redis_param_key(key)
+            value = cast(Optional[bytes], _redis.get(redis_key))
+            if value is None:
+                return default
+            # Deserialize the value
+            return msgspec.json.decode(value)
+        else:
+            return self._mem_params.get(key, default)
+
+    def set_param(
+        self, key: str, value: Any, param_backend: ParamBackend = ParamBackend.MEMORY
+    ) -> None:
+        """
+        Set a parameter in the cache
+
+        Args:
+            key: Parameter key
+            value: Parameter value
+            param_backend: Storage backend (MEMORY or REDIS)
+
+        Raises:
+            RuntimeError: If param_backend is REDIS but Redis is not available
+        """
+        if param_backend == ParamBackend.REDIS:
+            self._ensure_redis_client()
+            _redis = self._redis_client
+            assert _redis is not None
+            redis_key = self._get_redis_param_key(key)
+            # Serialize the value
+            serialized_value = msgspec.json.encode(value)
+            _redis.set(redis_key, serialized_value)
+        else:
+            self._mem_params[key] = value
+
+    def get_all_params(
+        self, param_backend: ParamBackend = ParamBackend.MEMORY
+    ) -> Dict[str, Any]:
+        """
+        Get all parameters from the cache
+
+        Args:
+            param_backend: Storage backend (MEMORY or REDIS)
+
+        Returns:
+            Dictionary of all parameters
+
+        Raises:
+            RuntimeError: If param_backend is REDIS but Redis is not available
+        """
+        if param_backend == ParamBackend.REDIS:
+            self._ensure_redis_client()
+            _redis = self._redis_client
+            assert _redis is not None
+            pattern = f"nexus:param:{self.strategy_id}:{self.user_id}:*"
+            params = {}
+            # Use SCAN to iterate over keys matching the pattern
+            cursor = 0
+            while True:
+                cursor, keys = _redis.scan(cursor, match=pattern, count=100)
+                for redis_key in keys:
+                    # Extract the parameter key from the Redis key
+                    key = (
+                        redis_key.decode()
+                        if isinstance(redis_key, bytes)
+                        else str(redis_key)
+                    )
+                    param_key = key.split(":", 4)[-1]  # Get the part after the last ':'
+                    value = cast(Optional[bytes], _redis.get(redis_key))
+                    if value is not None:
+                        params[param_key] = msgspec.json.decode(value)
+                if cursor == 0:
+                    break
+            return params
+        else:
+            return self._mem_params.copy()
+
+    def clear_param(
+        self,
+        key: Optional[str] = None,
+        param_backend: ParamBackend = ParamBackend.MEMORY,
+    ) -> None:
+        """
+        Clear parameter(s) from the cache
+
+        Args:
+            key: Parameter key to clear (None to clear all)
+            param_backend: Storage backend (MEMORY or REDIS)
+
+        Raises:
+            RuntimeError: If param_backend is REDIS but Redis is not available
+        """
+        if param_backend == ParamBackend.REDIS:
+            self._ensure_redis_client()
+            _redis = self._redis_client
+            assert _redis is not None
+            if key is None:
+                # Clear all parameters matching the pattern
+                pattern = f"nexus:param:{self.strategy_id}:{self.user_id}:*"
+                cursor = 0
+                while True:
+                    cursor, keys = _redis.scan(cursor, match=pattern, count=100)
+                    if keys:
+                        _redis.delete(*keys)
+                    if cursor == 0:
+                        break
+            else:
+                # Clear specific parameter
+                redis_key = self._get_redis_param_key(key)
+                _redis.delete(redis_key)
+        else:
+            if key is None:
+                # Clear all parameters
+                self._mem_params.clear()
+            else:
+                # Clear specific parameter
+                self._mem_params.pop(key, None)
+
+    ################ # Redis direct access  ###################
+
+    @property
+    def redis(self) -> Optional[redis.Redis]:
+        """
+        Direct access to Redis client for calling any Redis methods.
+
+        Returns:
+            Redis client instance
+
+        Raises:
+            RuntimeError: If Redis is not available
+
+        Example:
+            # Set a key
+            cache.redis.set("mykey", "myvalue")
+
+            # Get a key
+            value = cache.redis.get("mykey")
+
+            # Use hash operations
+            cache.redis.hset("myhash", "field1", "value1")
+            cache.redis.hget("myhash", "field1")
+
+            # Use list operations
+            cache.redis.lpush("mylist", "item1")
+            cache.redis.rpop("mylist")
+
+            # Use set operations
+            cache.redis.sadd("myset", "member1")
+            cache.redis.smembers("myset")
+
+            # Use sorted set operations
+            cache.redis.zadd("myzset", {"member1": 1.0})
+            cache.redis.zrange("myzset", 0, -1)
+
+            # Pipeline operations
+            pipe = cache.redis.pipeline()
+            pipe.set("key1", "value1")
+            pipe.set("key2", "value2")
+            pipe.execute()
+
+            # All standard Redis commands are available
+        """
+        self._ensure_redis_client()
+        return self._redis_client