frank/sequencer support (#163)

* sequencer support

* remove debug print

* fix sequencer logic bug

* add some guardrails

* sequencing per subaccount

* remove this

* remove comment

src/driftpy/addresses.py

@@ -96,6 +96,14 @@ def get_prelaunch_oracle_public_key(program_id: Pubkey, market_index: int) -> Pu
     )[0]
 
 
+def get_sequencer_public_key_and_bump(
+    program_id: Pubkey, payer: Pubkey, subaccount_id: int
+) -> tuple[Pubkey, int]:
+    return Pubkey.find_program_address(
+        [(str(subaccount_id)).encode(), bytes(payer)], program_id
+    )
+
+
 # program = Pubkey("9jwr5nC2f9yAraXrg4UzHXmCX3vi9FQkjD6p9e8bRqNa")
 # auth = Pubkey("D78cqss3dbU1aJAs5qeuhLi8Rqa2CL4Kzkr3VzdgN5F6")
 # == EjQ8rFmR4hd9faX1TYLkqCTsAkyjJ4qUKBuagtmVG3cP

src/driftpy/constants/config.py

@@ -32,6 +32,12 @@
 DriftEnv = Literal["devnet", "mainnet"]
 
 DRIFT_PROGRAM_ID = Pubkey.from_string("dRiftyHA39MWEi3m9aunc5MzRF1JYuBsbn6VPcn33UH")
+SEQUENCER_PROGRAM_ID = Pubkey.from_string(
+    "GDDMwNyyx8uB6zrqwBFHjLLG3TBYk2F8Az4yrQC5RzMp"
+)
+DEVNET_SEQUENCER_PROGRAM_ID = Pubkey.from_string(
+    "FBngRHN4s5cmHagqy3Zd6xcK3zPJBeX5DixtHFbBhyCn"
+)
 
 
 @dataclass
