From 243c66833be6a4375fb81c7f0b8a305c12455d72 Mon Sep 17 00:00:00 2001
From: Jeffrey Wilcke <jeffrey@ethereum.org>
Date: Fri, 2 Dec 2016 17:30:17 +0100
Subject: [PATCH] core/vm: added RunWithContext to the EVM for cancellation

Implemented RunWithContext which allows the execution of the EVM to be
halted mid-flight. This is useful when using the EVM unmetered or when
using the Light Client, which requires control over the EVM run loop.
---
 core/vm/vm.go | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/core/vm/vm.go b/core/vm/vm.go
index 11d7b47d01..f49665723c 100644
--- a/core/vm/vm.go
+++ b/core/vm/vm.go
@@ -17,8 +17,10 @@
 package vm
 
 import (
+	"context"
 	"fmt"
 	"math/big"
+	"sync/atomic"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -54,6 +56,10 @@ type EVM struct {
 	jumpTable vmJumpTable
 	cfg       Config
 	gasTable  params.GasTable
+
+	// done is an atomic int and is used for
+	// cancellation during RunWithContext.
+	done int32
 }
 
 // New returns a new instance of the EVM.
@@ -66,6 +72,18 @@ func New(env *Environment, cfg Config) *EVM {
 	}
 }
 
+// RunWithContext allows the EVM to be ran with a cancellation method by passing in a context.Context. The EVM
+// behaves exactly the same as an EVM without a context.
+//
+// RunWithContext is only used for the initial call and shouldn't be called more than once.
+func (evm *EVM) RunWithContext(ctx context.Context, contract *Contract, input []byte) (ret []byte, err error) {
+	go func() {
+		<-ctx.Done()
+		atomic.StoreInt32(&evm.done, 1)
+	}()
+	return evm.Run(contract, input)
+}
+
 // Run loops and evaluates the contract's code with the given input data
 func (evm *EVM) Run(contract *Contract, input []byte) (ret []byte, err error) {
 	evm.env.Depth++
@@ -113,7 +131,11 @@ func (evm *EVM) Run(contract *Contract, input []byte) (ret []byte, err error) {
 		}()
 	}
 
-	for {
+	// The EVM main run loop (contextual). This loop runs until either an
+	// explicit STOP, RETURN or SUICIDE is executed, an error accured during
+	// the execution of one of the operations or until the evm.done is set by
+	// the parent context.Context.
+	for atomic.LoadInt32(&evm.done) == 0 {
 		// Get the memory location of pc
 		op = contract.GetOp(pc)
 
@@ -166,6 +188,7 @@ func (evm *EVM) Run(contract *Contract, input []byte) (ret []byte, err error) {
 			pc++
 		}
 	}
+	return nil, nil
 }
 
 // RunPrecompile runs and evaluate the output of a precompiled contract defined in contracts.go