From bcfa2c411fb4b4b09b7d69982ab40bfb58819a8a Mon Sep 17 00:00:00 2001
From: Gregory Schier <gschier1990@gmail.com>
Date: Tue, 6 Feb 2024 23:26:24 -0800
Subject: [PATCH] Fix pool management

---
 src-tauri/grpc/src/manager.rs | 16 +++++++++++-----
 src-tauri/grpc/src/proto.rs   |  2 ++
 src-tauri/src/main.rs         |  7 ++++---
 src-web/hooks/useGrpc.ts      |  3 ++-
 src-web/hooks/useHotKey.ts    |  8 ++++----
 5 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/src-tauri/grpc/src/manager.rs b/src-tauri/grpc/src/manager.rs
index 49812f3a..614cff3c 100644
--- a/src-tauri/grpc/src/manager.rs
+++ b/src-tauri/grpc/src/manager.rs
@@ -150,34 +150,41 @@ impl GrpcConnection {
 }
 
 pub struct GrpcHandle {
-    connections: HashMap<String, GrpcConnection>,
     pools: HashMap<String, DescriptorPool>,
 }
 
 impl Default for GrpcHandle {
     fn default() -> Self {
-        let connections = HashMap::new();
         let pools = HashMap::new();
-        Self { connections, pools }
+        Self { pools }
     }
 }
 
 impl GrpcHandle {
     pub async fn services_from_files(
-        &self,
+        &mut self,
+        id: &str,
+        uri: &Uri,
         paths: Vec<PathBuf>,
     ) -> Result<Vec<ServiceDefinition>, String> {
         let pool = fill_pool_from_files(paths).await?;
+        self.pools.insert(self.get_pool_key(id, uri), pool.clone());
         Ok(self.services_from_pool(&pool))
     }
     pub async fn services_from_reflection(
         &mut self,
+        id: &str,
         uri: &Uri,
     ) -> Result<Vec<ServiceDefinition>, String> {
         let pool = fill_pool(uri).await?;
+        self.pools.insert(self.get_pool_key(id, uri), pool.clone());
         Ok(self.services_from_pool(&pool))
     }
 
+    fn get_pool_key(&self, id: &str, uri: &Uri) -> String {
+        format!("{}-{}", id, uri)
+    }
+
     fn services_from_pool(&self, pool: &DescriptorPool) -> Vec<ServiceDefinition> {
         pool.services()
             .map(|s| {
@@ -268,7 +275,6 @@ impl GrpcHandle {
 
         let conn = get_transport();
         let connection = GrpcConnection { pool, conn, uri };
-        self.connections.insert(id.to_string(), connection.clone());
         Ok(connection)
     }
 }
diff --git a/src-tauri/grpc/src/proto.rs b/src-tauri/grpc/src/proto.rs
index 3b26d598..00f802d2 100644
--- a/src-tauri/grpc/src/proto.rs
+++ b/src-tauri/grpc/src/proto.rs
@@ -24,6 +24,7 @@ use tonic_reflection::pb::server_reflection_response::MessageResponse;
 use tonic_reflection::pb::ServerReflectionRequest;
 
 pub async fn fill_pool_from_files(paths: Vec<PathBuf>) -> Result<DescriptorPool, String> {
+    println!("FILL POOL FROM FILES");
     let mut pool = DescriptorPool::new();
     let random_file_name = format!("{}.desc", uuid::Uuid::new_v4());
     let desc_path = temp_dir().join(random_file_name);
@@ -88,6 +89,7 @@ pub fn get_transport() -> Client<HttpsConnector<HttpConnector>, BoxBody> {
 }
 
 pub async fn fill_pool(uri: &Uri) -> Result<DescriptorPool, String> {
+    println!("FILL POOL FROM URI");
     let mut pool = DescriptorPool::new();
     let mut client = ServerReflectionClient::with_origin(get_transport(), uri.clone());
 
diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs
index e0a169e6..699959e9 100644
--- a/src-tauri/src/main.rs
+++ b/src-tauri/src/main.rs
@@ -101,12 +101,15 @@ async fn cmd_grpc_reflect(
     let req = get_grpc_request(&app_handle, request_id)
         .await
         .map_err(|e| e.to_string())?;
+    let uri = safe_uri(&req.url).map_err(|e| e.to_string())?;
     if req.proto_files.0.len() > 0 {
         println!("REFLECT FROM FILES");
         grpc_handle
             .lock()
             .await
             .services_from_files(
+                &req.id,
+                &uri,
                 req.proto_files
                     .0
                     .iter()
@@ -115,12 +118,10 @@ async fn cmd_grpc_reflect(
             )
             .await
     } else {
-        println!("REFLECT FROM URI");
-        let uri = safe_uri(&req.url).map_err(|e| e.to_string())?;
         grpc_handle
             .lock()
             .await
-            .services_from_reflection(&uri)
+            .services_from_reflection(&req.id, &uri)
             .await
     }
 }
diff --git a/src-web/hooks/useGrpc.ts b/src-web/hooks/useGrpc.ts
index 796e0a76..de4ab684 100644
--- a/src-web/hooks/useGrpc.ts
+++ b/src-web/hooks/useGrpc.ts
@@ -54,9 +54,10 @@ export function useGrpc(req: GrpcRequest | null, conn: GrpcConnection | null) {
   const debouncedUrl = useDebouncedValue<string>(req?.url ?? 'n/a', 1000);
   const reflect = useQuery<ReflectResponseService[] | null, string>({
     enabled: req != null,
-    queryKey: ['grpc_reflect', debouncedUrl],
+    queryKey: ['grpc_reflect', req?.id ?? 'n/a', debouncedUrl],
     refetchOnWindowFocus: false,
     queryFn: async () => {
+      console.log('useGrpc.reflect', { requestId });
       return (await minPromiseMillis(
         invoke('cmd_grpc_reflect', { requestId }),
         300,
diff --git a/src-web/hooks/useHotKey.ts b/src-web/hooks/useHotKey.ts
index 67637624..e7de7261 100644
--- a/src-web/hooks/useHotKey.ts
+++ b/src-web/hooks/useHotKey.ts
@@ -112,11 +112,11 @@ export function useAnyHotkey(
       }
       currentKeys.current.delete(normalizeKey(e.key, os));
     };
-    window.addEventListener('keydown', down);
-    window.addEventListener('keyup', up);
+    document.addEventListener('keydown', down);
+    document.addEventListener('keyup', up);
     return () => {
-      window.removeEventListener('keydown', down);
-      window.removeEventListener('keyup', up);
+      document.removeEventListener('keydown', down);
+      document.removeEventListener('keyup', up);
     };
   }, [options.enable, os]);
 }