fix: release vm right after execution finished (#274)

* fix to release vm right after execution finished * fix to charge gas first before returnning errorg
initia-labs · Sep 26, 2024 · 30d4e29 · 30d4e29
1 parent 224e733
commit 30d4e29
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 48 deletions.
diff --git a/x/move/keeper/genesis.go b/x/move/keeper/genesis.go
@@ -46,13 +46,10 @@ func (k Keeper) Initialize(
 		_allowedPublishers[i] = addr
 	}
 
-	// TODO - remove this after loader v2 is installed
-	vm := k.acquireVM(ctx)
-	defer k.releaseVM()
-
-	// The default upgrade policy is compatible when it's not set.
 	vmStore := types.NewVMStore(ctx, k.VMStore)
-	execRes, err := vm.Initialize(vmStore, api, env, vmtypes.NewModuleBundle(modules...), _allowedPublishers)
+	execRes, err := execVM(ctx, k, func(vm types.VMEngine) (vmtypes.ExecutionResult, error) {
+		return vm.Initialize(vmStore, api, env, vmtypes.NewModuleBundle(modules...), _allowedPublishers)
+	})
 	if err != nil {
 		return err
 	}

diff --git a/x/move/keeper/handler.go b/x/move/keeper/handler.go
@@ -164,10 +164,6 @@ func (k Keeper) executeEntryFunction(
 	args [][]byte,
 	isJSON bool,
 ) error {
-	// TODO - remove this after loader v2 is installed
-	vm := k.acquireVM(ctx)
-	defer k.releaseVM()
-
 	payload, err := types.BuildExecuteEntryFunctionPayload(
 		moduleAddr,
 		moduleName,
@@ -194,24 +190,24 @@ func (k Keeper) executeEntryFunction(
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	gasMeter := sdkCtx.GasMeter()
 	gasForRuntime := gasMeter.Limit() - gasMeter.GasConsumedToLimit()
-
 	if isSimulation(ctx) {
 		gasForRuntime = k.config.ContractSimulationGasLimit
 	}
 
 	// delegate gas metering to move vm
 	sdkCtx = sdkCtx.WithGasMeter(storetypes.NewInfiniteGasMeter())
 
-	// run vm
 	gasBalance := gasForRuntime
-	execRes, err := vm.ExecuteEntryFunction(
-		&gasBalance,
-		types.NewVMStore(sdkCtx, k.VMStore),
-		NewApi(k, sdkCtx),
-		types.NewEnv(sdkCtx, ac, ec),
-		senders,
-		payload,
-	)
+	execRes, err := execVM(ctx, k, func(vm types.VMEngine) (vmtypes.ExecutionResult, error) {
+		return vm.ExecuteEntryFunction(
+			&gasBalance,
+			types.NewVMStore(sdkCtx, k.VMStore),
+			NewApi(k, sdkCtx),
+			types.NewEnv(sdkCtx, ac, ec),
+			senders,
+			payload,
+		)
+	})
 
 	// consume gas first and check error
 	gasUsed := gasForRuntime - gasBalance
@@ -284,10 +280,6 @@ func (k Keeper) executeScript(
 	args [][]byte,
 	isJSON bool,
 ) error {
-	// TODO - remove this after loader v2 is installed
-	vm := k.acquireVM(ctx)
-	defer k.releaseVM()
-
 	// prepare payload
 	payload, err := types.BuildExecuteScriptPayload(
 		byteCodes,
@@ -313,24 +305,24 @@ func (k Keeper) executeScript(
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	gasMeter := sdkCtx.GasMeter()
 	gasForRuntime := gasMeter.Limit() - gasMeter.GasConsumedToLimit()
-
 	if isSimulation(ctx) {
 		gasForRuntime = k.config.ContractSimulationGasLimit
 	}
 
 	// delegate gas metering to move vm
 	sdkCtx = sdkCtx.WithGasMeter(storetypes.NewInfiniteGasMeter())
 
-	// run vm
 	gasBalance := gasForRuntime
-	execRes, err := vm.ExecuteScript(
-		&gasBalance,
-		types.NewVMStore(sdkCtx, k.VMStore),
-		NewApi(k, sdkCtx),
-		types.NewEnv(sdkCtx, ac, ec),
-		senders,
-		payload,
-	)
+	execRes, err := execVM(ctx, k, func(vm types.VMEngine) (vmtypes.ExecutionResult, error) {
+		return vm.ExecuteScript(
+			&gasBalance,
+			types.NewVMStore(sdkCtx, k.VMStore),
+			NewApi(k, sdkCtx),
+			types.NewEnv(sdkCtx, ac, ec),
+			senders,
+			payload,
+		)
+	})
 
 	// consume gas first and check error
 	gasUsed := gasForRuntime - gasBalance
@@ -565,10 +557,6 @@ func (k Keeper) executeViewFunction(
 	args [][]byte,
 	isJSON bool,
 ) (vmtypes.ViewOutput, uint64, error) {
-	// TODO - remove this after loader v2 is installed
-	vm := k.acquireVM(ctx)
-	defer k.releaseVM()
-
 	payload, err := types.BuildExecuteViewFunctionPayload(
 		moduleAddr,
 		moduleName,
@@ -598,20 +586,33 @@ func (k Keeper) executeViewFunction(
 	gasForRuntime := gasMeter.Limit() - gasMeter.GasConsumedToLimit()
 
 	gasBalance := gasForRuntime
-	viewRes, err := vm.ExecuteViewFunction(
-		&gasBalance,
-		types.NewVMStore(ctx, k.VMStore),
-		api,
-		env,
-		payload,
-	)
-	if err != nil {
-		return vmtypes.ViewOutput{}, 0, err
-	}
+	viewRes, err := execVM(ctx, k, func(vm types.VMEngine) (vmtypes.ViewOutput, error) {
+		return vm.ExecuteViewFunction(
+			&gasBalance,
+			types.NewVMStore(ctx, k.VMStore),
+			api,
+			env,
+			payload,
+		)
+	})
 
 	// consume gas first and check error
 	gasUsed := gasForRuntime - gasBalance
 	gasMeter.ConsumeGas(gasUsed, "view; move runtime")
+	if err != nil {
+		return vmtypes.ViewOutput{}, gasUsed, err
+	}
 
 	return viewRes, gasUsed, nil
 }
+
+// execVM runs vm in separate function statement to release right after execution
+// to avoid deadlock even if the function panics
+//
+// TODO - remove this after loader v2 is installed
+func execVM[T any](ctx context.Context, k Keeper, f func(types.VMEngine) (T, error)) (T, error) {
+	vm := k.acquireVM(ctx)
+	defer k.releaseVM()
+
+	return f(vm)
+}