waddrmgr: add new DeriveFromKeyPathCache method for faster key retrieval

In this commit, we add a new method `DeriveFromKeyPathCache` that gives
callers a way to more quickly obtain a private key they know they'll be
using frequently. This method lets a caller avoid the write database
transaction as well as the EC operations to derive the key itself (BIP
32).

waddrmgr/error.go

@@ -139,6 +139,10 @@ const (
 	// ErrBlockNotFound is returned when we attempt to retrieve the hash for
 	// a block that we do not know of.
 	ErrBlockNotFound
+
+	// ErrAccountNotCached is returned when we attempt to perform an
+	// operation that relies on an account begin cached but it isn't.
+	ErrAccountNotCached
 )
 
 // Map of ErrorCode values back to their constant names for pretty printing.
@@ -165,6 +169,7 @@ var errorCodeStrings = map[ErrorCode]string{
 	ErrCallBackBreak:     "ErrCallBackBreak",
 	ErrEmptyPassphrase:   "ErrEmptyPassphrase",
 	ErrScopeNotFound:     "ErrScopeNotFound",
+	ErrAccountNotCached:  "ErrAccountNotCached",
 }
 
 // String returns the ErrorCode as a human-readable name.

waddrmgr/manager.go

@@ -17,6 +17,7 @@ import (
 	"github.com/btcsuite/btcwallet/internal/zero"
 	"github.com/btcsuite/btcwallet/snacl"
 	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/lightninglabs/neutrino/cache/lru"
 )
 
 const (
@@ -602,11 +603,12 @@ func (m *Manager) NewScopedKeyManager(ns walletdb.ReadWriteBucket,
 	// Finally, we'll register this new scoped manager with the root
 	// manager.
 	m.scopedManagers[scope] = &ScopedKeyManager{
-		scope:       scope,
-		addrSchema:  addrSchema,
-		rootManager: m,
-		addrs:       make(map[addrKey]ManagedAddress),
-		acctInfo:    make(map[uint32]*accountInfo),
+		scope:        scope,
+		addrSchema:   addrSchema,
+		rootManager:  m,
+		addrs:        make(map[addrKey]ManagedAddress),
+		acctInfo:     make(map[uint32]*accountInfo),
+		privKeyCache: lru.NewCache(defaultPrivKeyCacheSize),
 	}
 	m.externalAddrSchemas[addrSchema.ExternalAddrType] = append(
 		m.externalAddrSchemas[addrSchema.ExternalAddrType], scope,
@@ -1620,10 +1622,11 @@ func loadManager(ns walletdb.ReadBucket, pubPassphrase []byte,
 		}
 
 		scopedManagers[scope] = &ScopedKeyManager{
-			scope:      scope,
-			addrSchema: *scopeSchema,
-			addrs:      make(map[addrKey]ManagedAddress),
-			acctInfo:   make(map[uint32]*accountInfo),
+			scope:        scope,
+			addrSchema:   *scopeSchema,
+			addrs:        make(map[addrKey]ManagedAddress),
+			acctInfo:     make(map[uint32]*accountInfo),
+			privKeyCache: lru.NewCache(defaultPrivKeyCacheSize),
 		}
 
 		return nil

waddrmgr/manager_test.go

@@ -14,6 +14,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcutil"
@@ -21,6 +22,7 @@ import (
 	"github.com/btcsuite/btcwallet/snacl"
 	"github.com/btcsuite/btcwallet/walletdb"
 	"github.com/davecgh/go-spew/spew"
+	"github.com/stretchr/testify/require"
 )
 
 // failingCryptoKey is an implementation of the EncryptorDecryptor interface
@@ -2764,3 +2766,96 @@ func testNewRawAccount(t *testing.T, _ *Manager, db walletdb.DB,
 			accountTargetAddr.AddrHash())
 	}
 }
+
+// TestDeriveFromKeyPathCache tests that the DeriveFromKeyPathCache method will
+// properly cache items in the cache, and return corresponding errors if the
+// account isn't properly cached.
+func TestDeriveFromKeyPathCache(t *testing.T) {
+	t.Parallel()
+
+	teardown, db := emptyDB(t)
+	defer teardown()
+
+	// We'll start the test by creating a new root manager that will be
+	// used for the duration of the test.
+	var mgr *Manager
+	err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
+		ns, err := tx.CreateTopLevelBucket(waddrmgrNamespaceKey)
+		if err != nil {
+			return err
+		}
+		err = Create(
+			ns, seed, pubPassphrase, privPassphrase,
+			&chaincfg.MainNetParams, fastScrypt, time.Time{},
+		)
+		if err != nil {
+			return err
+		}
+		mgr, err = Open(ns, pubPassphrase, &chaincfg.MainNetParams)
+		if err != nil {
+			return err
+		}
+
+		return mgr.Unlock(ns, privPassphrase)
+	})
+	require.NoError(t, err, "create/open: unexpected error: %v", err)
+
+	defer mgr.Close()
+
+	// Now that we have the manager created, we'll fetch one of the default
+	// scopes for usage within this test.
+	scopedMgr, err := mgr.FetchScopedKeyManager(KeyScopeBIP0044)
+	require.NoError(
+		t, err, "unable to fetch scope %v: %v", KeyScopeBIP0044, err,
+	)
+
+	keyPath := DerivationPath{
+		InternalAccount: 0,
+		Account:         hdkeychain.HardenedKeyStart,
+		Branch:          10,
+		Index:           1,
+	}
+
+	// Our test starts here, we'll attempt to derive a new key using the
+	// cached method. This should fail at first since the account itself
+	// isn't cached.
+	_, err = scopedMgr.DeriveFromKeyPathCache(keyPath)
+	if !IsError(err, ErrAccountNotCached) {
+		t.Fatalf("didn't get account not cached error: %v", err)
+	}
+
+	// Now we'll attempt to derive the key using the normal method that
+	// requires a database transaction.
+	var derivedKey *btcec.PrivateKey
+	err = walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
+		ns := tx.ReadWriteBucket(waddrmgrNamespaceKey)
+
+		managedAddr, err := scopedMgr.DeriveFromKeyPath(
+			ns, keyPath,
+		)
+		if err != nil {
+			return err
+		}
+
+		derivedKey, err = managedAddr.(ManagedPubKeyAddress).PrivKey()
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+	require.NoError(t, err, "unable to derive addr: %v", err)
+
+	// Next attempt to read the key again from the cache, it should succeed
+	// this time.
+	cachedKey, err := scopedMgr.DeriveFromKeyPathCache(keyPath)
+	require.NoError(t, err, "account wasn't cached")
+
+	// We should be able to read the key again.
+	cachedKey2, err := scopedMgr.DeriveFromKeyPathCache(keyPath)
+	require.NoError(t, err, "account wasn't cached")
+
+	// All three keys we have now should match exactly.
+	require.Equal(t, cachedKey.Serialize(), cachedKey2.Serialize())
+	require.Equal(t, derivedKey.Serialize(), cachedKey2.Serialize())
+}

