arkd-wallet: source UTXOs from NBXplorer events and cache wallet key paths

pkg/arkd-wallet/core/application/wallet/keypath_cache.go

@@ -0,0 +1,52 @@
+package wallet
+
+import lru "github.com/hashicorp/golang-lru/v2"
+
+// keyPathCacheSize bounds the number of cached script -> key path entries.
+//
+// A signing operation only needs the scripts of the inputs being signed, which
+// are the most recently listed UTXOs, so evicting older entries is always safe:
+// a miss simply falls back to the NBXplorer lookup. The bound keeps memory flat
+// on a long-running server whose wallet churns through many addresses (e.g. the
+// connector account across many rounds).
+const keyPathCacheSize = 50_000
+
+// keyPathEntry records the derivation scheme and relative key path a script
+// belongs to.
+type keyPathEntry struct {
+	derivationScheme string
+	keyPath          string
+}
+
+// keyPathCache is a bounded, concurrency-safe cache of the script ->
+// (derivation scheme, key path) mapping.
+//
+// The mapping is immutable: a given script always derives from the same key
+// path under the same account, so cached entries are never stale. The cache lets
+// the wallet skip the per-input NBXplorer script lookups when signing inputs it
+// has already seen as UTXOs.
+type keyPathCache struct {
+	cache *lru.Cache[string, keyPathEntry]
+}
+
+func newKeyPathCache() *keyPathCache {
+	// lru.New only errors on a non-positive size, which is a positive constant
+	// here, so the error cannot occur.
+	cache, _ := lru.New[string, keyPathEntry](keyPathCacheSize)
+	return &keyPathCache{cache: cache}
+}
+
+// get returns the cached entry for the given script, if any.
+func (c *keyPathCache) get(script string) (keyPathEntry, bool) {
+	return c.cache.Get(script)
+}
+
+// set caches the derivation scheme and key path for the given script, evicting
+// the least-recently-used entry if the cache is full. Entries with an empty
+// script or key path are ignored, as they cannot be used to derive a key later.
+func (c *keyPathCache) set(script, derivationScheme, keyPath string) {
+	if script == "" || keyPath == "" {
+		return
+	}
+	c.cache.Add(script, keyPathEntry{derivationScheme: derivationScheme, keyPath: keyPath})
+}

pkg/arkd-wallet/core/application/wallet/keypath_cache_test.go

@@ -0,0 +1,97 @@
+package wallet
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestKeyPathCache(t *testing.T) {
+	t.Run("set and get", func(t *testing.T) {
+		c := newKeyPathCache()
+
+		c.set("script-a", "scheme-main", "0/0")
+		entry, ok := c.get("script-a")
+		require.True(t, ok)
+		require.Equal(t, "scheme-main", entry.derivationScheme)
+		require.Equal(t, "0/0", entry.keyPath)
+	})
+
+	t.Run("get miss", func(t *testing.T) {
+		c := newKeyPathCache()
+
+		_, ok := c.get("missing")
+		require.False(t, ok)
+	})
+
+	t.Run("empty script or key path is not cached", func(t *testing.T) {
+		c := newKeyPathCache()
+
+		c.set("", "scheme", "0/0")
+		_, ok := c.get("")
+		require.False(t, ok)
+
+		c.set("script-b", "scheme", "")
+		_, ok = c.get("script-b")
+		require.False(t, ok)
+	})
+
+	t.Run("overwrite", func(t *testing.T) {
+		c := newKeyPathCache()
+
+		c.set("script-c", "scheme-main", "0/0")
+		c.set("script-c", "scheme-connector", "1/2")
+
+		entry, ok := c.get("script-c")
+		require.True(t, ok)
+		require.Equal(t, "scheme-connector", entry.derivationScheme)
+		require.Equal(t, "1/2", entry.keyPath)
+	})
+}
+
+func TestKeyPathCache_Bounded(t *testing.T) {
+	c := newKeyPathCache()
+
+	total := keyPathCacheSize + 1000
+	for i := range total {
+		c.set(fmt.Sprintf("script-%d", i), "scheme", fmt.Sprintf("0/%d", i))
+	}
+
+	// the cache never grows past its bound
+	require.Equal(t, keyPathCacheSize, c.cache.Len())
+
+	// the oldest entries are evicted, the most recent are retained
+	_, ok := c.get("script-0")
+	require.False(t, ok)
+
+	entry, ok := c.get(fmt.Sprintf("script-%d", total-1))
+	require.True(t, ok)
+	require.Equal(t, fmt.Sprintf("0/%d", total-1), entry.keyPath)
+}
+
+func TestKeyPathCache_Concurrent(t *testing.T) {
+	c := newKeyPathCache()
+	numberOfRoutines := 100
+
+	wg := sync.WaitGroup{}
+	wg.Add(numberOfRoutines)
+
+	for i := range numberOfRoutines {
+		go func(i int) {
+			defer wg.Done()
+			script := fmt.Sprintf("script-%d", i)
+			c.set(script, "scheme", fmt.Sprintf("0/%d", i))
+			_, _ = c.get(script)
+		}(i)
+	}
+
+	wg.Wait()
+
+	for i := range numberOfRoutines {
+		entry, ok := c.get(fmt.Sprintf("script-%d", i))
+		require.True(t, ok)
+		require.Equal(t, fmt.Sprintf("0/%d", i), entry.keyPath)
+	}
+}

