Cache signer address

MystenLabs · May 25, 2024 · 69a187e · 69a187e
1 parent e28feec
commit 69a187e
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 27 deletions.
diff --git a/src/command.rs b/src/command.rs
@@ -51,9 +51,9 @@ impl Command {
         let _guard = telemetry_config.init();
         info!("Metrics server started at {:?}", metric_address);
 
-        let signer = signer_config.new_signer();
+        let signer = signer_config.new_signer().await;
         let storage_metrics = StorageMetrics::new(&prometheus_registry);
-        let sponsor_address = signer.get_address().await.unwrap();
+        let sponsor_address = signer.get_address();
         info!("Sponsor address: {:?}", sponsor_address);
         let storage = connect_storage(&gas_pool_config, sponsor_address, storage_metrics).await;
         let sui_client = SuiClient::new(&fullnode_url, fullnode_basic_auth).await;

diff --git a/src/config.rs b/src/config.rs
@@ -94,10 +94,10 @@ impl Default for TxSignerConfig {
 }
 
 impl TxSignerConfig {
-    pub fn new_signer(self) -> Arc<dyn TxSigner> {
+    pub async fn new_signer(self) -> Arc<dyn TxSigner> {
         match self {
             TxSignerConfig::Local { keypair } => TestTxSigner::new(keypair),
-            TxSignerConfig::Sidecar { sidecar_url } => SidecarTxSigner::new(sidecar_url),
+            TxSignerConfig::Sidecar { sidecar_url } => SidecarTxSigner::new(sidecar_url).await,
         }
     }
 }

diff --git a/src/gas_pool/gas_pool_core.rs b/src/gas_pool/gas_pool_core.rs
@@ -65,7 +65,7 @@ impl GasPool {
         duration: Duration,
     ) -> anyhow::Result<(SuiAddress, ReservationID, Vec<ObjectRef>)> {
         self.gas_usage_cap.check_usage().await?;
-        let sponsor = self.signer.get_address().await?;
+        let sponsor = self.signer.get_address();
         let (reservation_id, gas_coins) = self
             .gas_pool_store
             .reserve_gas_coins(gas_budget, duration.as_millis() as u64)
@@ -87,7 +87,7 @@ impl GasPool {
         user_sig: GenericSignature,
     ) -> anyhow::Result<SuiTransactionBlockEffects> {
         let sponsor = tx_data.gas_data().owner;
-        if !self.signer.is_valid_address(&sponsor).await? {
+        if !self.signer.is_valid_address(&sponsor) {
             bail!("Sponsor {:?} is not registered", sponsor);
         };
         Self::check_transaction_validity(&tx_data)?;

diff --git a/src/gas_pool_initializer.rs b/src/gas_pool_initializer.rs
@@ -238,7 +238,7 @@ impl GasPoolInitializer {
         target_init_coin_balance: u64,
         signer: &Arc<dyn TxSigner>,
     ) {
-        let sponsor_address = signer.get_address().await.unwrap();
+        let sponsor_address = signer.get_address();
         if storage
             .acquire_init_lock(MAX_INIT_DURATION_SEC)
             .await
@@ -343,7 +343,7 @@ mod tests {
         telemetry_subscribers::init_for_testing();
         let (cluster, signer) = start_sui_cluster(vec![1000 * MIST_PER_SUI]).await;
         let fullnode_url = cluster.fullnode_handle.rpc_url;
-        let storage = connect_storage_for_testing(signer.get_address().await.unwrap()).await;
+        let storage = connect_storage_for_testing(signer.get_address()).await;
         let sui_client = SuiClient::new(&fullnode_url, None).await;
         let _ = GasPoolInitializer::start(
             sui_client,
@@ -363,7 +363,7 @@ mod tests {
         telemetry_subscribers::init_for_testing();
         let (cluster, signer) = start_sui_cluster(vec![10000000 * MIST_PER_SUI]).await;
         let fullnode_url = cluster.fullnode_handle.rpc_url;
-        let storage = connect_storage_for_testing(signer.get_address().await.unwrap()).await;
+        let storage = connect_storage_for_testing(signer.get_address()).await;
         let target_init_balance = 12345 * MIST_PER_SUI;
         let sui_client = SuiClient::new(&fullnode_url, None).await;
         let _ = GasPoolInitializer::start(
@@ -383,9 +383,9 @@ mod tests {
     async fn test_add_new_funds_to_pool() {
         telemetry_subscribers::init_for_testing();
         let (cluster, signer) = start_sui_cluster(vec![1000 * MIST_PER_SUI]).await;
-        let sponsor = signer.get_address().await.unwrap();
+        let sponsor = signer.get_address();
         let fullnode_url = cluster.fullnode_handle.rpc_url.clone();
-        let storage = connect_storage_for_testing(signer.get_address().await.unwrap()).await;
+        let storage = connect_storage_for_testing(signer.get_address()).await;
         let sui_client = SuiClient::new(&fullnode_url, None).await;
         let _init_task = GasPoolInitializer::start(
             sui_client,

diff --git a/src/test_env.rs b/src/test_env.rs
@@ -47,7 +47,7 @@ pub async fn start_gas_station(
     debug!("Starting Sui cluster..");
     let (test_cluster, signer) = start_sui_cluster(init_gas_amounts).await;
     let fullnode_url = test_cluster.fullnode_handle.rpc_url.clone();
-    let sponsor_address = signer.get_address().await.unwrap();
+    let sponsor_address = signer.get_address();
     debug!("Starting storage. Sponsor address: {:?}", sponsor_address);
     let storage = connect_storage_for_testing(sponsor_address).await;
     let sui_client = SuiClient::new(&fullnode_url, None).await;

diff --git a/src/tx_signer.rs b/src/tx_signer.rs
@@ -18,9 +18,9 @@ use sui_types::transaction::TransactionData;
 pub trait TxSigner: Send + Sync {
     async fn sign_transaction(&self, tx_data: &TransactionData)
         -> anyhow::Result<GenericSignature>;
-    async fn get_address(&self) -> anyhow::Result<SuiAddress>;
-    async fn is_valid_address(&self, address: &SuiAddress) -> anyhow::Result<bool> {
-        Ok(self.get_address().await? == *address)
+    fn get_address(&self) -> SuiAddress;
+    fn is_valid_address(&self, address: &SuiAddress) -> bool {
+        self.get_address() == *address
     }
 }
 
@@ -39,13 +39,26 @@ struct SuiAddressResponse {
 pub struct SidecarTxSigner {
     sidecar_url: String,
     client: Client,
+    sui_address: SuiAddress,
 }
 
 impl SidecarTxSigner {
-    pub fn new(sidecar_url: String) -> Arc<Self> {
+    pub async fn new(sidecar_url: String) -> Arc<Self> {
+        let client = Client::new();
+        let resp = client
+            .get(format!("{}/{}", sidecar_url, "get-pubkey-address"))
+            .send()
+            .await
+            .unwrap_or_else(|err| panic!("Failed to get pubkey address: {}", err));
+        let sui_address = resp
+            .json::<SuiAddressResponse>()
+            .await
+            .unwrap_or_else(|err| panic!("Failed to parse address response: {}", err))
+            .sui_pubkey_address;
         Arc::new(Self {
             sidecar_url,
-            client: Client::new(),
+            client,
+            sui_address,
         })
     }
 }
@@ -70,14 +83,8 @@ impl TxSigner for SidecarTxSigner {
         Ok(sig)
     }
 
-    async fn get_address(&self) -> anyhow::Result<SuiAddress> {
-        let resp = self
-            .client
-            .get(format!("{}/{}", self.sidecar_url, "get-pubkey-address"))
-            .send()
-            .await?;
-        let address = resp.json::<SuiAddressResponse>().await?;
-        Ok(address.sui_pubkey_address)
+    fn get_address(&self) -> SuiAddress {
+        self.sui_address
     }
 }
 
@@ -102,7 +109,7 @@ impl TxSigner for TestTxSigner {
         Ok(sponsor_sig)
     }
 
-    async fn get_address(&self) -> anyhow::Result<SuiAddress> {
-        Ok((&self.keypair.public()).into())
+    fn get_address(&self) -> SuiAddress {
+        (&self.keypair.public()).into()
     }
 }