Merge pull request #43 from helium/refactor/mutexes

Refactor to remove mutexes

Author: ChewingGlass

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

@@ -0,0 +1,107 @@
+use std::{collections::HashMap, sync::Arc};
+
+use solana_client::nonblocking::rpc_client::RpcClient;
+use solana_sdk::{
+    address_lookup_table::{state::AddressLookupTable, AddressLookupTableAccount},
+    pubkey::Pubkey,
+};
+use tokio_graceful_shutdown::SubsystemHandle;
+use tracing::info;
+
+use crate::{cache::LookupTableRequest, sync};
+
+pub type LookupTablesSender = sync::MessageSender<LookupTableRequest>;
+pub type LookupTablesReceiver = sync::MessageReceiver<LookupTableRequest>;
+
+impl LookupTablesSender {
+    pub async fn get_lookup_tables(
+        &self,
+        lookup_table_keys: Vec<Pubkey>,
+    ) -> Result<Vec<AddressLookupTableAccount>, sync::Error> {
+        self.request(|resp| LookupTableRequest::Get {
+            lookup_table_keys,
+            resp,
+        })
+        .await
+    }
+}
+
+pub fn lookup_tables_channel() -> (LookupTablesSender, LookupTablesReceiver) {
+    let (tx, rx) = sync::message_channel(100);
+    (tx, rx)
+}
+
+pub struct LookupTablesCache {
+    rpc_client: Arc<RpcClient>,
+    cache: HashMap<Pubkey, AddressLookupTableAccount>,
+    receiver: LookupTablesReceiver,
+}
+
+impl LookupTablesCache {
+    pub fn new(rpc_client: Arc<RpcClient>, receiver: LookupTablesReceiver) -> Self {
+        Self {
+            rpc_client,
+            cache: HashMap::new(),
+            receiver,
+        }
+    }
+
+    pub async fn run(mut self, handle: SubsystemHandle) -> anyhow::Result<()> {
+        info!("starting lookup tables cache");
+        loop {
+            tokio::select! {
+                _ = handle.on_shutdown_requested() => {
+                    info!("shutting down lookup tables cache");
+                    break;
+                }
+                Some(req) = self.receiver.recv() => {
+                    match req {
+                        LookupTableRequest::Get {
+                            lookup_table_keys,
+                            resp,
+                        } => {
+                            let mut result = Vec::with_capacity(lookup_table_keys.len());
+                            let mut missing_keys = Vec::new();
+
+                            // First check cache
+                            for key in &lookup_table_keys {
+                                if let Some(table) = self.cache.get(key) {
+                                    result.push(table.clone());
+                                } else {
+                                    missing_keys.push(*key);
+                                }
+                            }
+
+                            // Fetch missing tables
+                            if !missing_keys.is_empty() {
+                                for key in missing_keys {
+                                    if let Ok(account) = self.rpc_client.get_account(&key).await {
+                                        if let Ok(lut) = AddressLookupTable::deserialize(&account.data) {
+                                            let table = AddressLookupTableAccount {
+                                                key,
+                                                addresses: lut.addresses.to_vec(),
+                                            };
+                                            self.cache.insert(key, table.clone());
+                                            result.push(table);
+                                        }
+                                    }
+                                }
+                            }
+
+                            // Sort result to match input order
+                            result.sort_by_key(|table| {
+                                lookup_table_keys.iter()
+                                    .position(|key| key == &table.key)
+                                    .unwrap_or(usize::MAX)
+                            });
+
+                            resp.send(result);
+                        }
+                    }
+                }
+                else => break,
+            }
+        }
+        Ok(())
+    }
+}

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

@@ -0,0 +1,51 @@
+use std::collections::HashSet;
+
+use solana_sdk::{address_lookup_table::AddressLookupTableAccount, pubkey::Pubkey};
+use tuktuk_program::TaskQueueV0;
+
+use crate::sync::ResponseSender;
+
+// Request types for state management
+#[derive(Debug)]
+pub enum TaskStateRequest {
+    AddInProgressTasks {
+        pubkey: Pubkey,
+        task_ids: HashSet<u16>,
+    },
+    RemoveInProgressTasks {
+        pubkey: Pubkey,
+        task_ids: HashSet<u16>,
+    },
+    GetInProgressTasks {
+        pubkey: Pubkey,
+        resp: ResponseSender<HashSet<u16>>,
+    },
+}
+
+#[derive(Debug)]
+pub enum LookupTableRequest {
+    Get {
+        lookup_table_keys: Vec<Pubkey>,
+        resp: ResponseSender<Vec<AddressLookupTableAccount>>,
+    },
+}
+
+#[derive(Debug)]
+pub enum TaskQueueRequest {
+    Get {
+        pubkey: Pubkey,
+        resp: ResponseSender<Option<TaskQueueV0>>,
+    },
+    Update {
+        pubkey: Pubkey,
+        queue: Box<TaskQueueV0>,
+    },
+}
+
+mod lookup_tables;
+mod task_queues;
+mod task_state;
+
+pub use lookup_tables::{lookup_tables_channel, LookupTablesCache, LookupTablesSender};
+pub use task_queues::{task_queues_channel, TaskQueueCache, TaskQueuesSender};
+pub use task_state::{task_state_channel, TaskStateCache, TaskStateSender};

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

