From d6362661b609fd9a1eae6dd038ac8bdeb8125469 Mon Sep 17 00:00:00 2001
From: Oscar Beaumont <oscar@otbeaumont.me>
Date: Sun, 17 Mar 2024 11:46:18 +0800
Subject: [PATCH] Fix #260

Co-authored-by: campbellcole <10430178+campbellcole@users.noreply.github.com>
---
 src/integrations/tauri.rs | 31 +++++++++++++++++++------------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/integrations/tauri.rs b/src/integrations/tauri.rs
index 4f6458d8..4e88db64 100644
--- a/src/integrations/tauri.rs
+++ b/src/integrations/tauri.rs
@@ -1,10 +1,10 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{borrow::Borrow, collections::HashMap, sync::Arc};
 
 use tauri::{
     plugin::{Builder, TauriPlugin},
     Manager, Runtime,
 };
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, Mutex};
 
 use crate::{
     internal::jsonrpc::{self, handle_json_rpc, Sender, SubscriptionMap},
@@ -22,19 +22,26 @@ where
     Builder::new("rspc")
         .setup(|app_handle| {
             let (tx, mut rx) = mpsc::unbounded_channel::<jsonrpc::Request>();
-            let (mut resp_tx, mut resp_rx) = mpsc::unbounded_channel::<jsonrpc::Response>();
-            let mut subscriptions = HashMap::new();
+            let (resp_tx, mut resp_rx) = mpsc::unbounded_channel::<jsonrpc::Response>();
+            // TODO: Don't keep using a tokio mutex. We don't need to hold it over the await point.
+            let subscriptions = Arc::new(Mutex::new(HashMap::new()));
 
             tokio::spawn(async move {
                 while let Some(req) = rx.recv().await {
-                    handle_json_rpc(
-                        ctx_fn(),
-                        req,
-                        &router,
-                        &mut Sender::ResponseChannel(&mut resp_tx),
-                        &mut SubscriptionMap::Ref(&mut subscriptions),
-                    )
-                    .await;
+                    let ctx = ctx_fn();
+                    let router = router.clone();
+                    let mut resp_tx = resp_tx.clone();
+                    let subscriptions = subscriptions.clone();
+                    tokio::spawn(async move {
+                        handle_json_rpc(
+                            ctx,
+                            req,
+                            &router,
+                            &mut Sender::ResponseChannel(&mut resp_tx),
+                            &mut SubscriptionMap::Mutex(subscriptions.borrow()),
+                        )
+                        .await;
+                    });
                 }
             });