refactor: simplify sync to async sqlx bridge (#354)

Also * reduce amount of threads and use multi-threaded tokio runtime. * close the connection when encrypting the db before copying it.
boxdot · Feb 20, 2025 · 1420648 · 1420648
1 parent e92d191
commit 1420648
Show file tree

Hide file tree

Showing 11 changed files with 269 additions and 380 deletions.
diff --git a/...73bc217fad6419b52bcad53b899de7fe6cf4.json → ...7e00c0c5da86af32c5894a41a40158e4210b.json b/...73bc217fad6419b52bcad53b899de7fe6cf4.json → ...7e00c0c5da86af32c5894a41a40158e4210b.json
diff --git a/...17a6c50e61a2b17e6983f330c339cb23d442.json → ...c88fd285a20c5c23a360a8bba9e5b18096b7.json b/...17a6c50e61a2b17e6983f330c339cb23d442.json → ...c88fd285a20c5c23a360a8bba9e5b18096b7.json
diff --git a/...0cfea8098a9c509887d20a0706bd2f339ead.json → ...d5bf0fa7b9ee6891fabde769defb9f6e38d3.json b/...0cfea8098a9c509887d20a0706bd2f339ead.json → ...d5bf0fa7b9ee6891fabde769defb9f6e38d3.json
diff --git a/...74e2c33824667c9a6fe9370456d659798b34.json → ...efb50b96d40a77d5d4277b31e21f2e572d5a.json b/...74e2c33824667c9a6fe9370456d659798b34.json → ...efb50b96d40a77d5d4277b31e21f2e572d5a.json
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -61,7 +61,6 @@ phonenumber = "0.3.6"
 postcard = { version = "1.1.1", features = ["alloc"] }
 qr2term = "0.3.3"
 ratatui = "0.29.0"
-rayon = "1.10.0"
 regex = "1.11.1"
 scopeguard = "1.2.0"
 serde = { version = "1.0.216", features = ["derive"] }
@@ -74,7 +73,6 @@ sqlx = { version = "0.8.2", features = [
 ] }
 textwrap = "0.16.1"
 thiserror = "2.0.9"
-thread_local = "1.1.8"
 tokio = { version = "1.42.0", default-features = false, features = [
     "rt-multi-thread",
     "macros",
@@ -94,6 +92,7 @@ tempfile = "3.14.0"
 crokey = "1.1.0"
 strum_macros = "0.26.4"
 strum = { version = "0.26.3", features = ["derive"] }
+tokio-util = { version = "0.7.13", features = ["rt"] }
 qrcode = { version = "0.14.1", default-features = false, features = ["image"] }
 
 [package.metadata.cargo-machete]

diff --git a/src/main.rs b/src/main.rs
@@ -23,6 +23,7 @@ use presage::libsignal_service::content::Content;
 use ratatui::{backend::CrosstermBackend, Terminal};
 use tokio::select;
 use tokio_stream::StreamExt;
+use tokio_util::task::LocalPoolHandle;
 use tracing::debug;
 use tracing::{error, info};
 
@@ -43,7 +44,7 @@ struct Args {
     relink: bool,
 }
 
-#[tokio::main(flavor = "current_thread")]
+#[tokio::main(worker_threads = 2)]
 async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
@@ -62,9 +63,7 @@ async fn main() -> anyhow::Result<()> {
 
     log_panics::init();
 
-    tokio::task::LocalSet::new()
-        .run_until(run_single_threaded(args.relink))
-        .await
+    run(args.relink).await
 }
 
 async fn is_online() -> bool {
@@ -88,8 +87,10 @@ pub enum Event {
     AppEvent(gurk::event::Event),
 }
 
-async fn run_single_threaded(relink: bool) -> anyhow::Result<()> {
-    let (mut signal_manager, config) = signal::ensure_linked_device(relink).await?;
+async fn run(relink: bool) -> anyhow::Result<()> {
+    let local_pool = LocalPoolHandle::new(2);
+    let (mut signal_manager, config) =
+        signal::ensure_linked_device(relink, local_pool.clone()).await?;
 
     let mut storage: Box<dyn Storage> = {
         debug!(
@@ -102,6 +103,7 @@ async fn run_single_threaded(relink: bool) -> anyhow::Result<()> {
             config.passphrase.clone(),
             config.sqlite.preserve_unencrypted,
         )
+        .await
         .with_context(|| {
             format!(
                 "failed to open sqlite data storage at: {}",
@@ -153,7 +155,8 @@ async fn run_single_threaded(relink: bool) -> anyhow::Result<()> {
     });
 
     let inner_tx = tx.clone();
-    tokio::task::spawn_local(async move {
+
+    local_pool.spawn_pinned(|| async move {
         let mut backoff = Backoff::new();
         loop {
             let mut messages = if !is_online().await {
@@ -168,7 +171,8 @@ async fn run_single_threaded(relink: bool) -> anyhow::Result<()> {
                     Err(e) => {
                         let e = e.context(
                             "failed to initialize the stream of Signal messages.\n\
-                            Maybe the device was unlinked? Please try to restart with '--relink` flag.",
+                            Maybe the device was unlinked? Please try to restart with \
+                            '--relink` flag.",
                         );
                         inner_tx
                             .send(Event::Quit(Some(e)))

diff --git a/src/signal/impl.rs b/src/signal/impl.rs
@@ -20,6 +20,7 @@ use presage::{
 use presage_store_sled::SledStore;
 use tokio::sync::oneshot;
 use tokio_stream::{Stream, StreamExt};
+use tokio_util::task::LocalPoolHandle;
 use tracing::{error, warn};
 use uuid::Uuid;
 
@@ -33,18 +34,25 @@ use super::{
 
 pub(super) struct PresageManager {
     manager: presage::Manager<SledStore, Registered>,
+    local_pool: LocalPoolHandle,
 }
 
 impl PresageManager {
-    pub(super) fn new(manager: presage::Manager<SledStore, Registered>) -> Self {
-        Self { manager }
+    pub(super) fn new(
+        manager: presage::Manager<SledStore, Registered>,
+        local_pool: LocalPoolHandle,
+    ) -> Self {
+        Self {
+            manager,
+            local_pool,
+        }
     }
 }
 
 #[async_trait(?Send)]
 impl SignalManager for PresageManager {
     fn clone_boxed(&self) -> Box<dyn SignalManager> {
-        Box::new(Self::new(self.manager.clone()))
+        Box::new(Self::new(self.manager.clone(), self.local_pool.clone()))
     }
 
     fn user_id(&self) -> Uuid {
@@ -102,7 +110,7 @@ impl SignalManager for PresageManager {
         };
 
         let mut manager = self.manager.clone();
-        tokio::task::spawn_local(async move {
+        self.local_pool.spawn_pinned(move || async move {
             let body = ContentBody::ReceiptMessage(data_message);
             if let Err(error) = manager
                 .send_message(ServiceId::Aci(sender_uuid.into()), body, now_timestamp)
@@ -162,7 +170,7 @@ impl SignalManager for PresageManager {
         match channel.id {
             ChannelId::User(uuid) => {
                 let mut manager = self.manager.clone();
-                tokio::task::spawn_local(async move {
+                self.local_pool.spawn_pinned(move || async move {
                     if let Err(error) =
                         upload_attachments(&manager, attachments, &mut data_message).await
                     {
@@ -202,7 +210,7 @@ impl SignalManager for PresageManager {
                         ..Default::default()
                     });
 
-                    tokio::task::spawn_local(async move {
+                    self.local_pool.spawn_pinned(move || async move {
                         if let Err(error) =
                             upload_attachments(&manager, attachments, &mut data_message).await
                         {
@@ -271,7 +279,7 @@ impl SignalManager for PresageManager {
             (ChannelId::User(uuid), _) => {
                 let mut manager = self.manager.clone();
                 let body = ContentBody::DataMessage(data_message);
-                tokio::task::spawn_local(async move {
+                self.local_pool.spawn_pinned(move || async move {
                     if let Err(e) = manager
                         .send_message(ServiceId::Aci(uuid.into()), body, timestamp)
                         .await
@@ -291,7 +299,7 @@ impl SignalManager for PresageManager {
                     ..Default::default()
                 });
 
-                tokio::task::spawn_local(async move {
+                self.local_pool.spawn_pinned(move || async move {
                     if let Err(e) = manager
                         .send_message_to_group(&master_key_bytes, data_message, timestamp)
                         .await

diff --git a/src/signal/mod.rs b/src/signal/mod.rs
@@ -10,6 +10,7 @@ use futures_channel::oneshot;
 use image::Luma;
 use presage::{libsignal_service::configuration::SignalServers, model::identity::OnNewIdentity};
 use presage_store_sled::{MigrationConflictStrategy, SledStore};
+use tokio_util::task::LocalPoolHandle;
 use tracing::{error, warn};
 use url::Url;
 
@@ -37,7 +38,8 @@ pub type GroupIdentifierBytes = [u8; GROUP_IDENTIFIER_LEN];
 ///    path.
 pub async fn ensure_linked_device(
     relink: bool,
-) -> anyhow::Result<(Box<dyn SignalManager>, Config)> {
+    local_pool: LocalPoolHandle,
+) -> anyhow::Result<(Box<dyn SignalManager + Send>, Config)> {
     let config = Config::load_installed()?;
 
     // warn about deprecated keys
@@ -74,7 +76,7 @@ pub async fn ensure_linked_device(
             match presage::Manager::load_registered(store.clone()).await {
                 Ok(manager) => {
                     // done loading manager from store
-                    return Ok((Box::new(PresageManager::new(manager)), config));
+                    return Ok((Box::new(PresageManager::new(manager, local_pool)), config));
                 }
                 Err(e) => {
                     bail!("error loading manager. Try again later or run with --relink to force relink: {}", e)
@@ -147,7 +149,7 @@ pub async fn ensure_linked_device(
         config
     };
 
-    Ok((Box::new(PresageManager::new(manager)), config))
+    Ok((Box::new(PresageManager::new(manager, local_pool)), config))
 }
 
 async fn gen_qr_code(rx: oneshot::Receiver<Url>, path: &Path) -> anyhow::Result<()> {