Eliminate gap between generating next key in wallet and saving output for this key (#302)

wallet/src/receiver.rs

@@ -90,34 +90,27 @@ impl Handler for WalletReceiver {
 	}
 }
 
-// Read wallet data without acquiring the write lock.
 fn retrieve_existing_key(
-	config: &WalletConfig,
+	wallet_data: &WalletData,
 	key_id: Identifier,
-) -> Result<(Identifier, u32), Error> {
+) -> (Identifier, u32) {
-	let res = WalletData::read_wallet(&config.data_file_dir, |wallet_data| {
+	if let Some(existing) = wallet_data.get_output(&key_id) {
-		if let Some(existing) = wallet_data.get_output(&key_id) {
+		let key_id = existing.key_id.clone();
-			let key_id = existing.key_id.clone();
+		let derivation = existing.n_child;
-			let derivation = existing.n_child;
+		(key_id, derivation)
-			(key_id, derivation)
+	} else {
-		} else {
+		panic!("should never happen");
-			panic!("should never happen");
+	}
-		}
-	})?;
-	Ok(res)
 }
 
 fn next_available_key(
-	config: &WalletConfig,
+	wallet_data: &WalletData,
 	keychain: &Keychain,
-) -> Result<(Identifier, u32), Error> {
+) -> (Identifier, u32) {
-	let res = WalletData::read_wallet(&config.data_file_dir, |wallet_data| {
+	let root_key_id = keychain.root_key_id();
-		let root_key_id = keychain.root_key_id();
+	let derivation = wallet_data.next_child(root_key_id.clone());
-		let derivation = wallet_data.next_child(root_key_id.clone());
+	let key_id = keychain.derive_key_id(derivation).unwrap();
-		let key_id = keychain.derive_key_id(derivation).unwrap();
+	(key_id, derivation)
-		(key_id, derivation)
-	})?;
-	Ok(res)
 }
 
 /// Build a coinbase output and the corresponding kernel
