implement add_mapping

Makefile

@@ -8,4 +8,4 @@ install:
 	poetry install
 
 test:
-	poetry run pytest
+	poetry run pytest -rx

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,13 +29,15 @@
     ReferencePublishers,
 )
 from program_admin.util import (
+    MAPPING_ACCOUNT_PRODUCT_LIMIT,
     MAPPING_ACCOUNT_SIZE,
     PRICE_ACCOUNT_V1_SIZE,
     PRICE_ACCOUNT_V2_SIZE,
     PRODUCT_ACCOUNT_SIZE,
     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,

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,

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,

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,
     )