keychain: cache derived priv key in BtcWalletKeyRing.ECDH to avoid per-call DB txn

keychain/btcwallet.go

@@ -2,6 +2,7 @@ package keychain
 
 import (
 	"crypto/sha256"
+	"errors"
 	"fmt"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -12,6 +13,8 @@ import (
 	"github.com/btcsuite/btcwallet/waddrmgr"
 	"github.com/btcsuite/btcwallet/wallet"
 	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/lightninglabs/neutrino/cache"
+	"github.com/lightninglabs/neutrino/cache/lru"
 )
 
 const (
@@ -22,6 +25,13 @@ const (
 	// CoinTypeTestnet specifies the BIP44 coin type for all testnet key
 	// derivation.
 	CoinTypeTestnet = 1
+
+	// ecdhPrivKeyCacheSize bounds the number of derived ECDH private keys
+	// that we'll keep memoized at any one time. In practice ECDH is only
+	// invoked against a handful of our own keys (one per key family used
+	// for onion decryption), so this size is well above the expected
+	// working set.
+	ecdhPrivKeyCacheSize = 1000
 )
 
 var (
@@ -57,6 +67,28 @@ type BtcWalletKeyRing struct {
 	// lightningScope is a pointer to the scope that we'll be using as a
 	// sub key manager to derive all the keys that we require.
 	lightningScope *waddrmgr.ScopedKeyManager
+
+	// ecdhPrivKeyCache memoizes the private keys derived for ECDH so that
+	// each subsequent ECDH call against the same key avoids the read-write
+	// wallet DB transaction (and the bbolt fdatasync that transaction
+	// forces). The cache is keyed by the compressed-serialized public key
+	// and bounded by ecdhPrivKeyCacheSize via an LRU eviction policy.
+	ecdhPrivKeyCache *lru.Cache[
+		[btcec.PubKeyBytesLenCompressed]byte, *cachedPrivKey,
+	]
+}
+
+// cachedPrivKey wraps a *btcec.PrivateKey so it can be stored in the LRU
+// cache, which requires values to implement the cache.Value interface.
+type cachedPrivKey struct {
+	privKey *btcec.PrivateKey
+}
+
+// Size returns the "size" of an entry. We return 1 so that the LRU cache
+// bounds the total number of entries rather than doing byte-accurate
+// accounting.
+func (c *cachedPrivKey) Size() (uint64, error) {
+	return 1, nil
 }
 
 // NewBtcWalletKeyRing creates a new implementation of the
@@ -76,6 +108,9 @@ func NewBtcWalletKeyRing(w wallet.Interface, coinType uint32) SecretKeyRing {
 	return &BtcWalletKeyRing{
 		wallet:        w,
 		chainKeyScope: chainKeyScope,
+		ecdhPrivKeyCache: lru.NewCache[
+			[btcec.PubKeyBytesLenCompressed]byte, *cachedPrivKey,
+		](ecdhPrivKeyCacheSize),
 	}
 }
 
@@ -385,11 +420,15 @@ func (b *BtcWalletKeyRing) DerivePrivKey(keyDesc KeyDescriptor) (
 //
 //	sx := k*P s := sha256(sx.SerializeCompressed())
 //
+// The derived private key for keyDesc is memoized after the first call so
+// repeated ECDH operations against the same key avoid reopening a read-write
+// wallet DB transaction (and the bbolt fdatasync per call) on the hot path.
+//
 // NOTE: This is part of the keychain.ECDHRing interface.
 func (b *BtcWalletKeyRing) ECDH(keyDesc KeyDescriptor,
 	pub *btcec.PublicKey) ([32]byte, error) {
 
-	privKey, err := b.DerivePrivKey(keyDesc)
+	privKey, err := b.derivePrivKeyForECDH(keyDesc)
 	if err != nil {
 		return [32]byte{}, err
 	}
@@ -408,6 +447,43 @@ func (b *BtcWalletKeyRing) ECDH(keyDesc KeyDescriptor,
 	return h, nil
 }
 
+// derivePrivKeyForECDH returns the private key for keyDesc, consulting and
+// populating the per-keyring ECDH cache. The cache is keyed by the
+// compressed-serialized public key, which uniquely identifies the private key
+// regardless of whether DerivePrivKey takes the path-based or the PubKey-scan
+// branch. When the descriptor has no public key set we cannot cache (we have
+// nothing collision-free to key on without first deriving) and forward to
+// DerivePrivKey directly. Production ECDH callers always supply a PubKey, so
+// the no-PubKey path is not on the hot path.
+func (b *BtcWalletKeyRing) derivePrivKeyForECDH(keyDesc KeyDescriptor) (
+	*btcec.PrivateKey, error) {
+
+	if keyDesc.PubKey == nil {
+		return b.DerivePrivKey(keyDesc)
+	}
+
+	var cacheKey [btcec.PubKeyBytesLenCompressed]byte
+	copy(cacheKey[:], keyDesc.PubKey.SerializeCompressed())
+
+	if v, err := b.ecdhPrivKeyCache.Get(cacheKey); err == nil {
+		return v.privKey, nil
+	} else if !errors.Is(err, cache.ErrElementNotFound) {
+		return nil, err
+	}
+
+	priv, err := b.DerivePrivKey(keyDesc)
+	if err != nil {
+		return nil, err
+	}
+
+	// Insertion is best-effort: a Put failure (e.g. the entry exceeds
+	// capacity, which can't happen here since each entry is size 1) only
+	// means the next ECDH against this key takes the slow path again.
+	_, _ = b.ecdhPrivKeyCache.Put(cacheKey, &cachedPrivKey{privKey: priv})
+
+	return priv, nil
+}
+
 // SignMessage signs the given message, single or double SHA256 hashing it
 // first, with the private key described in the key locator.
 //

onionmessage/actor_bench_test.go

@@ -0,0 +1,182 @@
+package onionmessage
+
+import (
+	"bytes"
+	"testing"
+	"time"
+
+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/chaincfg"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcwallet/snacl"
+	"github.com/btcsuite/btcwallet/waddrmgr"
+	"github.com/btcsuite/btcwallet/wallet"
+	"github.com/btcsuite/btcwallet/walletdb"
+	_ "github.com/btcsuite/btcwallet/walletdb/bdb"
+	sphinx "github.com/lightningnetwork/lightning-onion"
+	"github.com/lightningnetwork/lnd/keychain"
+	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/record"
+	"github.com/lightningnetwork/lnd/routing/route"
+	"github.com/stretchr/testify/require"
+)
+
+// noopDispatcher is a zero-overhead OnionMessageUpdateDispatcher used for
+// benchmarks so the dispatch path does not perturb timing.
+type noopDispatcher struct{}
+
+func (noopDispatcher) SendUpdate(any) error { return nil }
+
+// noopSender is a zero-overhead PeerMessageSender used for benchmarks.
+type noopSender struct{}
+
+func (noopSender) SendToPeer([33]byte, *lnwire.OnionMessage) error {
+	return nil
+}
+
+// testHDSeed is the deterministic seed used to create the benchmark wallet.
+var testHDSeed = chainhash.Hash{
+	0xb7, 0x94, 0x38, 0x5f, 0x2d, 0x1e, 0xf7, 0xab,
+	0x4d, 0x92, 0x73, 0xd1, 0x90, 0x63, 0x81, 0xb4,
+	0x4f, 0x2f, 0x6f, 0x25, 0x98, 0xa3, 0xef, 0xb9,
+	0x69, 0x49, 0x18, 0x83, 0x31, 0x98, 0x47, 0x53,
+}
+
+// lightningAddrSchema mirrors keychain/btcwallet.go (unexported there).
+var lightningAddrSchema = waddrmgr.ScopeAddrSchema{
+	ExternalAddrType: waddrmgr.WitnessPubKey,
+	InternalAddrType: waddrmgr.WitnessPubKey,
+}
+
+// waddrmgrNamespaceKey mirrors keychain/btcwallet.go (unexported there).
+var waddrmgrNamespaceKey = []byte("waddrmgr")
+
+// createBenchBtcWallet builds a real on-disk btcwallet (simnet, fast scrypt)
+// suitable for driving keychain.BtcWalletKeyRing in benchmarks. It mirrors
+// the unexported createTestBtcWallet in keychain/interface_test.go so the
+// onionmessage bench can construct keychain.PubKeyECDH against a real
+// SecretKeyRing — exactly as server.go does in production.
+func createBenchBtcWallet(b *testing.B) *wallet.Wallet {
+	b.Helper()
+
+	fast := waddrmgr.FastScryptOptions
+	waddrmgr.SetSecretKeyGen(func(passphrase *[]byte,
+		_ *waddrmgr.ScryptOptions) (*snacl.SecretKey, error) {
+
+		return snacl.NewSecretKey(passphrase, fast.N, fast.R, fast.P)
+	})
+
+	loader := wallet.NewLoader(
+		&chaincfg.SimNetParams, b.TempDir(), true, time.Second*10, 0,
+	)
+
+	pass := []byte("test")
+	w, err := loader.CreateNewWallet(
+		pass, pass, testHDSeed[:], time.Time{},
+	)
+	require.NoError(b, err)
+	require.NoError(b, w.Unlock(pass, nil))
+
+	scope := waddrmgr.KeyScope{
+		Purpose: keychain.BIP0043Purpose,
+		Coin:    keychain.CoinTypeBitcoin,
+	}
+	if _, err := w.Manager.FetchScopedKeyManager(scope); err != nil {
+		require.NoError(b, walletdb.Update(w.Database(),
+			func(tx walletdb.ReadWriteTx) error {
+				ns := tx.ReadWriteBucket(waddrmgrNamespaceKey)
+				_, err := w.Manager.NewScopedKeyManager(
+					ns, scope, lightningAddrSchema,
+				)
+
+				return err
+			}))
+	}
+
+	b.Cleanup(func() { w.Lock() })
+
+	return w
+}
+
+// BenchmarkOnionMessagePipeline measures OnionPeerActor.Receive end to end on
+// the deliver path: decode -> sphinx process -> routing decision -> dispatch.
+// It isolates the actor + pipeline overhead on top of sphinx.
+// ProcessOnionPacket.
+func BenchmarkOnionMessagePipeline(b *testing.B) {
+	// Mirror server.go: derive the node key from a real BtcWalletKeyRing
+	// and wrap it in keychain.NewPubKeyECDH. This is the exact onion-key
+	// type the production sphinx router uses.
+	w := createBenchBtcWallet(b)
+	keyRing := keychain.NewBtcWalletKeyRing(w, keychain.CoinTypeBitcoin)
+
+	nodeKeyDesc, err := keyRing.DeriveKey(keychain.KeyLocator{
+		Family: keychain.KeyFamilyNodeKey,
+		Index:  0,
+	})
+	require.NoError(b, err)
+
+	nodeKeyECDH := keychain.NewPubKeyECDH(nodeKeyDesc, keyRing)
+
+	router := sphinx.NewRouter(nodeKeyECDH, sphinx.NewNoOpReplayLog())
+	require.NoError(b, router.Start())
+	b.Cleanup(func() { router.Stop() })
+
+	var peerPubKey [33]byte
+	copy(peerPubKey[:], nodeKeyDesc.PubKey.SerializeCompressed())
+
+	a := &OnionPeerActor{
+		peerPubKey:       peerPubKey,
+		peerSender:       noopSender{},
+		router:           router,
+		resolver:         newMockNodeIDResolver(),
+		updateDispatcher: noopDispatcher{},
+	}
+
+	// Build a single-hop "deliver" onion message addressed to nodeKey.
+	plain, err := record.EncodeBlindedRouteData(&record.BlindedRouteData{})
+	require.NoError(b, err)
+
+	hops := []*sphinx.HopInfo{
+		{NodePub: nodeKeyDesc.PubKey, PlainText: plain},
+	}
+
+	sessionKey, err := btcec.NewPrivateKey()
+	require.NoError(b, err)
+
+	blindedPath, err := sphinx.BuildBlindedPath(sessionKey, hops)
+	require.NoError(b, err)
+
+	sphinxPath, err := route.OnionMessageBlindedPathToSphinxPath(
+		blindedPath.Path, nil, nil,
+	)
+	require.NoError(b, err)
+
+	onionSessionKey, err := btcec.NewPrivateKey()
+	require.NoError(b, err)
+
+	pkt, err := sphinx.NewOnionPacket(
+		sphinxPath, onionSessionKey, nil,
+		sphinx.DeterministicPacketFiller,
+		sphinx.WithMaxPayloadSize(sphinx.MaxRoutingPayloadSize),
+	)
+	require.NoError(b, err)
+
+	var buf bytes.Buffer
+	require.NoError(b, pkt.Encode(&buf))
+
+	onionBlob := buf.Bytes()
+	pathKey := blindedPath.SessionKey.PubKey()
+
+	ctx := b.Context()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		req := &Request{msg: lnwire.OnionMessage{
+			PathKey:   pathKey,
+			OnionBlob: onionBlob,
+		}}
+		_ = a.Receive(ctx, req)
+	}
+}