diff --git a/src/driftpy/market_map/market_map.py b/src/driftpy/market_map/market_map.py index fc5a4708..2ef9f5da 100644 --- a/src/driftpy/market_map/market_map.py +++ b/src/driftpy/market_map/market_map.py @@ -103,7 +103,7 @@ def clear(self): def get_last_dump_filepath(self) -> str: return f"{market_type_to_string(self.market_type)}_{self.latest_slot}.pkl" - async def dump(self): + async def pre_dump(self) -> dict[str, bytes]: try: filters = [] if is_variant(self.market_type, "Perp"): @@ -142,13 +142,20 @@ async def dump(self): raw_bytes = base64.b64decode(market["account"]["data"][0]) raw[str(pubkey)] = raw_bytes + return raw + + except Exception as e: + print(f"error in marketmap pre-dump: {e}") + + def dump(self, raw: dict[str, bytes], filename: Optional[str] = None): + try: markets = [] for pubkey, market in raw.items(): markets.append(PickledData(pubkey=pubkey, data=compress(market))) - filename = ( + path = filename or ( f"{market_type_to_string(self.market_type)}_{self.latest_slot}.pkl" ) - with open(filename, "wb") as f: + with open(path, "wb") as f: pickle.dump(markets, f) except Exception as e: diff --git a/src/driftpy/pickle/vat.py b/src/driftpy/pickle/vat.py index f6726109..02e881a8 100644 --- a/src/driftpy/pickle/vat.py +++ b/src/driftpy/pickle/vat.py @@ -1,3 +1,4 @@ +import asyncio import pickle import os from dataclasses import dataclass @@ -15,7 +16,7 @@ def __init__( self, drift_client: DriftClient, users: UserMap, - user_stats: Optional[UserStatsMap], + user_stats: UserStatsMap, spot_markets: MarketMap, perp_markets: MarketMap, ): @@ -28,18 +29,37 @@ def __init__( self.perp_oracles = {} self.spot_oracles = {} - async def pickle(self): - await self.users.sync() - self.users.dump() + async def pickle(self, file_prefix: Optional[str] = None) -> dict[str, str]: + users_sync = asyncio.create_task(self.users.sync()) + user_stats_sync = asyncio.create_task(self.user_stats.sync()) + spot_markets_pre_dump = asyncio.create_task(self.spot_markets.pre_dump()) + perp_markets_pre_dump = asyncio.create_task(self.perp_markets.pre_dump()) + register_oracle_slot = asyncio.create_task(self.register_oracle_slot()) + + await asyncio.gather( + users_sync, + user_stats_sync, + spot_markets_pre_dump, + perp_markets_pre_dump, + register_oracle_slot, + ) + + spot_market_raw = spot_markets_pre_dump.result() + perp_market_raw = perp_markets_pre_dump.result() + + filenames = self.get_filenames(file_prefix) + + self.users.dump(filenames["users"]) - if self.user_stats is not None: - await self.user_stats.sync() - self.user_stats.dump() + self.user_stats.dump(filenames["userstats"]) - await self.spot_markets.dump() - await self.perp_markets.dump() + self.spot_markets.dump(spot_market_raw, filenames["spot_markets"]) - await self.dump_oracles() + self.perp_markets.dump(perp_market_raw, filenames["perp_markets"]) + + self.dump_oracles(filenames["spot_oracles"], filenames["perp_oracles"]) + + return filenames async def unpickle( self, @@ -56,8 +76,7 @@ async def unpickle( self.perp_markets.clear() await self.users.load(users_filename) - if self.user_stats is not None: - await self.user_stats.load(user_stats_filename) + await self.user_stats.load(user_stats_filename) await self.spot_markets.load(spot_markets_filename) await self.perp_markets.load(perp_markets_filename) @@ -67,7 +86,12 @@ async def unpickle( self.spot_markets, self.perp_markets, self.spot_oracles, self.perp_oracles ) - async def dump_oracles(self): + async def register_oracle_slot(self): + self.last_oracle_slot = (await self.drift_client.connection.get_slot()).value + + def dump_oracles( + self, spot_filepath: Optional[str] = None, perp_filepath: Optional[str] = None + ): perp_oracles = [] for market in self.drift_client.get_perp_market_accounts(): oracle_price_data = self.drift_client.get_oracle_price_data_for_perp_market( @@ -86,12 +110,12 @@ async def dump_oracles(self): PickledData(pubkey=market.market_index, data=oracle_price_data) ) - self.last_oracle_slot = (await self.drift_client.connection.get_slot()).value - - with open(f"perporacles_{self.last_oracle_slot}.pkl", "wb") as f: + perp_path = perp_filepath or f"perporacles_{self.last_oracle_slot}.pkl" + with open(perp_path, "wb") as f: pickle.dump(perp_oracles, f) - with open(f"spotoracles_{self.last_oracle_slot}.pkl", "wb") as f: + spot_path = spot_filepath or f"spotoracles_{self.last_oracle_slot}.pkl" + with open(spot_path, "wb") as f: pickle.dump(spot_oracles, f) def load_oracles( @@ -121,3 +145,29 @@ def load_oracles( ) # oracle.pubkey is actually a market index else: raise FileNotFoundError(f"File {spot_filename} not found") + + def get_filenames(self, prefix: Optional[str]) -> dict[str, str]: + filenames = {} + + usermap_slot = self.users.get_slot() + userstats_slot = self.user_stats.latest_slot + spot_markets_slot = self.spot_markets.latest_slot + perp_markets_slot = self.perp_markets.latest_slot + oracle_slot = self.last_oracle_slot + + if prefix: + filenames["users"] = f"{prefix}usermap_{usermap_slot}.pkl" + filenames["userstats"] = f"{prefix}userstats_{userstats_slot}.pkl" + filenames["spot_markets"] = f"{prefix}spot_{spot_markets_slot}.pkl" + filenames["perp_markets"] = f"{prefix}perp_{perp_markets_slot}.pkl" + filenames["spot_oracles"] = f"{prefix}spotoracles_{oracle_slot}.pkl" + filenames["perp_oracles"] = f"{prefix}perporacles_{oracle_slot}.pkl" + else: + filenames["users"] = f"usermap_{usermap_slot}.pkl" + filenames["userstats"] = f"userstats_{userstats_slot}.pkl" + filenames["spot_markets"] = f"spot_{spot_markets_slot}.pkl" + filenames["perp_markets"] = f"perp_{perp_markets_slot}.pkl" + filenames["spot_oracles"] = f"spotoracles_{oracle_slot}.pkl" + filenames["perp_oracles"] = f"perporacles_{oracle_slot}.pkl" + + return filenames diff --git a/src/driftpy/user_map/user_map.py b/src/driftpy/user_map/user_map.py index 567e7918..f10d8cd7 100644 --- a/src/driftpy/user_map/user_map.py +++ b/src/driftpy/user_map/user_map.py @@ -249,6 +249,6 @@ def dump(self, filename: Optional[str] = None): for pubkey, user in self.raw.items(): users.append(PickledData(pubkey=pubkey, data=compress(user))) self.last_dumped_slot = self.get_slot() - filename = filename or f"usermap_{self.last_dumped_slot}.pkl" - with open(filename, "wb") as f: + path = filename or f"usermap_{self.last_dumped_slot}.pkl" + with open(path, "wb") as f: pickle.dump(users, f, pickle.HIGHEST_PROTOCOL) diff --git a/src/driftpy/user_map/userstats_map.py b/src/driftpy/user_map/userstats_map.py index 56ef0320..686eb820 100644 --- a/src/driftpy/user_map/userstats_map.py +++ b/src/driftpy/user_map/userstats_map.py @@ -266,13 +266,13 @@ async def load(self, filename: Optional[str] = None): Pubkey.from_string(str(user_stat.pubkey)), DataAndSlot(slot, data) ) - def dump(self): + def dump(self, filename: Optional[str] = None): user_stats = [] for _pubkey, user_stat in self.raw.items(): decoded: UserStatsAccount = decode_user_stat(user_stat) auth = decoded.authority user_stats.append(PickledData(pubkey=auth, data=compress(user_stat))) self.last_dumped_slot = self.latest_slot - filename = f"userstats_{self.last_dumped_slot}.pkl" - with open(filename, "wb") as f: + path = filename or f"userstats_{self.last_dumped_slot}.pkl" + with open(path, "wb") as f: pickle.dump(user_stats, f, pickle.HIGHEST_PROTOCOL)