diff --git a/Makefile b/Makefile index 82728d1..31839ce 100644 --- a/Makefile +++ b/Makefile @@ -8,4 +8,4 @@ install: poetry install test: - poetry run pytest + poetry run pytest -rx diff --git a/program_admin/__init__.py b/program_admin/__init__.py index 4561252..9db8d68 100644 --- a/program_admin/__init__.py +++ b/program_admin/__init__.py @@ -1,5 +1,6 @@ import asyncio import json +import math import os from pathlib import Path from typing import Dict, List, Literal, Optional, Tuple @@ -28,6 +29,7 @@ ReferencePublishers, ) from program_admin.util import ( + MAPPING_ACCOUNT_PRODUCT_LIMIT, MAPPING_ACCOUNT_SIZE, PRICE_ACCOUNT_V1_SIZE, PRICE_ACCOUNT_V2_SIZE, @@ -35,6 +37,7 @@ account_exists, compute_transaction_size, get_actual_signers, + get_available_mapping_account_key, recent_blockhash, sort_mapping_account_keys, ) @@ -246,8 +249,13 @@ async def sync( await self.send_transaction(authority_instructions, authority_signers) # Sync mapping accounts + + # Create all the mapping accounts we need for the the number of product accounts + num_mapping_accounts = math.ceil( + len(ref_products) / MAPPING_ACCOUNT_PRODUCT_LIMIT + ) mapping_instructions, mapping_keypairs = await self.sync_mapping_instructions( - generate_keys + generate_keys, num_mapping_accounts ) if mapping_instructions: @@ -257,10 +265,6 @@ async def sync( await self.refresh_program_accounts() - # FIXME: We should check if the mapping account has enough space to - # add/remove new products. That is not urgent because we are around 10% - # of the first mapping account capacity. - # Sync product/price accounts product_transactions: List[asyncio.Task[None]] = [] @@ -334,28 +338,26 @@ async def sync( return instructions async def sync_mapping_instructions( - self, - generate_keys: bool, + self, generate_keys: bool, num_mapping_accounts: int = 1 ) -> Tuple[List[TransactionInstruction], List[Keypair]]: - mapping_chain = sort_mapping_account_keys(list(self._mapping_accounts.values())) funding_keypair = load_keypair("funding", key_dir=self.key_dir) - mapping_0_keypair = load_keypair( + mapping_keypair_0 = load_keypair( "mapping_0", key_dir=self.key_dir, generate=generate_keys ) - instructions: List[TransactionInstruction] = [] - if not mapping_chain: - logger.info("Creating new mapping account") + instructions: List[TransactionInstruction] = [] - if not ( - await account_exists(self.rpc_endpoint, mapping_0_keypair.public_key) + # Create initial mapping account + if len(self._mapping_accounts) < 1: + if not await account_exists( + self.rpc_endpoint, mapping_keypair_0.public_key ): logger.debug("Building system.program.create_account instruction") instructions.append( system_program.create_account( system_program.CreateAccountParams( from_pubkey=funding_keypair.public_key, - new_account_pubkey=mapping_0_keypair.public_key, + new_account_pubkey=mapping_keypair_0.public_key, # FIXME: Change to minimum rent-exempt amount lamports=await self.fetch_minimum_balance( MAPPING_ACCOUNT_SIZE @@ -371,11 +373,55 @@ async def sync_mapping_instructions( pyth_program.init_mapping( self.program_key, funding_keypair.public_key, - mapping_0_keypair.public_key, + mapping_keypair_0.public_key, ) ) - return (instructions, [funding_keypair, mapping_0_keypair]) + # Add extra mapping accounts + mapping_keypairs: List[Keypair] = [] + if len(self._mapping_accounts) < num_mapping_accounts: + if num_mapping_accounts > 1: + mapping_keypairs: List[Keypair] = [ + load_keypair( + f"mapping_{n}", key_dir=self.key_dir, generate=generate_keys + ) + for n in range(1, num_mapping_accounts) + ] + + tail_mapping_keypair = mapping_keypair_0 + for mapping_keypair in mapping_keypairs: + if not ( + await account_exists(self.rpc_endpoint, mapping_keypair.public_key) + ): + logger.debug("Building system.program.create_account instruction") + instructions.append( + system_program.create_account( + system_program.CreateAccountParams( + from_pubkey=funding_keypair.public_key, + new_account_pubkey=mapping_keypair.public_key, + # FIXME: Change to minimum rent-exempt amount + lamports=await self.fetch_minimum_balance( + MAPPING_ACCOUNT_SIZE + ), + space=MAPPING_ACCOUNT_SIZE, + program_id=self.program_key, + ) + ) + ) + + logger.debug("Building pyth_program.add_mapping instruction") + instructions.append( + pyth_program.add_mapping( + self.program_key, + funding_keypair.public_key, + tail_mapping_keypair.public_key, + mapping_keypair.public_key, + ) + ) + + tail_mapping_keypair = mapping_keypair + + return (instructions, [funding_keypair, mapping_keypair_0] + mapping_keypairs) async def sync_product_instructions( self, @@ -385,8 +431,10 @@ async def sync_product_instructions( ) -> Tuple[List[TransactionInstruction], List[Keypair]]: instructions: List[TransactionInstruction] = [] funding_keypair = load_keypair("funding", key_dir=self.key_dir) - mapping_chain = sort_mapping_account_keys(list(self._mapping_accounts.values())) - mapping_keypair = load_keypair(mapping_chain[-1], key_dir=self.key_dir) + mapping_keypair = load_keypair( + get_available_mapping_account_key(list(self._mapping_accounts.values())), + key_dir=self.key_dir, + ) product_keypair = load_keypair( f"product_{product['jump_symbol']}", key_dir=self.key_dir, diff --git a/program_admin/instructions.py b/program_admin/instructions.py index 0f1d59c..5b8532b 100644 --- a/program_admin/instructions.py +++ b/program_admin/instructions.py @@ -8,9 +8,8 @@ from program_admin.types import ReferenceAuthorityPermissions from program_admin.util import encode_product_metadata, get_permissions_account -# TODO: Implement add_mapping instruction - COMMAND_INIT_MAPPING = 0 +COMMAND_ADD_MAPPING = 1 COMMAND_ADD_PRODUCT = 2 COMMAND_UPD_PRODUCT = 3 COMMAND_ADD_PRICE = 4 @@ -59,6 +58,39 @@ def init_mapping( ) +def add_mapping( + program_key: PublicKey, + funding_key: PublicKey, + tail_mapping_key: PublicKey, + mapping_key: PublicKey, +) -> TransactionInstruction: + """ + Pyth program init_mapping instruction + + accounts: + - funding account (signer, writable) + - tail mapping account (signer, writable) + - mapping account (signer, writable) + """ + layout = Struct("version" / Int32ul, "command" / Int32sl) + data = layout.build(dict(version=PROGRAM_VERSION, command=COMMAND_ADD_MAPPING)) + + permissions_account = get_permissions_account( + program_key, AUTHORITY_PERMISSIONS_PDA_SEED + ) + + return TransactionInstruction( + data=data, + keys=[ + AccountMeta(pubkey=funding_key, is_signer=True, is_writable=True), + AccountMeta(pubkey=tail_mapping_key, is_signer=True, is_writable=True), + AccountMeta(pubkey=mapping_key, is_signer=True, is_writable=True), + AccountMeta(pubkey=permissions_account, is_signer=False, is_writable=True), + ], + program_id=program_key, + ) + + def add_product( program_key: PublicKey, funding_key: PublicKey, diff --git a/program_admin/util.py b/program_admin/util.py index 3b7b515..fac6359 100644 --- a/program_admin/util.py +++ b/program_admin/util.py @@ -103,6 +103,26 @@ def sort_mapping_account_keys(accounts: List[PythMappingAccount]) -> List[Public return sorted_keys +def get_available_mapping_account_key( + mapping_accounts: List[PythMappingAccount], +) -> PublicKey: + """ + Returns the first non-full mapping account. + """ + # Create a dictionary with account key as key and product count as value + num_products_by_key = { + account.public_key: account.data.product_count for account in mapping_accounts + } + + sorted_keys = sort_mapping_account_keys(mapping_accounts) + + for key in sorted_keys: + if num_products_by_key[key] < MAPPING_ACCOUNT_PRODUCT_LIMIT: + return key + + raise RuntimeError("All mapping accounts are full") + + def apply_overrides( ref_permissions: ReferencePermissions, ref_overrides: ReferenceOverrides, diff --git a/tests/test_sync.py b/tests/test_sync.py index 4a0370c..3145a04 100644 --- a/tests/test_sync.py +++ b/tests/test_sync.py @@ -21,6 +21,7 @@ from program_admin.util import apply_overrides LOGGER = logging.getLogger(__name__) +LOGGER.setLevel(logging.getLevelName(os.getenv("LOG_LEVEL", "INFO").upper())) BTC_USD = { "account": "", @@ -221,7 +222,7 @@ def localhost_overrides_json(): @pytest.fixture async def validator(): process = await asyncio.create_subprocess_shell( - "solana-test-validator", + "solana-test-validator --reset", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, )