@@ -127,15 +120,15 @@ pub fn receive_coinbase(
 	block_fees: &BlockFees,
 ) -> Result<(Output, TxKernel, BlockFees), Error> {
 	let root_key_id = keychain.root_key_id();
-	let key_id = block_fees.key_id();
-
-	let (key_id, derivation) = match key_id {
-		Some(key_id) => retrieve_existing_key(config, key_id)?,
-		None => next_available_key(config, keychain)?,
-	};
 
 	// Now acquire the wallet lock and write the new output.
-	WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
+	let (key_id, derivation) = WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
+		let key_id = block_fees.key_id();
+		let (key_id, derivation) = match key_id {
+			Some(key_id) => retrieve_existing_key(&wallet_data, key_id),
+			None => next_available_key(&wallet_data, keychain),
+		};
+
 		// track the new output and return the stuff needed for reward
 		wallet_data.add_output(OutputData {
 			root_key_id: root_key_id.clone(),
@@ -147,6 +140,8 @@ pub fn receive_coinbase(
 			lock_height: 0,
 			is_coinbase: true,
 		});
+
+		(key_id, derivation)
 	})?;
 
 	debug!(
@@ -178,8 +173,6 @@ fn receive_transaction(
 ) -> Result<Transaction, Error> {
 	let root_key_id = keychain.root_key_id();
 
-	let (key_id, derivation) = next_available_key(config, keychain)?;
-
 	// double check the fee amount included in the partial tx
  // we don't necessarily want to just trust the sender
  // we could just overwrite the fee here (but we won't) due to the ecdsa sig
@@ -193,6 +186,24 @@ fn receive_transaction(
 
 	let out_amount = amount - fee;
 
+	// operate within a lock on wallet data
+	let (key_id, derivation) = WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
+		let (key_id, derivation) = next_available_key(&wallet_data, keychain);
+
+		wallet_data.add_output(OutputData {
+			root_key_id: root_key_id.clone(),
+			key_id: key_id.clone(),
+			n_child: derivation,
+			value: out_amount,
+			status: OutputStatus::Unconfirmed,
+			height: 0,
+			lock_height: 0,
+			is_coinbase: false,
+		});
+
+		(key_id, derivation)
+	})?;
+
 	let (tx_final, _) = build::transaction(
 		vec![
 			build::initial_tx(partial),
@@ -207,20 +218,6 @@ fn receive_transaction(
  // excess).
 	tx_final.validate()?;
 
-	// operate within a lock on wallet data
-	WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
-		wallet_data.add_output(OutputData {
-			root_key_id: root_key_id.clone(),
-			key_id: key_id.clone(),
-			n_child: derivation,
-			value: out_amount,
-			status: OutputStatus::Unconfirmed,
-			height: 0,
-			lock_height: 0,
-			is_coinbase: false,
-		});
-	})?;
-
 	debug!(
 		LOGGER,
 		"Received txn and built output - {:?}, {:?}, {}",

wallet/src/sender.rs

@@ -47,7 +47,7 @@ pub fn issue_send_tx(
 	// proof of concept - set lock_height on the tx
 	let lock_height = chain_tip.height;
 
-	let (tx, blind_sum, coins, change_key, change_derivation, change) = build_send_tx(
+	let (tx, blind_sum, coins) = build_send_tx(
 		config,
 		keychain,
 		amount,
@@ -60,23 +60,9 @@ pub fn issue_send_tx(
 
 	let partial_tx = build_partial_tx(amount, blind_sum, tx);
 
-	let root_key_id = keychain.clone().root_key_id();
+	// Acquire wallet lock and lock the coins being spent
-	// Acquire wallet lock, add the new change output and lock coins being spent.
+	// so we avoid accidental double spend attempt
 	let update_wallet = || WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
-		// we got that far, time to start tracking the output representing our change
-		wallet_data.add_output(OutputData {
-			root_key_id: root_key_id.clone(),
-			key_id: change_key.clone(),
-			n_child: change_derivation,
-			value: change as u64,
-			status: OutputStatus::Unconfirmed,
-			height: 0,
-			lock_height: 0,
-			is_coinbase: false,
-		});
-
-		// now lock the ouputs we're spending so we avoid accidental double spend
-		// attempt
 		for coin in coins {
 			wallet_data.lock_output(&coin);
 		}
@@ -117,7 +103,7 @@ fn build_send_tx(
 	lock_height: u64,
 	max_outputs: usize,
 	default_strategy: bool,
-) -> Result<(Transaction, BlindingFactor, Vec<OutputData>, Identifier, u32, u64), Error> {
+) -> Result<(Transaction, BlindingFactor, Vec<OutputData>), Error> {
 	let key_id = keychain.clone().root_key_id();
 
 	// select some spendable coins from the wallet
@@ -149,7 +135,7 @@ fn build_send_tx(
 
 	let (tx, blind) = build::transaction(parts.0, &keychain)?;
 
-	Ok((tx, blind, coins, parts.1, parts.2, parts.3))
+	Ok((tx, blind, coins))
 }
 
 pub fn issue_burn_tx(
@@ -199,19 +185,6 @@ pub fn issue_burn_tx(
 	Ok(())
 }
 
-fn next_available_key(
-	config: &WalletConfig,
-	keychain: &Keychain,
-) -> Result<(Identifier, u32), Error> {
-	let res = WalletData::read_wallet(&config.data_file_dir, |wallet_data| {
-		let root_key_id = keychain.root_key_id();
-		let derivation = wallet_data.next_child(root_key_id.clone());
-		let key_id = keychain.derive_key_id(derivation).unwrap();
-		(key_id, derivation)
-	})?;
-	Ok(res)
-}
-
 fn inputs_and_change(
 	coins: &Vec<OutputData>,
 	config: &WalletConfig,
@@ -245,7 +218,25 @@ fn inputs_and_change(
 		parts.push(build::input(coin.value, key_id));
 	}
 
-	let (change_key, change_derivation) = next_available_key(config, keychain)?;
+	// track the output representing our change
+	let (change_key, change_derivation) = WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
+		let root_key_id = keychain.root_key_id();
+		let change_derivation = wallet_data.next_child(root_key_id.clone());
+		let change_key = keychain.derive_key_id(change_derivation).unwrap();
+
+		wallet_data.add_output(OutputData {
+			root_key_id: root_key_id.clone(),
+			key_id: change_key.clone(),
+			n_child: change_derivation,
+			value: change as u64,
+			status: OutputStatus::Unconfirmed,
+			height: 0,
+			lock_height: 0,
+			is_coinbase: false,
+		});
+
+		(change_key, change_derivation)
+	})?;
 
 	parts.push(build::output(change, change_key.clone()));