waddrmgr/scoped_manager.go

@@ -14,6 +14,7 @@ import (
 	"github.com/btcsuite/btcwallet/internal/zero"
 	"github.com/btcsuite/btcwallet/netparams"
 	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/lightninglabs/neutrino/cache/lru"
 )
 
 // HDVersion represents the different supported schemes of hierarchical
@@ -51,6 +52,14 @@ const (
 	HDVersionSimNetBIP0044 HDVersion = 0x0420bd3a // spub
 )
 
+const (
+	// privKeyCacheSize is the default size of the LRU cache that we'll use
+	// to cache private keys to avoid DB and EC operations within the
+	// wallet. With the default sisize, we'll allocate up to 320 KB to
+	// caching private keys (ignoring pointer overhead, etc).
+	defaultPrivKeyCacheSize = 10_000
+)
+
 // DerivationPath represents a derivation path from a particular key manager's
 // scope.  Each ScopedKeyManager starts key derivation from the end of their
 // cointype hardened key: m/purpose'/cointype'. The fields in this struct allow
@@ -220,7 +229,7 @@ type ScopedKeyManager struct {
 	rootManager *Manager
 
 	// addrs is a cached map of all the addresses that we currently
-	// manager.
+	// manage.
 	addrs map[addrKey]ManagedAddress
 
 	// acctInfo houses information about accounts including what is needed
@@ -234,6 +243,11 @@ type ScopedKeyManager struct {
 	// order to encrypt it.
 	deriveOnUnlock []*unlockDeriveInfo
 
+	// privKeyCache stores the set of private keys that have been marked as
+	// items to be cached to allow us to avoid the database and EC
+	// operations each time a key need to be obtained.
+	privKeyCache *lru.Cache
+
 	mtx sync.RWMutex
 }
 
@@ -579,6 +593,77 @@ func (s *ScopedKeyManager) AccountProperties(ns walletdb.ReadBucket,
 	return props, nil
 }
 
+// cachedKey is an entry within the LRU map that stores private keys that are
+// to be used frequently. We use this wrapper struct to be able too report the
+// size of a given element to the cache.
+type cachedKey struct {
+	key *btcec.PrivateKey
+}
+
+// Size returns the size of this element. Rather than have the cache limit
+// based on bytes, we simply report that each element is of size 1, meaning we
+// can set our cached based on the amount of keys we want to store, rather than
+// the total size of all the keys.
+func (c *cachedKey) Size() (uint64, error) {
+	return 1, nil
+}
+
+// DeriveFromKeyPathCache is identical to DeriveFromKeyPath, however it'll fail
+// if the account refracted in the DerivationPath isn't already in the
+// in-memory cache. Callers looking for faster private key retrieval can opt to
+// call this method, which may fail if things aren't in the cache, then fall
+// back to the normal variant. The account can information can be drawn into
+// the cache if the normal DeriveFromKeyPath method is used, or the account is
+// looked up via any other means.
+func (s *ScopedKeyManager) DeriveFromKeyPathCache(
+	kp DerivationPath) (*btcec.PrivateKey, error) {
+
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	// First, try to look up the key itself in the proper cache, if the key
+	// is here, then we don't need to do anything further.
+	privKeyVal, err := s.privKeyCache.Get(kp)
+	if err == nil {
+		return privKeyVal.(*cachedKey).key, nil
+	}
+
+	// If the key isn't already in the cache, then we'll try to look up the
+	// account info in the cache, if this fails, then we exit here as we
+	// can't move forward without creating a DB transaction, and the point
+	// of this method is to avoid that.
+	acctInfo, ok := s.acctInfo[kp.InternalAccount]
+	if !ok {
+		return nil, managerError(
+			ErrAccountNotCached,
+			"", fmt.Errorf("acct %v not cached", kp.InternalAccount),
+		)
+	}
+
+	watchOnly := s.rootManager.WatchOnly()
+	private := !s.rootManager.IsLocked() && !watchOnly
+
+	// Now that we have the account information, we can derive the key
+	// directly.
+	addrKey, err := s.deriveKey(acctInfo, kp.Branch, kp.Index, private)
+	if err != nil {
+		return nil, err
+	}
+
+	// Now that we have the key, we'll attempt to insert it into the cache,
+	// and return it as is.
+	privKey, err := addrKey.ECPrivKey()
+	if err != nil {
+		return nil, err
+	}
+	_, err = s.privKeyCache.Put(kp, &cachedKey{key: privKey})
+	if err != nil {
+		return nil, err
+	}
+
+	return privKey, nil
+}
+
 // DeriveFromKeyPath attempts to derive a maximal child key (under the BIP0044
 // scheme) from a given key path. If key derivation isn't possible, then an
 // error will be returned.