From d6524ea517b5b3905908546777edcf140d69a4e3 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Fri, 13 Aug 2021 16:30:13 -0700
Subject: [PATCH] keychain+lnwallet: when fetching priv keys or signing try to
 use cache

In this commit, we start to optimistically use the new private key cache
that was added to btcwallet. As is, btcwallet will cache the decrypted
account keys for each scope in memory. However, the existing methods
to derive a child key from those account keys requires a write database
transaction, and will re-derive the private key using BIP-32 each time.

The newly added `DeriveFromKeyPathCache` lets us skip all this and
directly use a cache assuming the account info is already cached. The
new logic will try to use this method, but if it fails fall back to the
existing `DeriveFromKeyPath` method. All calls after this will use this
new cached key.

Fixes https://github.com/lightningnetwork/lnd/issues/5125.

Basic benchmark before the btcwallet change and after:
```
benchmark                    old ns/op     new ns/op     delta
BenchmarkDerivePrivKey-8     22840583      125           -100.00%

benchmark                    old allocs     new allocs     delta
BenchmarkDerivePrivKey-8     89             2              -97.75%

benchmark                    old bytes     new bytes     delta
BenchmarkDerivePrivKey-8     10225         24            -99.77%
```
---
 keychain/btcwallet.go        | 28 ++++++++++++++++++++++------
 lnwallet/btcwallet/signer.go | 14 ++++++++++++++
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/keychain/btcwallet.go b/keychain/btcwallet.go
index 02b721399..8d5dab3bd 100644
--- a/keychain/btcwallet.go
+++ b/keychain/btcwallet.go
@@ -252,14 +252,30 @@ func (b *BtcWalletKeyRing) DerivePrivKey(keyDesc KeyDescriptor) (
 
 	var key *btcec.PrivateKey
 
-	db := b.wallet.Database()
-	err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
-		addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)
+	scope, err := b.keyScope()
+	if err != nil {
+		return nil, err
+	}
 
-		scope, err := b.keyScope()
-		if err != nil {
-			return err
+	// First, attempt to see if we can read the key directly from
+	// btcwallet's internal cache, if we can then we can skip all the
+	// operations below (fast path).
+	if keyDesc.PubKey == nil {
+		keyPath := waddrmgr.DerivationPath{
+			InternalAccount: uint32(keyDesc.Family),
+			Account:         uint32(keyDesc.Family),
+			Branch:          0,
+			Index:           keyDesc.Index,
 		}
+		privKey, err := scope.DeriveFromKeyPathCache(keyPath)
+		if err == nil {
+			return privKey, nil
+		}
+	}
+
+	db := b.wallet.Database()
+	err = walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
+		addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)
 
 		// If the account doesn't exist, then we may need to create it
 		// for the first time in order to derive the keys that we
diff --git a/lnwallet/btcwallet/signer.go b/lnwallet/btcwallet/signer.go
index 953ec1006..bc48360e6 100644
--- a/lnwallet/btcwallet/signer.go
+++ b/lnwallet/btcwallet/signer.go
@@ -75,6 +75,20 @@ func (b *BtcWallet) deriveKeyByLocator(keyLoc keychain.KeyLocator) (*btcec.Priva
 		return nil, err
 	}
 
+	// First try to read the key from the cached store, if this fails, then
+	// we'll fall through to the method below that requires a database
+	// transaction.
+	path := waddrmgr.DerivationPath{
+		InternalAccount: uint32(keyLoc.Family),
+		Account:         uint32(keyLoc.Family),
+		Branch:          0,
+		Index:           keyLoc.Index,
+	}
+	privKey, err := scopedMgr.DeriveFromKeyPathCache(path)
+	if err == nil {
+		return privKey, nil
+	}
+
 	var key *btcec.PrivateKey
 	err = walletdb.Update(b.db, func(tx walletdb.ReadWriteTx) error {
 		addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)