fix(wallet): Fix race condition in concurrent wallet update (#1163)

server/core_wallet.go

@@ -119,7 +119,7 @@ func updateWallets(ctx context.Context, logger *zap.Logger, tx pgx.Tx, updates [
 		ids = append(ids, update.UserID)
 	}
 
-	initialQuery := "SELECT id, wallet FROM users WHERE id = ANY($1::UUID[])"
+	initialQuery := "SELECT id, wallet FROM users WHERE id = ANY($1::UUID[]) FOR UPDATE"
 
 	// Select the wallets from the DB and decode them.
 	wallets := make(map[string]map[string]int64, len(updates))

server/core_wallet_test.go

@@ -17,6 +17,8 @@ package server
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/gofrs/uuid/v5"
@@ -77,13 +79,27 @@ func TestUpdateWalletSingleUser(t *testing.T) {
 		t.Fatalf("error creating user: %v", err.Error())
 	}
 
-	for _, val := range values {
+	for _, val := range values[:len(values)/2] {
 		_, _, err := nk.WalletUpdate(context.Background(), userID, map[string]int64{"value": val}, nil, true)
 		if err != nil {
 			t.Fatalf("error updating wallet: %v", err.Error())
 		}
 	}
 
+	var wg sync.WaitGroup
+	for _, val := range values[len(values)/2:] {
+		v := val
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_, _, err := nk.WalletUpdate(context.Background(), userID, map[string]int64{"value": v}, nil, true)
+			if err != nil {
+				panic(fmt.Sprintf("error updating wallet: %v", err.Error()))
+			}
+		}()
+	}
+	wg.Wait()
+
 	account, err := GetAccount(context.Background(), logger, db, nil, uuid.FromStringOrNil(userID))
 	if err != nil {
 		t.Fatalf("error getting user: %v", err.Error())