@@ -214,7 +220,7 @@ async def find_all_market_and_oracles(
         return perp_ds, spot_ds, full_oracle_wrappers
 
 
-def decode_account(account_data: str, program: Program) -> PerpMarketAccount:
+def decode_account(account_data: str, program: Program):
     decoded_data = b64decode(account_data)
     return program.coder.accounts.decode(decoded_data)
 

src/driftpy/drift_client.py

@@ -1,5 +1,6 @@
 import json
 import os
+import anchorpy
 from deprecated import deprecated
 import requests
 from solders.pubkey import Pubkey
@@ -31,7 +32,14 @@
 from driftpy.drift_user import DriftUser
 from driftpy.accounts import *
 
-from driftpy.constants.config import DriftEnv, DRIFT_PROGRAM_ID, configs
+from driftpy.constants.config import (
+    DriftEnv,
+    DRIFT_PROGRAM_ID,
+    configs,
+    SEQUENCER_PROGRAM_ID,
+    DEVNET_SEQUENCER_PROGRAM_ID,
+    decode_account,
+)
 
 from typing import Tuple, Union, Optional, List
 from driftpy.math.perp_position import is_available
@@ -40,6 +48,7 @@
 from driftpy.name import encode_name
 from driftpy.tx.standard_tx_sender import StandardTxSender
 from driftpy.tx.types import TxSender, TxSigAndSlot
+from driftpy.addresses import get_sequencer_public_key_and_bump
 from spl.token.constants import ASSOCIATED_TOKEN_PROGRAM_ID
 from solders.system_program import ID as SYS_PROGRAM_ID
 
@@ -82,6 +91,7 @@ def __init__(
         sub_account_ids: Optional[list[int]] = None,
         market_lookup_table: Optional[Pubkey] = None,
         jito_params: Optional[JitoParams] = None,
+        enforce_tx_sequencing: bool = False,
     ):
         """Initializes the drift client object
 
@@ -145,6 +155,30 @@ def __init__(
 
         self.tx_version = tx_version if tx_version is not None else Legacy
 
+        self.enforce_tx_sequencing = enforce_tx_sequencing
+        if self.enforce_tx_sequencing is True:
+            file = Path(str(driftpy.__path__[0]) + "/idl/sequence_enforcer.json")
+            with file.open() as f:
+                raw = file.read_text()
+            idl = Idl.from_json(raw)
+
+            provider = Provider(connection, wallet, opts)
+            self.sequence_enforcer_pid = (
+                SEQUENCER_PROGRAM_ID
+                if env == "mainnet"
+                else DEVNET_SEQUENCER_PROGRAM_ID
+            )
+            self.sequence_enforcer_program = Program(
+                idl,
+                self.sequence_enforcer_pid,
+                provider,
+            )
+            self.sequence_number_by_subaccount = {}
+            self.sequence_bump_by_subaccount = {}
+            self.sequence_initialized_by_subaccount = {}
+            self.sequence_address_by_subaccount = {}
+            self.resetting_sequence = False
+
         if jito_params is not None:
             from driftpy.tx.jito_tx_sender import JitoTxSender
 
@@ -165,6 +199,8 @@ def __init__(
 
     async def subscribe(self):
         await self.account_subscriber.subscribe()
+        if self.enforce_tx_sequencing:
+            await self.load_sequence_info()
         for sub_account_id in self.sub_account_ids:
             await self.add_user(sub_account_id)
 
@@ -297,18 +333,38 @@ async def send_ixs(
         signers=None,
         lookup_tables: list[AddressLookupTableAccount] = None,
         tx_version: Optional[Union[Legacy, int]] = None,
+        sequencer_subaccount: Optional[int] = None,
     ) -> TxSigAndSlot:
         if isinstance(ixs, Instruction):
             ixs = [ixs]
 
         if not tx_version:
             tx_version = self.tx_version
 
+        compute_unit_instructions = []
         if self.tx_params.compute_units is not None:
-            ixs.insert(0, set_compute_unit_limit(self.tx_params.compute_units))
+            compute_unit_instructions.append(
+                set_compute_unit_limit(self.tx_params.compute_units)
+            )
 
         if self.tx_params.compute_units_price is not None:
-            ixs.insert(1, set_compute_unit_price(self.tx_params.compute_units_price))
+            compute_unit_instructions.append(
+                set_compute_unit_price(self.tx_params.compute_units_price)
+            )
+
+        ixs[0:0] = compute_unit_instructions
+
+        subaccount = sequencer_subaccount or self.active_sub_account_id
+
+        if (
+            self.enforce_tx_sequencing
+            and self.sequence_initialized_by_subaccount[subaccount]
+            and not self.resetting_sequence
+        ):
+            sequence_instruction = self.get_check_and_set_sequence_number_ix(
+                self.sequence_number_by_subaccount[subaccount], subaccount
+            )
+            ixs.insert(len(compute_unit_instructions), sequence_instruction)
 
         if tx_version == Legacy:
             tx = await self.tx_sender.get_legacy_tx(ixs, self.wallet.payer, signers)
@@ -2788,3 +2844,108 @@ def get_update_prelaunch_oracle_ix(self, market_index: int):
                 }
             )
         )
+
+    async def init_sequence(self, subaccount: int = 0) -> Signature:
+        try:
+            sig = (await self.send_ixs([self.get_sequence_init_ix(subaccount)])).tx_sig
+            self.sequence_initialized_by_subaccount[subaccount] = True
+            return sig
+        except Exception as e:
+            print(f"WARNING: failed to initialize sequence: {e}")
+
+    def get_sequence_init_ix(self, subaccount: int = 0) -> Instruction:
+        if self.enforce_tx_sequencing is False:
+            raise ValueError("tx sequencing is disabled")
+        return self.sequence_enforcer_program.instruction["initialize"](
+            self.sequence_bump_by_subaccount[subaccount],
+            str(subaccount),
+            ctx=Context(
+                accounts={
+                    "sequence_account": self.sequence_address_by_subaccount[subaccount],
+                    "authority": self.wallet.payer.pubkey(),
+                    "system_program": ID,
+                }
+            ),
+        )
+
+    async def reset_sequence_number(
+        self, sequence_number: int = 0, subaccount: int = 0
+    ) -> Signature:
+        try:
+            ix = self.get_reset_sequence_number_ix(sequence_number)
+            self.resetting_sequence = True
+            sig = (await self.send_ixs(ix)).tx_sig
+            self.resetting_sequence = False
+            self.sequence_number_by_subaccount[subaccount] = sequence_number
+            return sig
+        except Exception as e:
+            print(f"WARNING: failed to reset sequence number: {e}")
+
+    def get_reset_sequence_number_ix(
+        self, sequence_number: int, subaccount: int = 0
+    ) -> Instruction:
+        if self.enforce_tx_sequencing is False:
+            raise ValueError("tx sequencing is disabled")
+        return self.sequence_enforcer_program.instruction["reset_sequence_number"](
+            sequence_number,
+            ctx=Context(
+                accounts={
+                    "sequence_account": self.sequence_address_by_subaccount[subaccount],
+                    "authority": self.wallet.payer.pubkey(),
+                }
+            ),
+        )
+
+    def get_check_and_set_sequence_number_ix(
+        self, sequence_number: Optional[int] = None, subaccount: int = 0
+    ):
+        if self.enforce_tx_sequencing is False:
+            raise ValueError("tx sequencing is disabled")
+        sequence_number = (
+            sequence_number or self.sequence_number_by_subaccount[subaccount]
+        )
+
+        if (
+            sequence_number < self.sequence_number_by_subaccount[subaccount] - 1
+        ):  # we increment after creating the ix, so we check - 1
+            print(
+                f"WARNING: sequence number {sequence_number} < last used {self.sequence_number_by_subaccount[subaccount] - 1}"
+            )
+
+        ix = self.sequence_enforcer_program.instruction[
+            "check_and_set_sequence_number"
+        ](
+            sequence_number,
+            ctx=Context(
+                accounts={
+                    "sequence_account": self.sequence_address_by_subaccount[subaccount],
+                    "authority": self.wallet.payer.pubkey(),
+                }
+            ),
+        )
+
+        self.sequence_number_by_subaccount[subaccount] += 1
+        return ix
+
+    async def load_sequence_info(self):
+        for subaccount in self.sub_account_ids:
+            address, bump = get_sequencer_public_key_and_bump(
+                self.sequence_enforcer_pid, self.wallet.payer.pubkey(), subaccount
+            )
+            try:
+                sequence_account_raw = await self.sequence_enforcer_program.account[
+                    "SequenceAccount"
+                ].fetch(address)
+            except anchorpy.error.AccountDoesNotExistError as e:
+                self.sequence_address_by_subaccount[subaccount] = address
+                self.sequence_number_by_subaccount[subaccount] = 1
+                self.sequence_bump_by_subaccount[subaccount] = bump
+                self.sequence_initialized_by_subaccount[subaccount] = False
+                continue
+            sequence_account = cast(SequenceAccount, sequence_account_raw)
+            self.sequence_number_by_subaccount[subaccount] = (
+                sequence_account.sequence_num + 1
+            )
+            self.sequence_bump_by_subaccount[subaccount] = bump
+            self.sequence_initialized_by_subaccount[subaccount] = True
+            self.sequence_address_by_subaccount[subaccount] = address

src/driftpy/idl/sequence_enforcer.json

@@ -0,0 +1,103 @@
+{
+    "version": "0.1.0",
+    "name": "sequence_enforcer",
+    "instructions": [
+      {
+        "name": "initialize",
+        "accounts": [
+          {
+            "name": "sequenceAccount",
+            "isMut": true,
+            "isSigner": false
+          },
+          {
+            "name": "authority",
+            "isMut": false,
+            "isSigner": true
+          },
+          {
+            "name": "systemProgram",
+            "isMut": false,
+            "isSigner": false
+          }
+        ],
+        "args": [
+          {
+            "name": "bump",
+            "type": "u8"
+          },
+          {
+            "name": "sym",
+            "type": "string"
+          }
+        ]
+      },
+      {
+        "name": "resetSequenceNumber",
+        "accounts": [
+          {
+            "name": "sequenceAccount",
+            "isMut": true,
+            "isSigner": false
+          },
+          {
+            "name": "authority",
+            "isMut": false,
+            "isSigner": true
+          }
+        ],
+        "args": [
+          {
+            "name": "sequenceNum",
+            "type": "u64"
+          }
+        ]
+      },
+      {
+        "name": "checkAndSetSequenceNumber",
+        "accounts": [
+          {
+            "name": "sequenceAccount",
+            "isMut": true,
+            "isSigner": false
+          },
+          {
+            "name": "authority",
+            "isMut": false,
+            "isSigner": true
+          }
+        ],
+        "args": [
+          {
+            "name": "sequenceNum",
+            "type": "u64"
+          }
+        ]
+      }
+    ],
+    "accounts": [
+      {
+        "name": "SequenceAccount",
+        "type": {
+          "kind": "struct",
+          "fields": [
+            {
+              "name": "sequenceNum",
+              "type": "u64"
+            },
+            {
+              "name": "authority",
+              "type": "publicKey"
+            }
+          ]
+        }
+      }
+    ],
+    "errors": [
+      {
+        "code": 6000,
+        "name": "SequenceOutOfOrder",
+        "msg": "Sequence out of order"
+      }
+    ]
+  }

src/driftpy/types.py

@@ -1324,3 +1324,9 @@ class PrelaunchOracle:
     amm_last_update_slot: int
     last_update_slot: int
     perp_market_index: int
+
+
+@dataclass
+class SequenceAccount:
+    sequence_num: int
+    authority: Pubkey