From 62309d340492567f4f02f8cefddcdacf6a237828 Mon Sep 17 00:00:00 2001
From: beer-1 <beer-1@codingcats.xyz>
Date: Thu, 26 Sep 2024 18:17:26 +0900
Subject: [PATCH] use atomic

---
 x/move/keeper/keeper.go | 9 +++------
 x/move/keeper/vmpool.go | 8 +++-----
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/x/move/keeper/keeper.go b/x/move/keeper/keeper.go
index 0498babc..652fd31b 100644
--- a/x/move/keeper/keeper.go
+++ b/x/move/keeper/keeper.go
@@ -3,7 +3,6 @@ package keeper
 import (
 	"context"
 	"errors"
-	"sync"
 
 	"golang.org/x/sync/semaphore"
 
@@ -46,9 +45,8 @@ type Keeper struct {
 
 	// TODO - remove after loader v2
 	moveVMs         []types.VMEngine
-	moveVMMutx      *sync.Mutex
+	moveVMIdx       *uint64
 	moveVMSemaphore *semaphore.Weighted
-	moveVMIdx       *int
 
 	feeCollector string
 	authority    string
@@ -92,7 +90,7 @@ func NewKeeper(
 	}
 
 	vmCount := 10
-	moveVMIdx := int(0)
+	moveVMIdx := uint64(0)
 	vms := make([]types.VMEngine, vmCount)
 	for i := 0; i < vmCount; i++ {
 		moveVM, err := vm.NewVM(vmtypes.InitiaVMConfig{
@@ -117,9 +115,8 @@ func NewKeeper(
 		grpcRouter:          grpcRouter,
 		config:              moveConfig,
 		moveVMs:             vms,
-		moveVMMutx:          new(sync.Mutex),
-		moveVMSemaphore:     semaphore.NewWeighted(int64(vmCount)),
 		moveVMIdx:           &moveVMIdx,
+		moveVMSemaphore:     semaphore.NewWeighted(int64(vmCount)),
 		distrKeeper:         distrKeeper,
 		StakingKeeper:       stakingKeeper,
 		RewardKeeper:        rewardKeeper,
diff --git a/x/move/keeper/vmpool.go b/x/move/keeper/vmpool.go
index d4675a41..a5425486 100644
--- a/x/move/keeper/vmpool.go
+++ b/x/move/keeper/vmpool.go
@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"context"
+	"sync/atomic"
 
 	"github.com/initia-labs/initia/x/move/types"
 )
@@ -12,11 +13,8 @@ func (k Keeper) acquireVM(ctx context.Context) (vm types.VMEngine) {
 		panic(err)
 	}
 
-	k.moveVMMutx.Lock()
-	idx := *k.moveVMIdx
-	*k.moveVMIdx = (idx + 1) % len(k.moveVMs)
-	vm = k.moveVMs[idx]
-	k.moveVMMutx.Unlock()
+	idx := atomic.AddUint64(k.moveVMIdx, 1)
+	vm = k.moveVMs[idx%uint64(len(k.moveVMs))]
 
 	return
 }