From 6a6f2d322102c71a13f3fab5b9b1ef406af1fbbc Mon Sep 17 00:00:00 2001
From: Diana Carvalho <diana@propellerheads.xyz>
Date: Thu, 27 Feb 2025 11:16:21 +0000
Subject: [PATCH] fix: Get correct runtime everywhere

Made an utils function for it to make sure it is used correctly

--- don't change below this line ---
ENG-4260 Took 9 minutes
---
 src/encoding/evm/approvals/permit2.rs         | 12 ++-----
 .../approvals/protocol_approvals_manager.rs   | 32 +++++++++++++------
 src/encoding/evm/utils.rs                     | 15 ++++++++-
 3 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/src/encoding/evm/approvals/permit2.rs b/src/encoding/evm/approvals/permit2.rs
index 0374afd..9eb6c3e 100644
--- a/src/encoding/evm/approvals/permit2.rs
+++ b/src/encoding/evm/approvals/permit2.rs
@@ -21,7 +21,7 @@ use crate::encoding::{
     errors::EncodingError,
     evm::{
         approvals::protocol_approvals_manager::get_client,
-        utils::{biguint_to_u256, bytes_to_address, encode_input},
+        utils::{biguint_to_u256, bytes_to_address, encode_input, get_runtime},
     },
     models::Chain,
 };
@@ -70,15 +70,7 @@ sol! {
 
 impl Permit2 {
     pub fn new(swapper_pk: String, chain: Chain) -> Result<Self, EncodingError> {
-        let (handle, runtime) = match Handle::try_current() {
-            Ok(h) => (h, None),
-            Err(_) => {
-                let rt = Arc::new(Runtime::new().map_err(|_| {
-                    EncodingError::FatalError("Failed to create a new tokio runtime".to_string())
-                })?);
-                (rt.handle().clone(), Some(rt))
-            }
-        };
+        let (handle, runtime) = get_runtime()?;
         let client = block_in_place(|| handle.block_on(get_client()))?;
         let pk = B256::from_str(&swapper_pk).map_err(|_| {
             EncodingError::FatalError("Failed to convert swapper private key to B256".to_string())
diff --git a/src/encoding/evm/approvals/protocol_approvals_manager.rs b/src/encoding/evm/approvals/protocol_approvals_manager.rs
index 9af2355..d239f73 100644
--- a/src/encoding/evm/approvals/protocol_approvals_manager.rs
+++ b/src/encoding/evm/approvals/protocol_approvals_manager.rs
@@ -8,21 +8,32 @@ use alloy::{
 use alloy_primitives::{Address, Bytes, TxKind, U256};
 use alloy_sol_types::SolValue;
 use dotenv::dotenv;
-use tokio::runtime::Runtime;
+use tokio::{
+    runtime::{Handle, Runtime},
+    task::block_in_place,
+};
 
-use crate::encoding::{errors::EncodingError, evm::utils::encode_input};
+use crate::encoding::{
+    errors::EncodingError,
+    evm::utils::{encode_input, get_runtime},
+};
 
 /// A manager for checking if an approval is needed for interacting with a certain spender.
 pub struct ProtocolApprovalsManager {
     client: Arc<RootProvider<BoxTransport>>,
-    runtime: Runtime,
+    runtime_handle: Handle,
+    // Store the runtime to prevent it from being dropped before use.
+    // This is required since tycho-execution does not have a pre-existing runtime.
+    // However, if the library is used in a context where a runtime already exists, it is not
+    // necessary to store it.
+    #[allow(dead_code)]
+    runtime: Option<Arc<Runtime>>,
 }
 impl ProtocolApprovalsManager {
     pub fn new() -> Result<Self, EncodingError> {
-        let runtime = Runtime::new()
-            .map_err(|_| EncodingError::FatalError("Failed to create runtime".to_string()))?;
-        let client = runtime.block_on(get_client())?;
-        Ok(Self { client, runtime })
+        let (handle, runtime) = get_runtime()?;
+        let client = block_in_place(|| handle.block_on(get_client()))?;
+        Ok(Self { client, runtime_handle: handle, runtime })
     }
 
     /// Checks the current allowance for the given token, owner, and spender, and returns true
@@ -41,9 +52,10 @@ impl ProtocolApprovalsManager {
             ..Default::default()
         };
 
-        let output = self
-            .runtime
-            .block_on(async { self.client.call(&tx).await });
+        let output = block_in_place(|| {
+            self.runtime_handle
+                .block_on(async { self.client.call(&tx).await })
+        });
         match output {
             Ok(response) => {
                 let allowance: U256 = U256::abi_decode(&response, true).map_err(|_| {
diff --git a/src/encoding/evm/utils.rs b/src/encoding/evm/utils.rs
index a59abff..96e04ea 100644
--- a/src/encoding/evm/utils.rs
+++ b/src/encoding/evm/utils.rs
@@ -1,7 +1,8 @@
-use std::cmp::max;
+use std::{cmp::max, sync::Arc};
 
 use alloy_primitives::{aliases::U24, keccak256, Address, FixedBytes, Keccak256, U256, U8};
 use num_bigint::BigUint;
+use tokio::runtime::{Handle, Runtime};
 use tycho_core::Bytes;
 
 use crate::encoding::{
@@ -120,3 +121,15 @@ pub fn get_static_attribute(swap: &Swap, attribute_name: &str) -> Result<Vec<u8>
         })?
         .to_vec())
 }
+
+pub fn get_runtime() -> Result<(Handle, Option<Arc<Runtime>>), EncodingError> {
+    match Handle::try_current() {
+        Ok(h) => Ok((h, None)),
+        Err(_) => {
+            let rt = Arc::new(Runtime::new().map_err(|_| {
+                EncodingError::FatalError("Failed to create a new tokio runtime".to_string())
+            })?);
+            Ok((rt.handle().clone(), Some(rt)))
+        }
+    }
+}