pkg/arkd-wallet/core/application/wallet/service.go

@@ -54,14 +54,20 @@ type WalletOptions struct {
 type wallet struct {
 	WalletOptions
 
-	locker  *outpointLocker
-	keyMgr  *keyManager
-	readyCh chan bool
+	locker   *outpointLocker
+	keyMgr   *keyManager
+	keyPaths *keyPathCache
+	readyCh  chan bool
 }
 
 // New creates a new WalletService service
 func New(opts WalletOptions) application.WalletService {
-	return &wallet{opts, newOutpointLocker(time.Minute), nil, make(chan bool)}
+	return &wallet{
+		WalletOptions: opts,
+		locker:        newOutpointLocker(time.Minute),
+		keyPaths:      newKeyPathCache(),
+		readyCh:       make(chan bool),
+	}
 }
 
 func (w *wallet) GetReadyUpdate(ctx context.Context) <-chan bool {
@@ -336,6 +342,7 @@ func (w *wallet) ListConnectorUtxos(ctx context.Context, connectorAddress string
 	if err != nil {
 		return nil, err
 	}
+	w.cacheKeyPaths(w.keyMgr.connectorAccountDerivationScheme, connectorAccountUtxos)
 
 	lockedOutpoints, err := w.locker.get(ctx)
 	if err != nil {
@@ -388,6 +395,7 @@ func (w *wallet) SelectUtxos(ctx context.Context, amount uint64, confirmedOnly b
 	if err != nil {
 		return nil, 0, err
 	}
+	w.cacheKeyPaths(w.keyMgr.mainAccountDerivationScheme, mainAccountUtxos)
 
 	lockedOutpoints, err := w.locker.get(ctx)
 	if err != nil {
@@ -832,6 +840,7 @@ func (w *wallet) withdrawAll(ctx context.Context, feeRate chainfee.SatPerKVByte,
 	if err != nil {
 		return nil, err
 	}
+	w.cacheKeyPaths(w.keyMgr.mainAccountDerivationScheme, mainAccountUtxos)
 
 	utxos = append(utxos, mainAccountUtxos...)
 
@@ -840,6 +849,7 @@ func (w *wallet) withdrawAll(ctx context.Context, feeRate chainfee.SatPerKVByte,
 		if err != nil {
 			return nil, err
 		}
+		w.cacheKeyPaths(w.keyMgr.connectorAccountDerivationScheme, connectorAccountUtxos)
 		utxos = append(utxos, connectorAccountUtxos...)
 	}
 
@@ -985,6 +995,12 @@ func (w *wallet) getPrivateKeyFromScript(ctx context.Context, scriptPubKey strin
 		return nil, ErrWalletLocked
 	}
 
+	// A cache hit lets us derive the key without any NBXplorer lookup. The cache
+	// is populated whenever we list UTXOs, which already carry their key path.
+	if entry, ok := w.keyPaths.get(scriptPubKey); ok {
+		return w.keyMgr.deriveKey(entry.derivationScheme, entry.keyPath)
+	}
+
 	accountsDerivationSchemes := []string{
 		w.keyMgr.mainAccountDerivationScheme,
 		w.keyMgr.connectorAccountDerivationScheme,
@@ -995,18 +1011,33 @@ func (w *wallet) getPrivateKeyFromScript(ctx context.Context, scriptPubKey strin
 		if err != nil {
 			continue
 		}
+		// a script tracked under a scheme always has a key path; if it is empty
+		// the script does not belong to this account, so try the next scheme.
+		if scriptPubKeyDetails.KeyPath == "" {
+			continue
+		}
 
+		w.keyPaths.set(scriptPubKey, derivationScheme, scriptPubKeyDetails.KeyPath)
 		return w.keyMgr.deriveKey(derivationScheme, scriptPubKeyDetails.KeyPath)
 	}
 
 	return nil, nil
 }
 
+// cacheKeyPaths records the script -> key path mapping for the given UTXOs so
+// that later signing of these scripts can skip the per-input NBXplorer lookup.
+func (w *wallet) cacheKeyPaths(derivationScheme string, utxos []ports.Utxo) {
+	for _, utxo := range utxos {
+		w.keyPaths.set(utxo.Script, derivationScheme, utxo.KeyPath)
+	}
+}
+
 func (w *wallet) getBalance(ctx context.Context, derivationScheme string) (uint64, uint64, error) {
 	utxos, err := w.Nbxplorer.GetUtxos(ctx, derivationScheme)
 	if err != nil {
 		return 0, 0, err
 	}
+	w.cacheKeyPaths(derivationScheme, utxos)
 
 	lockedOutpoints, err := w.locker.get(ctx)
 	if err != nil {