Funding cell exclusion

src/ckb/actor.rs

@@ -7,7 +7,10 @@ use ractor::{
 
 use crate::ckb::contracts::{get_script_by_contract, Contract};
 
-use super::{funding::FundingContext, CkbConfig, FundingError, FundingRequest, FundingTx};
+use super::{
+    funding::{FundingContext, FundingExclusion},
+    CkbConfig, FundingError, FundingRequest, FundingTx,
+};
 
 pub struct CkbChainActor {}
 
@@ -16,6 +19,7 @@ pub struct CkbChainState {
     config: CkbConfig,
     secret_key: secp256k1::SecretKey,
     funding_source_lock_script: packed::Script,
+    funding_exclusion: FundingExclusion,
 }
 
 #[derive(Debug, Clone)]
@@ -27,12 +31,19 @@ pub struct TraceTxRequest {
 
 #[derive(Debug)]
 pub enum CkbChainMessage {
+    // Funding management
     Fund(
         FundingTx,
         FundingRequest,
         RpcReplyPort<Result<FundingTx, FundingError>>,
     ),
     Sign(FundingTx, RpcReplyPort<Result<FundingTx, FundingError>>),
+
+    // AddTxs and RemoveTx are used to manage the funding exclusion list.
+    AddTxs(Vec<FundingTx>),
+    RemoveTx(packed::Byte32),
+
+    // Interacts with CKB chain
     SendTx(TransactionView, RpcReplyPort<Result<(), RpcError>>),
     TraceTx(TraceTxRequest, RpcReplyPort<TraceTxResponse>),
     GetCurrentBlockNumber((), RpcReplyPort<Result<u64, RpcError>>),
@@ -80,6 +91,7 @@ impl Actor for CkbChainActor {
             config,
             secret_key,
             funding_source_lock_script,
+            funding_exclusion: Default::default(),
         })
     }
 
@@ -89,7 +101,9 @@ impl Actor for CkbChainActor {
         message: Self::Msg,
         state: &mut Self::State,
     ) -> Result<(), ActorProcessingErr> {
-        use CkbChainMessage::{Fund, GetCurrentBlockNumber, SendTx, Sign, TraceTx};
+        use CkbChainMessage::{
+            AddTxs, Fund, GetCurrentBlockNumber, RemoveTx, SendTx, Sign, TraceTx,
+        };
         match message {
             GetCurrentBlockNumber(_, reply) => {
                 // Have to use block_in_place here, see https://github.com/seanmonstar/reqwest/issues/1017.
@@ -104,7 +118,7 @@ impl Actor for CkbChainActor {
                 let context = state.build_funding_context(&request);
                 if !reply_port.is_closed() {
                     tokio::task::block_in_place(move || {
-                        let result = tx.fulfill(request, context);
+                        let result = tx.fulfill(request, context, &mut state.funding_exclusion);
                         if !reply_port.is_closed() {
                             // ignore error
                             let _ = reply_port.send(result);
@@ -125,6 +139,16 @@ impl Actor for CkbChainActor {
                     });
                 }
             }
+
+            AddTxs(txs) => {
+                for tx in txs {
+                    state.funding_exclusion.insert(tx);
+                }
+            }
+            RemoveTx(tx_hash) => {
+                state.funding_exclusion.remove(&tx_hash);
+            }
+
             SendTx(tx, reply_port) => {
                 let rpc_url = state.config.rpc_url.clone();
                 tokio::task::block_in_place(move || {

src/ckb/error.rs

@@ -15,6 +15,9 @@ pub enum FundingError {
     #[error("Failed to sign CKB tx: {0}")]
     CkbTxUnlockError(#[from] UnlockError),
 
+    #[error("Failed to manage live cells exclusion list")]
+    FundingExclusionError,
+
     #[error("Dead cell found in the tx")]
     DeadCell,
 

src/ckb/funding/funding_exclusion.rs

@@ -0,0 +1,40 @@
+use std::collections::HashMap;
+
+use ckb_sdk::traits::CellCollector;
+use ckb_types::{core::TransactionView, packed::Byte32};
+use tracing::error;
+
+use crate::ckb::FundingError;
+
+use super::FundingTx;
+
+#[derive(Default, Clone, Debug)]
+pub struct FundingExclusion {
+    pending_funding_txs: HashMap<Byte32, TransactionView>,
+}
+
+impl FundingExclusion {
+    pub fn insert(&mut self, tx: FundingTx) {
+        if let Some(tx) = tx.into_inner() {
+            let tx_hash = tx.hash();
+            self.pending_funding_txs.insert(tx_hash, tx);
+        }
+    }
+
+    pub fn remove(&mut self, tx_hash: &Byte32) {
+        self.pending_funding_txs.remove(tx_hash);
+    }
+
+    pub fn apply_to_cell_collector(
+        &self,
+        collector: &mut dyn CellCollector,
+    ) -> Result<(), FundingError> {
+        for tx in self.pending_funding_txs.values() {
+            collector.apply_tx(tx.data(), u64::MAX).map_err(|err| {
+                error!("Failed to apply exclusion list to cell collector: {}", err);
+                FundingError::FundingExclusionError
+            })?;
+        }
+        Ok(())
+    }
+}

src/ckb/funding/funding_tx.rs

@@ -1,4 +1,4 @@
-use super::super::FundingError;
+use super::{super::FundingError, FundingExclusion};
 use crate::{ckb::contracts::get_udt_cell_deps, fiber::serde_utils::EntityHex};
 use anyhow::anyhow;
 use ckb_sdk::{
@@ -280,7 +280,7 @@ impl FundingTxBuilder {
         )));
     }
 
-    fn build(self) -> Result<FundingTx, FundingError> {
+    fn build(self, funding_exclusion: &mut FundingExclusion) -> Result<FundingTx, FundingError> {
         // Build ScriptUnlocker
         let signer = SecpCkbRawKeySigner::new_with_secret_keys(vec![]);
         let sighash_unlocker = SecpSighashUnlocker::from(Box::new(signer) as Box<_>);
@@ -320,6 +320,7 @@ impl FundingTxBuilder {
 
         let header_dep_resolver = DefaultHeaderDepResolver::new(&self.context.rpc_url);
         let mut cell_collector = DefaultCellCollector::new(&self.context.rpc_url);
+        funding_exclusion.apply_to_cell_collector(&mut cell_collector)?;
         let tx_dep_provider = DefaultTransactionDependencyProvider::new(&self.context.rpc_url, 10);
 
         let (tx, _) = self.build_unlocked(
@@ -334,7 +335,12 @@ impl FundingTxBuilder {
         let mut funding_tx = self.funding_tx;
         let tx_builder = tx.as_advanced_builder();
         debug!("final tx_builder: {:?}", tx_builder);
+        let old_tx_hash_opt = funding_tx.tx.as_ref().map(|tx| tx.hash());
         funding_tx.update_for_self(tx)?;
+        if let Some(tx_hash) = old_tx_hash_opt {
+            funding_exclusion.remove(&tx_hash);
+        }
+        funding_exclusion.insert(funding_tx.clone());
         Ok(funding_tx)
     }
 }
@@ -360,13 +366,14 @@ impl FundingTx {
         self,
         request: FundingRequest,
         context: FundingContext,
+        funding_exclusion: &mut FundingExclusion,
     ) -> Result<Self, FundingError> {
         let builder = FundingTxBuilder {
             funding_tx: self,
             request,
             context,
         };
-        builder.build()
+        builder.build(funding_exclusion)
     }
 
     pub fn sign(

src/ckb/funding/mod.rs

@@ -1,4 +1,7 @@
+mod funding_exclusion;
 mod funding_tx;
 
 pub(crate) use funding_tx::FundingContext;
 pub use funding_tx::{FundingRequest, FundingTx};
+
+pub(crate) use funding_exclusion::FundingExclusion;

src/ckb/tests/test_utils.rs

@@ -379,6 +379,14 @@ impl Actor for MockChainActor {
                     );
                 }
             }
+
+            AddTxs(_txs) => {
+                unimplemented!()
+            }
+            RemoveTx(_tx_hash) => {
+                unimplemented!()
+            }
+
             Sign(tx, reply_port) => {
                 // We don't need to sign the funding transaction in mock chain actor,
                 // as any funding transaction is considered correct if we can successfully

src/fiber/network.rs

@@ -3745,6 +3745,7 @@ where
         );
         let network = self.network.clone();
         self.broadcast_tx_with_callback(transaction, move |result| {
+            // TODO: remove tx from exclusion list
             debug!("Funding transaction broadcast result: {:?}", &result);
             let message = match result {
                 Ok(TraceTxResponse {
@@ -3858,6 +3859,11 @@ where
             )),
         )
         .await;
+
+        // Cleanup exculsion list
+        self.chain_actor
+            .send_message(CkbChainMessage::RemoveTx(outpoint.tx_hash()))
+            .expect(ASSUME_CHAIN_ACTOR_ALWAYS_ALIVE_FOR_NOW);
     }
 
     async fn on_commitment_transaction_confirmed(&mut self, tx_hash: Hash256, channel_id: Hash256) {