feat: Allow remote servers to return additional LUTs

Author: ChewingGlass

Cargo.lock

@@ -64,6 +64,7 @@ solana-account-decoder = { version = "2.2.3" }
 solana-clock = { version = "2.2.1" }
 solana-transaction-status = "2.2.3"
 bincode = { version = "1.3.3" }
+async-trait = "0.1.89"
 
 [profile.release]
 debug = true

Cargo.toml

@@ -123,6 +123,7 @@ pub enum Cmd {
 async fn simulate_task(client: &CliClient, task_key: Pubkey) -> Result<Option<SimulationResult>> {
     // Get the run instruction
     let run_ix_res = tuktuk_sdk::compiled_transaction::run_ix(
+        client.as_ref(),
         client.as_ref(),
         task_key,
         client.payer.pubkey(),
@@ -598,6 +599,7 @@ impl TaskCmd {
                 };
                 for (task_key, _) in tasks {
                     let run_ix_result = tuktuk_sdk::compiled_transaction::run_ix(
+                        client.as_ref(),
                         client.as_ref(),
                         task_key,
                         client.payer.pubkey(),

tuktuk-cli/src/cmd/task.rs

@@ -15,6 +15,7 @@ path = "src/main.rs"
 [dependencies]
 anchor-lang = { workspace = true }
 anchor-client = { workspace = true, features = ["async"] }
+async-trait = { workspace = true }
 bincode = { workspace = true }
 solana-sdk = { workspace = true }
 tokio = { workspace = true }

tuktuk-crank-turner/Cargo.toml

@@ -7,6 +7,7 @@ use solana_sdk::{
 };
 use tokio_graceful_shutdown::SubsystemHandle;
 use tracing::info;
+use tuktuk_sdk::client::LookupTableResolver;
 
 use crate::{cache::LookupTableRequest, sync};
 
@@ -26,6 +27,18 @@ impl LookupTablesSender {
     }
 }
 
+#[async_trait::async_trait]
+impl LookupTableResolver for LookupTablesSender {
+    async fn resolve_lookup_tables(
+        &self,
+        lookup_tables: Vec<Pubkey>,
+    ) -> Result<Vec<AddressLookupTableAccount>, tuktuk_sdk::error::Error> {
+        self.get_lookup_tables(lookup_tables).await.map_err(|_| {
+            tuktuk_sdk::error::Error::InvalidTransaction("lookup tables channel closed")
+        })
+    }
+}
+
 pub fn lookup_tables_channel() -> (LookupTablesSender, LookupTablesReceiver) {
     let (tx, rx) = sync::message_channel(100);
     (tx, rx)

tuktuk-crank-turner/src/cache/lookup_tables.rs

@@ -117,11 +117,6 @@ impl TimedTask {
 
         let task_queue = self.get_task_queue(ctx.clone()).await?;
 
-        let lookup_tables = ctx
-            .lookup_tables_client
-            .get_lookup_tables(task_queue.lookup_tables)
-            .await
-            .map_err(|_| anyhow::anyhow!("lookup tables channel closed"))?;
         let maybe_next_available = self.get_available_task_ids(ctx.clone()).await;
         let next_available = match maybe_next_available {
             Ok(next_available) => next_available,
@@ -163,7 +158,8 @@ impl TimedTask {
                 &self.task,
                 payer.pubkey(),
                 next_available,
-                lookup_tables,
+                task_queue.lookup_tables.clone(),
+                ctx.rpc_client.as_ref(),
             )
             .await
         };

tuktuk-crank-turner/src/task_processor.rs

@@ -68,7 +68,7 @@ where
     };
 
     let mut total_tasks = 0;
-    let mut has_unprocessed_tasks = true;
+    let mut has_unprocessed_tasks: bool;
     for AccountWithSeeds { account, seeds } in accounts.iter() {
         // Store original size before any reallocation
         original_sizes.push(account.data_len());

tuktuk-program/src/write_return_tasks.rs

@@ -28,12 +28,12 @@ tokio = { workspace = true }
 tokio-graceful-shutdown = { workspace = true }
 bincode = { workspace = true }
 itertools = { workspace = true }
-async-trait = { version = "0" }
 spl-associated-token-account = { workspace = true }
 spl-token = { workspace = true }
 tuktuk-program = { workspace = true }
 base64 = "0.22.1"
 serde_json = "1.0.135"
+async-trait = { workspace = true }
 bytemuck = "1.21.0"
 rand = "0.9"
 # For pubsub client

tuktuk-sdk/Cargo.toml

@@ -4,7 +4,10 @@ use anchor_lang::AccountDeserialize;
 use futures::{stream, StreamExt, TryFutureExt, TryStreamExt};
 use itertools::Itertools;
 pub use solana_client::nonblocking::rpc_client::RpcClient as SolanaRpcClient;
-use solana_sdk::{account::Account, pubkey::Pubkey};
+use solana_sdk::{
+    account::Account, address_lookup_table::state::AddressLookupTable,
+    message::AddressLookupTableAccount, pubkey::Pubkey,
+};
 
 use crate::error::Error;
 
@@ -61,6 +64,37 @@ impl GetAccount for SolanaRpcClient {
     }
 }
 
+#[async_trait::async_trait]
+pub trait LookupTableResolver {
+    async fn resolve_lookup_tables(
+        &self,
+        lookup_tables: Vec<Pubkey>,
+    ) -> Result<Vec<AddressLookupTableAccount>, Error>;
+}
+
+#[async_trait::async_trait]
+impl LookupTableResolver for SolanaRpcClient {
+    async fn resolve_lookup_tables(
+        &self,
+        lookup_tables: Vec<Pubkey>,
+    ) -> Result<Vec<AddressLookupTableAccount>, Error> {
+        let accounts = self.get_multiple_accounts(&lookup_tables).await?;
+        Ok(accounts
+            .into_iter()
+            .zip(lookup_tables.iter())
+            .filter_map(|(maybe_acc, pubkey)| {
+                maybe_acc.map(|acc| {
+                    let lut = AddressLookupTable::deserialize(&acc.data).map_err(Error::from)?;
+                    Ok(AddressLookupTableAccount {
+                        key: *pubkey,
+                        addresses: lut.addresses.to_vec(),
+                    })
+                })
+            })
+            .collect::<Result<Vec<_>, Error>>()?)
+    }
+}
+
 #[async_trait::async_trait]
 impl GetAnchorAccount for SolanaRpcClient {
     async fn anchor_account<T: AccountDeserialize>(

tuktuk-sdk/src/client.rs

@@ -6,15 +6,18 @@ use bytemuck::{bytes_of, Pod, Zeroable};
 use serde::Deserialize;
 use solana_client::client_error::reqwest;
 use solana_sdk::{
-    address_lookup_table::{state::AddressLookupTable, AddressLookupTableAccount},
+    address_lookup_table::AddressLookupTableAccount,
     ed25519_instruction::{DATA_START, PUBKEY_SERIALIZED_SIZE, SIGNATURE_SERIALIZED_SIZE},
     ed25519_program,
     instruction::Instruction,
     pubkey::Pubkey,
 };
 use tuktuk_program::{tuktuk, TaskQueueV0, TaskV0, TransactionSourceV0};
 
-use crate::{client::GetAnchorAccount, error::Error};
+use crate::{
+    client::{GetAnchorAccount, LookupTableResolver},
+    error::Error,
+};
 
 pub fn next_available_task_ids_excluding_in_progress(
     capacity: u16,
@@ -68,7 +71,8 @@ pub async fn run_ix_with_free_tasks(
     task: &TaskV0,
     payer: Pubkey,
     next_available: Vec<u16>,
-    lookup_tables: Vec<AddressLookupTableAccount>,
+    lookup_table_pubkeys: Vec<Pubkey>,
+    lut_resolver: &impl LookupTableResolver,
 ) -> Result<RunTaskResult, Error> {
     let transaction = &task.transaction;
 
@@ -125,6 +129,10 @@ pub async fn run_ix_with_free_tasks(
             ]
             .concat();
 
+            let lookup_tables = lut_resolver
+                .resolve_lookup_tables(lookup_table_pubkeys)
+                .await?;
+
             Ok(RunTaskResult {
                 instructions: vec![Instruction {
                     program_id: tuktuk_program::tuktuk::ID,
@@ -176,6 +184,10 @@ pub async fn run_ix_with_free_tasks(
             instruction_data.extend_from_slice(&signature);
             instruction_data.extend_from_slice(&message);
 
+            // Combine lookup table pubkeys from task queue and remote transaction
+            let all_lut_pubkeys = [lookup_table_pubkeys, remote_transaction.lookup_tables].concat();
+            let lookup_tables = lut_resolver.resolve_lookup_tables(all_lut_pubkeys).await?;
+
             Ok(RunTaskResult {
                 lookup_tables,
                 instructions: vec![
@@ -216,6 +228,7 @@ pub async fn run_ix_with_free_tasks(
 
 pub async fn run_ix(
     client: &impl GetAnchorAccount,
+    lut_resolver: &impl LookupTableResolver,
     task_key: Pubkey,
     payer: Pubkey,
     in_progress_task_ids: &HashSet<u16>,
@@ -239,22 +252,15 @@ pub async fn run_ix(
         rand::random_range(0..task_queue.task_bitmap.len()),
     )?;
 
-    let lookup_tables = client
-        .accounts(&task_queue.lookup_tables)
-        .await?
-        .into_iter()
-        .filter_map(|(addr, raw)| {
-            raw.map(|acc| {
-                let lut = AddressLookupTable::deserialize(&acc.data).map_err(Error::from)?;
-                Ok::<AddressLookupTableAccount, Error>(AddressLookupTableAccount {
-                    key: addr,
-                    addresses: lut.addresses.to_vec(),
-                })
-            })
-        })
-        .collect::<Result<Vec<_>, _>>()?;
-
-    run_ix_with_free_tasks(task_key, &task, payer, next_available, lookup_tables).await
+    run_ix_with_free_tasks(
+        task_key,
+        &task,
+        payer,
+        next_available,
+        task_queue.lookup_tables.clone(),
+        lut_resolver,
+    )
+    .await
 }
 
 #[derive(Default, Debug, Copy, Clone, Zeroable, Pod, Eq, PartialEq)]
@@ -281,12 +287,14 @@ struct RemoteResponse {
     transaction: String,
     remaining_accounts: Vec<RemoteAccountMeta>,
     signature: String,
+    lookup_tables: Option<Vec<String>>,
 }
 
 struct FetchedRemoteResponse {
     transaction: Vec<u8>,
     remaining_accounts: Vec<AccountMeta>,
     signature: Vec<u8>,
+    lookup_tables: Vec<Pubkey>,
 }
 
 async fn fetch_remote_transaction(
@@ -327,9 +335,17 @@ async fn fetch_remote_transaction(
         .decode(&json.signature)
         .map_err(Error::from)?;
 
+    let lookup_tables = json
+        .lookup_tables
+        .unwrap_or_default()
+        .into_iter()
+        .map(|key| Pubkey::from_str(&key).map_err(Error::from))
+        .collect::<Result<Vec<_>, Error>>()?;
+
     Ok(FetchedRemoteResponse {
         transaction: transaction_bytes,
         remaining_accounts,
         signature: signature_bytes,
+        lookup_tables,
     })
 }