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

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

@@ -0,0 +1,45 @@
+package wallet
+
+import "sync"
+
+// keyPathEntry records the derivation scheme and relative key path a script
+// belongs to.
+type keyPathEntry struct {
+	derivationScheme string
+	keyPath          string
+}
+
+// keyPathCache caches 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 never need invalidation. The
+// cache lets the wallet skip the per-input NBXplorer script lookups when signing
+// inputs it has already seen as UTXOs.
+type keyPathCache struct {
+	mu    sync.RWMutex
+	cache map[string]keyPathEntry
+}
+
+func newKeyPathCache() *keyPathCache {
+	return &keyPathCache{cache: make(map[string]keyPathEntry)}
+}
+
+// get returns the cached entry for the given script, if any.
+func (c *keyPathCache) get(script string) (keyPathEntry, bool) {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	entry, ok := c.cache[script]
+	return entry, ok
+}
+
+// set caches the derivation scheme and key path for the given script. 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.mu.Lock()
+	defer c.mu.Unlock()
+	c.cache[script] = keyPathEntry{derivationScheme: derivationScheme, keyPath: keyPath}
+}

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

@@ -0,0 +1,77 @@
+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_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,15 @@ 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{opts, newOutpointLocker(time.Minute), nil, newKeyPathCache(), make(chan bool)}
 }
 
 func (w *wallet) GetReadyUpdate(ctx context.Context) <-chan bool {
@@ -336,6 +337,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 +390,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 +835,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 +844,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 +990,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 +1006,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 {