@@ -0,0 +1,69 @@
+use std::collections::HashMap;
+
+use solana_sdk::pubkey::Pubkey;
+use tokio_graceful_shutdown::SubsystemHandle;
+use tracing::info;
+use tuktuk_program::TaskQueueV0;
+
+use crate::{cache::TaskQueueRequest, sync};
+
+pub type TaskQueuesSender = sync::MessageSender<TaskQueueRequest>;
+pub type TaskQueuesReceiver = sync::MessageReceiver<TaskQueueRequest>;
+
+impl TaskQueuesSender {
+    pub async fn get_task_queue(&self, pubkey: Pubkey) -> Result<Option<TaskQueueV0>, sync::Error> {
+        self.request(|resp| TaskQueueRequest::Get { pubkey, resp })
+            .await
+    }
+
+    pub async fn update_task_queue(&self, pubkey: Pubkey, queue: TaskQueueV0) {
+        self.send(TaskQueueRequest::Update {
+            pubkey,
+            queue: Box::new(queue),
+        })
+        .await
+    }
+}
+
+pub fn task_queues_channel() -> (TaskQueuesSender, TaskQueuesReceiver) {
+    let (tx, rx) = sync::message_channel(100);
+    (tx, rx)
+}
+
+pub struct TaskQueueCache {
+    cache: HashMap<Pubkey, TaskQueueV0>,
+    receiver: TaskQueuesReceiver,
+}
+
+impl TaskQueueCache {
+    pub fn new(receiver: TaskQueuesReceiver) -> Self {
+        Self {
+            cache: HashMap::new(),
+            receiver,
+        }
+    }
+
+    pub async fn run(mut self, handle: SubsystemHandle) -> anyhow::Result<()> {
+        info!("starting task queue cache");
+        loop {
+            tokio::select! {
+                _ = handle.on_shutdown_requested() => {
+                    info!("shutting down task queue cache");
+                    break;
+                }
+                Some(req) = self.receiver.recv() => {
+                    match req {
+                        TaskQueueRequest::Get { pubkey, resp } => {
+                            resp.send(self.cache.get(&pubkey).cloned());
+                        }
+                        TaskQueueRequest::Update { pubkey, queue } => {
+                            self.cache.insert(pubkey, *queue);
+                        }
+                    }
+                }
+                else => break,
+            }
+        }
+        Ok(())
+    }
+}

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

@@ -0,0 +1,80 @@
+use std::collections::{HashMap, HashSet};
+
+use solana_sdk::pubkey::Pubkey;
+use tokio_graceful_shutdown::SubsystemHandle;
+use tracing::info;
+
+use crate::{cache::TaskStateRequest, sync};
+
+pub type TaskStateSender = sync::MessageSender<TaskStateRequest>;
+pub type TaskStateReceiver = sync::MessageReceiver<TaskStateRequest>;
+
+impl TaskStateSender {
+    pub async fn get_in_progress_tasks(&self, pubkey: Pubkey) -> Result<HashSet<u16>, sync::Error> {
+        self.request(|resp| TaskStateRequest::GetInProgressTasks { pubkey, resp })
+            .await
+    }
+
+    pub async fn add_in_progress_tasks(&self, pubkey: Pubkey, task_ids: HashSet<u16>) {
+        self.send(TaskStateRequest::AddInProgressTasks { pubkey, task_ids })
+            .await
+    }
+
+    pub async fn remove_in_progress_tasks(&self, pubkey: Pubkey, task_ids: HashSet<u16>) {
+        self.send(TaskStateRequest::RemoveInProgressTasks { pubkey, task_ids })
+            .await
+    }
+}
+
+pub fn task_state_channel() -> (TaskStateSender, TaskStateReceiver) {
+    let (tx, rx) = sync::message_channel(100);
+    (tx, rx)
+}
+
+pub struct TaskStateCache {
+    cache: HashMap<Pubkey, HashSet<u16>>,
+    receiver: TaskStateReceiver,
+}
+
+impl TaskStateCache {
+    pub fn new(receiver: TaskStateReceiver) -> Self {
+        Self {
+            cache: HashMap::new(),
+            receiver,
+        }
+    }
+
+    pub async fn run(mut self, handle: SubsystemHandle) -> anyhow::Result<()> {
+        info!("starting task state cache");
+        loop {
+            tokio::select! {
+                _ = handle.on_shutdown_requested() => {
+                    info!("shutting down task state cache");
+                    break;
+                }
+                Some(req) = self.receiver.recv() => {
+                    match req {
+                        TaskStateRequest::AddInProgressTasks { pubkey, task_ids } => {
+                            self.cache.entry(pubkey).or_default().extend(task_ids);
+                        }
+                        TaskStateRequest::RemoveInProgressTasks { pubkey, task_ids } => {
+                            if let Some(tasks) = self.cache.get_mut(&pubkey) {
+                                for task_id in task_ids {
+                                    tasks.remove(&task_id);
+                                }
+                                if tasks.is_empty() {
+                                    self.cache.remove(&pubkey);
+                                }
+                            }
+                        }
+                        TaskStateRequest::GetInProgressTasks { pubkey, resp } => {
+                            resp.send(self.cache.get(&pubkey).cloned().unwrap_or_default());
+                        }
+                    }
+                }
+                else => break,
+            }
+        }
+        Ok(())
+    }
+}