Added UTXOIDs to the getBalance call to enable atomic reads

vms/avm/service.go

@@ -188,7 +188,8 @@ type GetBalanceArgs struct {
 
 // GetBalanceReply defines the GetBalance replies returned from the API
 type GetBalanceReply struct {
-	Balance json.Uint64 `json:"balance"`
+	Balance json.Uint64  `json:"balance"`
+	UTXOIDs []ava.UTXOID `json:"utxoIDs"`
 }
 
 // GetBalance returns the amount of an asset that an address at least partially owns
@@ -217,18 +218,21 @@ func (service *Service) GetBalance(r *http.Request, args *GetBalanceArgs, reply
 	}
 
 	for _, utxo := range utxos {
-		if utxo.AssetID().Equals(assetID) {
-			transferable, ok := utxo.Out.(ava.Transferable)
-			if !ok {
-				continue
-			}
-			amt, err := safemath.Add64(transferable.Amount(), uint64(reply.Balance))
-			if err != nil {
-				return err
-			}
-			reply.Balance = json.Uint64(amt)
+		if !utxo.AssetID().Equals(assetID) {
+			continue
 		}
+		transferable, ok := utxo.Out.(ava.Transferable)
+		if !ok {
+			continue
+		}
+		amt, err := safemath.Add64(transferable.Amount(), uint64(reply.Balance))
+		if err != nil {
+			return err
+		}
+		reply.Balance = json.Uint64(amt)
+		reply.UTXOIDs = append(reply.UTXOIDs, utxo.UTXOID)
 	}
+
 	return nil
 }
 

vms/avm/service_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/ava-labs/gecko/snow/choices"
+	"github.com/stretchr/testify/assert"
 
 	"github.com/ava-labs/gecko/database/memdb"
 	"github.com/ava-labs/gecko/ids"
@@ -107,6 +108,27 @@ func TestServiceGetTxStatus(t *testing.T) {
 	}
 }
 
+func TestServiceGetBalance(t *testing.T) {
+	genesisBytes, vm, s := setup(t)
+	defer ctx.Lock.Unlock()
+	defer vm.Shutdown()
+
+	genesisTx := GetFirstTxFromGenesisTest(genesisBytes, t)
+	assetID := genesisTx.ID()
+	addr := keys[0].PublicKey().Address()
+
+	balanceArgs := &GetBalanceArgs{
+		Address: fmt.Sprintf("%s-%s", vm.ctx.ChainID, addr),
+		AssetID: assetID.String(),
+	}
+	balanceReply := &GetBalanceReply{}
+	err := s.GetBalance(nil, balanceArgs, balanceReply)
+	assert.NoError(t, err)
+	assert.Equal(t, uint64(balanceReply.Balance), uint64(300000))
+
+	assert.Len(t, balanceReply.UTXOIDs, 4, "should have only returned four utxoIDs")
+}
+
 func TestServiceGetUTXOsInvalidAddress(t *testing.T) {
 	_, vm, s := setup(t)
 	defer ctx.Lock.Unlock()