Make kaspawallet store the utxos sorted by amount (#1947)

* Make kaspawallet store the utxos sorted by amount, so that the bigger utxos are spent first - making it less likely a compound will be required * Start refactor addEntryToUTXOSet * Add GetUTXOsByBalances command to rpc * Store list of addresses, updated with the collectAddresses methods (replacing collectUTXOs methods) * Fix wrong commands in GetBalanceByAddress * Rename: refreshExistingUTXOs -> refreshUTXOs Co-authored-by: Ori Newman <[email protected]>
kaspanet · Feb 18, 2022 · be3a660 · be3a660
1 parent f452531
commit be3a660
Show file tree

Hide file tree

Showing 19 changed files with 906 additions and 379 deletions.
diff --git a/app/appmessage/error.go b/app/appmessage/error.go
@@ -38,6 +38,10 @@ type RPCError struct {
 	Message string
 }
 
+func (err RPCError) Error() string {
+	return err.Message
+}
+
 // RPCErrorf formats according to a format specifier and returns the string
 // as an RPCError.
 func RPCErrorf(format string, args ...interface{}) *RPCError {

diff --git a/app/appmessage/message.go b/app/appmessage/message.go
@@ -150,6 +150,8 @@ const (
 	CmdNotifyVirtualDaaScoreChangedRequestMessage
 	CmdNotifyVirtualDaaScoreChangedResponseMessage
 	CmdVirtualDaaScoreChangedNotificationMessage
+	CmdGetBalancesByAddressesRequestMessage
+	CmdGetBalancesByAddressesResponseMessage
 )
 
 // ProtocolMessageCommandToString maps all MessageCommands to their string representation
@@ -274,6 +276,8 @@ var RPCMessageCommandToString = map[MessageCommand]string{
 	CmdNotifyVirtualDaaScoreChangedRequestMessage:                 "NotifyVirtualDaaScoreChangedRequest",
 	CmdNotifyVirtualDaaScoreChangedResponseMessage:                "NotifyVirtualDaaScoreChangedResponse",
 	CmdVirtualDaaScoreChangedNotificationMessage:                  "VirtualDaaScoreChangedNotification",
+	CmdGetBalancesByAddressesRequestMessage:                       "GetBalancesByAddressesRequest",
+	CmdGetBalancesByAddressesResponseMessage:                      "GetBalancesByAddressesResponse",
 }
 
 // Message is an interface that describes a kaspa message. A type that

diff --git a/app/appmessage/rpc_get_balances_by_addresses.go b/app/appmessage/rpc_get_balances_by_addresses.go
@@ -0,0 +1,47 @@
+package appmessage
+
+// GetBalancesByAddressesRequestMessage is an appmessage corresponding to
+// its respective RPC message
+type GetBalancesByAddressesRequestMessage struct {
+	baseMessage
+	Addresses []string
+}
+
+// Command returns the protocol command string for the message
+func (msg *GetBalancesByAddressesRequestMessage) Command() MessageCommand {
+	return CmdGetBalancesByAddressesRequestMessage
+}
+
+// NewGetBalancesByAddressesRequest returns a instance of the message
+func NewGetBalancesByAddressesRequest(addresses []string) *GetBalancesByAddressesRequestMessage {
+	return &GetBalancesByAddressesRequestMessage{
+		Addresses: addresses,
+	}
+}
+
+// BalancesByAddressesEntry represents the balance of some address
+type BalancesByAddressesEntry struct {
+	Address string
+	Balance uint64
+}
+
+// GetBalancesByAddressesResponseMessage is an appmessage corresponding to
+// its respective RPC message
+type GetBalancesByAddressesResponseMessage struct {
+	baseMessage
+	Entries []*BalancesByAddressesEntry
+
+	Error *RPCError
+}
+
+// Command returns the protocol command string for the message
+func (msg *GetBalancesByAddressesResponseMessage) Command() MessageCommand {
+	return CmdGetBalancesByAddressesResponseMessage
+}
+
+// NewGetBalancesByAddressesResponse returns an instance of the message
+func NewGetBalancesByAddressesResponse(entries []*BalancesByAddressesEntry) *GetBalancesByAddressesResponseMessage {
+	return &GetBalancesByAddressesResponseMessage{
+		Entries: entries,
+	}
+}
diff --git a/app/rpc/rpc.go b/app/rpc/rpc.go
@@ -38,6 +38,7 @@ var handlers = map[appmessage.MessageCommand]handler{
 	appmessage.CmdNotifyUTXOsChangedRequestMessage:                          rpchandlers.HandleNotifyUTXOsChanged,
 	appmessage.CmdStopNotifyingUTXOsChangedRequestMessage:                   rpchandlers.HandleStopNotifyingUTXOsChanged,
 	appmessage.CmdGetUTXOsByAddressesRequestMessage:                         rpchandlers.HandleGetUTXOsByAddresses,
+	appmessage.CmdGetBalancesByAddressesRequestMessage:                      rpchandlers.HandleGetBalancesByAddresses,
 	appmessage.CmdGetVirtualSelectedParentBlueScoreRequestMessage:           rpchandlers.HandleGetVirtualSelectedParentBlueScore,
 	appmessage.CmdNotifyVirtualSelectedParentBlueScoreChangedRequestMessage: rpchandlers.HandleNotifyVirtualSelectedParentBlueScoreChanged,
 	appmessage.CmdBanRequestMessage:                                         rpchandlers.HandleBan,

diff --git a/app/rpc/rpchandlers/get_balance_by_address.go b/app/rpc/rpchandlers/get_balance_by_address.go
@@ -6,6 +6,7 @@ import (
 	"github.com/kaspanet/kaspad/domain/consensus/utils/txscript"
 	"github.com/kaspanet/kaspad/infrastructure/network/netadapter/router"
 	"github.com/kaspanet/kaspad/util"
+	"github.com/pkg/errors"
 )
 
 // HandleGetBalanceByAddress handles the respectively named RPC command
@@ -18,30 +19,39 @@ func HandleGetBalanceByAddress(context *rpccontext.Context, _ *router.Router, re
 
 	getBalanceByAddressRequest := request.(*appmessage.GetBalanceByAddressRequestMessage)
 
-	var balance uint64 = 0
-	addressString := getBalanceByAddressRequest.Address
-
-	address, err := util.DecodeAddress(addressString, context.Config.ActiveNetParams.Prefix)
+	balance, err := getBalanceByAddress(context, getBalanceByAddressRequest.Address)
 	if err != nil {
+		rpcError := &appmessage.RPCError{}
+		if !errors.As(err, rpcError) {
+			return nil, err
+		}
 		errorMessage := &appmessage.GetUTXOsByAddressesResponseMessage{}
-		errorMessage.Error = appmessage.RPCErrorf("Could decode address '%s': %s", addressString, err)
+		errorMessage.Error = rpcError
 		return errorMessage, nil
 	}
 
+	response := appmessage.NewGetBalanceByAddressResponse(balance)
+	return response, nil
+}
+
+func getBalanceByAddress(context *rpccontext.Context, addressString string) (uint64, error) {
+	address, err := util.DecodeAddress(addressString, context.Config.ActiveNetParams.Prefix)
+	if err != nil {
+		return 0, appmessage.RPCErrorf("Couldn't decode address '%s': %s", addressString, err)
+	}
+
 	scriptPublicKey, err := txscript.PayToAddrScript(address)
 	if err != nil {
-		errorMessage := &appmessage.GetUTXOsByAddressesResponseMessage{}
-		errorMessage.Error = appmessage.RPCErrorf("Could not create a scriptPublicKey for address '%s': %s", addressString, err)
-		return errorMessage, nil
+		return 0, appmessage.RPCErrorf("Could not create a scriptPublicKey for address '%s': %s", addressString, err)
 	}
 	utxoOutpointEntryPairs, err := context.UTXOIndex.UTXOs(scriptPublicKey)
 	if err != nil {
-		return nil, err
+		return 0, err
 	}
+
+	balance := uint64(0)
 	for _, utxoOutpointEntryPair := range utxoOutpointEntryPairs {
 		balance += utxoOutpointEntryPair.Amount()
 	}
-
-	response := appmessage.NewGetBalanceByAddressResponse(balance)
-	return response, nil
+	return balance, nil
 }
diff --git a/app/rpc/rpchandlers/get_balances_by_addresses.go b/app/rpc/rpchandlers/get_balances_by_addresses.go
@@ -0,0 +1,41 @@
+package rpchandlers
+
+import (
+	"github.com/kaspanet/kaspad/app/appmessage"
+	"github.com/kaspanet/kaspad/app/rpc/rpccontext"
+	"github.com/kaspanet/kaspad/infrastructure/network/netadapter/router"
+	"github.com/pkg/errors"
+)
+
+// HandleGetBalancesByAddresses handles the respectively named RPC command
+func HandleGetBalancesByAddresses(context *rpccontext.Context, _ *router.Router, request appmessage.Message) (appmessage.Message, error) {
+	if !context.Config.UTXOIndex {
+		errorMessage := &appmessage.GetBalancesByAddressesResponseMessage{}
+		errorMessage.Error = appmessage.RPCErrorf("Method unavailable when kaspad is run without --utxoindex")
+		return errorMessage, nil
+	}
+
+	getBalancesByAddressesRequest := request.(*appmessage.GetBalancesByAddressesRequestMessage)
+
+	allEntries := make([]*appmessage.BalancesByAddressesEntry, len(getBalancesByAddressesRequest.Addresses))
+	for i, address := range getBalancesByAddressesRequest.Addresses {
+		balance, err := getBalanceByAddress(context, address)
+
+		if err != nil {
+			rpcError := &appmessage.RPCError{}
+			if !errors.As(err, rpcError) {
+				return nil, err
+			}
+			errorMessage := &appmessage.GetUTXOsByAddressesResponseMessage{}
+			errorMessage.Error = rpcError
+			return errorMessage, nil
+		}
+		allEntries[i] = &appmessage.BalancesByAddressesEntry{
+			Address: address,
+			Balance: balance,
+		}
+	}
+
+	response := appmessage.NewGetBalancesByAddressesResponse(allEntries)
+	return response, nil
+}
diff --git a/cmd/kaspawallet/daemon/server/balance.go b/cmd/kaspawallet/daemon/server/balance.go
@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+
 	"github.com/kaspanet/kaspad/cmd/kaspawallet/daemon/pb"
 	"github.com/kaspanet/kaspad/cmd/kaspawallet/libkaspawallet"
 )
@@ -21,7 +22,7 @@ func (s *server) GetBalance(_ context.Context, _ *pb.GetBalanceRequest) (*pb.Get
 	maturity := s.params.BlockCoinbaseMaturity
 
 	balancesMap := make(balancesMapType, 0)
-	for _, entry := range s.utxos {
+	for _, entry := range s.utxosSortedByAmount {
 		amount := entry.UTXOEntry.Amount()
 		address := entry.address
 		balances, ok := balancesMap[address]

diff --git a/cmd/kaspawallet/daemon/server/create_unsigned_transaction.go b/cmd/kaspawallet/daemon/server/create_unsigned_transaction.go
@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+
 	"github.com/kaspanet/kaspad/cmd/kaspawallet/daemon/pb"
 	"github.com/kaspanet/kaspad/cmd/kaspawallet/libkaspawallet"
 	"github.com/kaspanet/kaspad/domain/consensus/utils/constants"
@@ -17,7 +18,7 @@ func (s *server) CreateUnsignedTransaction(_ context.Context, request *pb.Create
 		return nil, errors.New("server is not synced")
 	}
 
-	err := s.refreshExistingUTXOs()
+	err := s.refreshUTXOs()
 	if err != nil {
 		return nil, err
 	}
@@ -66,7 +67,7 @@ func (s *server) selectUTXOs(spendAmount uint64, feePerInput uint64) (
 		return nil, 0, err
 	}
 
-	for _, utxo := range s.utxos {
+	for _, utxo := range s.utxosSortedByAmount {
 		if !isUTXOSpendable(utxo, dagInfo.VirtualDAAScore, s.params.BlockCoinbaseMaturity) {
 			continue
 		}

diff --git a/cmd/kaspawallet/daemon/server/server.go b/cmd/kaspawallet/daemon/server/server.go
@@ -2,15 +2,15 @@ package server
 
 import (
 	"fmt"
-	"github.com/kaspanet/kaspad/util/profiling"
 	"net"
 	"os"
 	"sync"
 	"time"
 
+	"github.com/kaspanet/kaspad/util/profiling"
+
 	"github.com/kaspanet/kaspad/cmd/kaspawallet/daemon/pb"
 	"github.com/kaspanet/kaspad/cmd/kaspawallet/keys"
-	"github.com/kaspanet/kaspad/domain/consensus/model/externalapi"
 	"github.com/kaspanet/kaspad/domain/dagconfig"
 	"github.com/kaspanet/kaspad/infrastructure/network/rpcclient"
 	"github.com/kaspanet/kaspad/infrastructure/os/signal"
@@ -26,11 +26,12 @@ type server struct {
 	rpcClient *rpcclient.RPCClient
 	params    *dagconfig.Params
 
-	lock               sync.RWMutex
-	utxos              map[externalapi.DomainOutpoint]*walletUTXO
-	nextSyncStartIndex uint32
-	keysFile           *keys.File
-	shutdown           chan struct{}
+	lock                sync.RWMutex
+	utxosSortedByAmount []*walletUTXO
+	nextSyncStartIndex  uint32
+	keysFile            *keys.File
+	shutdown            chan struct{}
+	addressSet          walletAddressSet
 }
 
 // Start starts the kaspawalletd server
@@ -61,12 +62,13 @@ func Start(params *dagconfig.Params, listen, rpcServer string, keysFilePath stri
 	}
 
 	serverInstance := &server{
-		rpcClient:          rpcClient,
-		params:             params,
-		utxos:              make(map[externalapi.DomainOutpoint]*walletUTXO),
-		nextSyncStartIndex: 0,
-		keysFile:           keysFile,
-		shutdown:           make(chan struct{}),
+		rpcClient:           rpcClient,
+		params:              params,
+		utxosSortedByAmount: []*walletUTXO{},
+		nextSyncStartIndex:  0,
+		keysFile:            keysFile,
+		shutdown:            make(chan struct{}),
+		addressSet:          make(walletAddressSet),
 	}
 
 	spawn("serverInstance.sync", func() {