add aws kms client caching (#58)

* add aws kms client caching

* don't inject cache into credential in signer abstraction

* add justfile

* refactor cache injection

* script updates

* remove bun declaration

* remove unused code path

* reduce noisy logs

* atomic cache inserts

* increase webhook workers count

Author: d4mr

Cargo.lock

@@ -23,3 +23,4 @@ alloy-signer-aws = { version = "1.0.23", features = ["eip712"] }
 aws-config = "1.8.2"
 aws-sdk-kms = "1.79.0"
 aws-credential-types = "1.2.4"
+moka = { version = "0.12", features = ["future"] }

core/Cargo.toml

@@ -5,17 +5,33 @@ use aws_config::{BehaviorVersion, Region};
 use aws_credential_types::provider::future::ProvideCredentials as ProvideCredentialsFuture;
 use aws_sdk_kms::config::{Credentials, ProvideCredentials};
 use serde::{Deserialize, Serialize};
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+use std::ops::Deref;
 use thirdweb_core::auth::ThirdwebAuth;
 use thirdweb_core::iaw::AuthToken;
 use vault_types::enclave::auth::Auth as VaultAuth;
 
 use crate::error::EngineError;
 
+/// Cache for AWS KMS clients to avoid recreating connections
+pub type KmsClientCache = moka::future::Cache<u64, aws_sdk_kms::Client>;
+
 impl SigningCredential {
     /// Create a random private key credential for testing
     pub fn random_local() -> Self {
         SigningCredential::PrivateKey(PrivateKeySigner::random())
     }
+
+    /// Inject KMS cache into AWS KMS credentials (useful after deserialization)
+    pub fn with_aws_kms_cache(self, kms_client_cache: &KmsClientCache) -> Self {
+        match self {
+            SigningCredential::AwsKms(creds) => {
+                SigningCredential::AwsKms(creds.with_cache(kms_client_cache.clone()))
+            }
+            other => other,
+        }
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -38,6 +54,18 @@ pub struct AwsKmsCredential {
     pub secret_access_key: String,
     pub key_id: String,
     pub region: String,
+    #[serde(skip)]
+    pub kms_client_cache: Option<KmsClientCache>,
+}
+
+impl Hash for AwsKmsCredential {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.access_key_id.hash(state);
+        self.secret_access_key.hash(state);
+        self.key_id.hash(state);
+        self.region.hash(state);
+        // Don't hash the cache - it's not part of the credential identity
+    }
 }
 
 impl ProvideCredentials for AwsKmsCredential {
@@ -57,14 +85,71 @@ impl ProvideCredentials for AwsKmsCredential {
 }
 
 impl AwsKmsCredential {
-    pub async fn get_signer(&self, chain_id: Option<ChainId>) -> Result<AwsSigner, EngineError> {
+    /// Create a new AwsKmsCredential with cache
+    pub fn new(
+        access_key_id: String,
+        secret_access_key: String,
+        key_id: String,
+        region: String,
+        kms_client_cache: KmsClientCache,
+    ) -> Self {
+        Self {
+            access_key_id,
+            secret_access_key,
+            key_id,
+            region,
+            kms_client_cache: Some(kms_client_cache),
+        }
+    }
+
+    /// Inject cache into this credential (useful after deserialization)
+    pub fn with_cache(mut self, kms_client_cache: KmsClientCache) -> Self {
+        self.kms_client_cache = Some(kms_client_cache);
+        self
+    }
+
+    /// Create a cache key from the credential
+    fn cache_key(&self) -> u64 {
+        let mut hasher = DefaultHasher::new();
+        self.hash(&mut hasher);
+        hasher.finish()
+    }
+
+    /// Create a new AWS KMS client (without caching)
+    async fn create_kms_client(&self) -> Result<aws_sdk_kms::Client, EngineError> {
         let config = aws_config::defaults(BehaviorVersion::latest())
             .credentials_provider(self.clone())
             .region(Region::new(self.region.clone()))
             .load()
             .await;
-        let client = aws_sdk_kms::Client::new(&config);
+        Ok(aws_sdk_kms::Client::new(&config))
+    }
 
+    /// Get a cached AWS KMS client, creating one if it doesn't exist
+    async fn get_cached_kms_client(&self) -> Result<aws_sdk_kms::Client, EngineError> {
+        match &self.kms_client_cache {
+            Some(cache) => {
+                let cache_key = self.cache_key();
+
+                cache
+                    .try_get_with(cache_key, async {
+                        tracing::debug!("Creating new KMS client for key: {}", cache_key);
+                        self.create_kms_client().await
+                    })
+                    .await
+                    .map_err(|e| e.deref().clone())
+            }
+            None => {
+                // Fallback to creating a new client without caching
+                tracing::debug!("No cache available, creating new KMS client");
+                self.create_kms_client().await
+            }
+        }
+    }
+
+    /// Get signer (uses cache if available)
+    pub async fn get_signer(&self, chain_id: Option<ChainId>) -> Result<AwsSigner, EngineError> {
+        let client = self.get_cached_kms_client().await?;
         let signer = AwsSigner::new(client, self.key_id.clone(), chain_id).await?;
         Ok(signer)
     }

core/src/credentials.rs

@@ -3,7 +3,7 @@ use alloy::primitives::{Address, U256};
 use alloy::providers::Provider;
 use engine_core::{
     chain::{Chain, ChainService},
-    credentials::SigningCredential,
+    credentials::{SigningCredential, KmsClientCache},
     error::AlloyRpcErrorToEngineError,
     signer::EoaSigner,
 };
@@ -127,6 +127,9 @@ where
 
     // EOA metrics abstraction with encapsulated configuration
     pub eoa_metrics: EoaMetrics,
+    
+    // KMS client cache for AWS KMS credentials
+    pub kms_client_cache: KmsClientCache,
 }
 
 impl<CS> DurableExecution for EoaExecutorJobHandler<CS>
@@ -182,12 +185,15 @@ where
 
         let chain_id = chain.chain_id();
 
+        // Inject KMS cache into the noop signing credential (after deserialization from Redis)
+        let noop_signing_credential = data.noop_signing_credential.clone().with_aws_kms_cache(&self.kms_client_cache);
+
         let worker = EoaExecutorWorker {
             store: scoped,
             chain,
             eoa: data.eoa_address,
             chain_id: data.chain_id,
-            noop_signing_credential: data.noop_signing_credential.clone(),
+            noop_signing_credential,
 
             max_inflight: if is_minimal_account {
                 1
@@ -199,6 +205,7 @@ where
             max_recycled_nonces: self.max_recycled_nonces,
             webhook_queue: self.webhook_queue.clone(),
             signer: self.eoa_signer.clone(),
+            kms_client_cache: self.kms_client_cache.clone(),
         };
 
         let job_start_time = current_timestamp_ms();
@@ -311,6 +318,7 @@ pub struct EoaExecutorWorker<C: Chain> {
 
     pub webhook_queue: Arc<Queue<WebhookJobHandler>>,
     pub signer: Arc<EoaSigner>,
+    pub kms_client_cache: KmsClientCache,
 }
 
 impl<C: Chain> EoaExecutorWorker<C> {

executors/src/eoa/worker/mod.rs

@@ -176,7 +176,7 @@ impl<C: Chain> EoaExecutorWorker<C> {
         // Try EIP-1559 fees first, fall back to legacy if unsupported
         match self.chain.provider().estimate_eip1559_fees().await {
             Ok(eip1559_fees) => {
-                tracing::debug!(
+                tracing::trace!(
                     "Using EIP-1559 fees: max_fee={}, max_priority_fee={}",
                     eip1559_fees.max_fee_per_gas,
                     eip1559_fees.max_priority_fee_per_gas
@@ -397,7 +397,11 @@ impl<C: Chain> EoaExecutorWorker<C> {
         nonce: u64,
     ) -> Result<Signed<TypedTransaction>, EoaExecutorWorkerError> {
         let typed_tx = self.build_typed_transaction(request, nonce).await?;
-        self.sign_transaction(typed_tx, &request.signing_credential)
+        
+        // Inject KMS cache into the signing credential (after deserialization from Redis)
+        let credential_with_cache = request.signing_credential.clone().with_aws_kms_cache(&self.kms_client_cache);
+        
+        self.sign_transaction(typed_tx, &credential_with_cache)
             .await
     }
 

executors/src/eoa/worker/transaction.rs

@@ -0,0 +1,42 @@
+# Start local chain with anvil + speedbump proxy (300ms variable latency)
+local-chain:
+    #!/usr/bin/env bash
+    set -e
+    
+    echo "🧹 Cleaning up existing processes..."
+    # Kill any existing anvil or speedbump processes
+    pkill -f "anvil.*8546" || true
+    pkill -f "speedbump.*8545" || true
+    lsof -ti:8545 | xargs kill -9 2>/dev/null || true
+    lsof -ti:8546 | xargs kill -9 2>/dev/null || true
+    
+    echo "🔨 Starting anvil on port 8546 (1s blocktime)..."
+    anvil --port 8546 --block-time 1 &
+    ANVIL_PID=$!
+    
+    # Cleanup function
+    cleanup() {
+        echo ""
+        echo "🛑 Stopping services..."
+        kill $ANVIL_PID 2>/dev/null || true
+        pkill -f "speedbump.*8545" || true
+        exit 0
+    }
+    
+    trap cleanup INT TERM EXIT
+    
+    # Wait a moment for anvil to start
+    sleep 2
+    
+    echo "🐌 Starting speedbump proxy on port 8545 (→ localhost:8546)"
+    echo "   Latency: 300ms base + 150ms sine wave (150-450ms variable)"
+    echo "   Connect to: http://localhost:8545"
+    echo ""
+    speedbump --port=8545 --latency=300ms --sine-amplitude=150ms --sine-period=1m localhost:8546
+
+# Fund an address with 1 ETH (bypasses speedbump for faster setup)
+fund address:
+    @echo "💰 Funding {{address}} with 1 ETH..."
+    @cast rpc anvil_setBalance {{address}} $(cast to-wei 1) --rpc-url http://localhost:8546
+    @echo "✅ Done!"
+

justfile

@@ -1,14 +1,3 @@
-/// <reference lib="dom" />
-
-// Bun globals (runtime will provide these)
-declare const Bun: any;
-declare const process: any;
-
-// Extend ImportMeta for Bun
-interface ImportMeta {
-  dir: string;
-}
-
 // Types based on events.rs
 interface WebhookEvent {
   transactionId: string;