Merge pull request #2752 from OffchainLabs/destroy_val_entry

Do not use SetFinalizer to remove entries from preimageResolvers

util/containers/syncmap.go

@@ -22,3 +22,13 @@ func (m *SyncMap[K, V]) Store(key K, val V) {
 func (m *SyncMap[K, V]) Delete(key K) {
 	m.internal.Delete(key)
 }
+
+// Only used for testing
+func (m *SyncMap[K, V]) Keys() []K {
+	s := make([]K, 0)
+	m.internal.Range(func(k, v interface{}) bool {
+		s = append(s, k.(K))
+		return true
+	})
+	return s
+}

validator/server_arb/machine.go

@@ -51,16 +51,32 @@ type MachineInterface interface {
 type ArbitratorMachine struct {
 	mutex     sync.Mutex // needed because go finalizers don't synchronize (meaning they aren't thread safe)
 	ptr       *C.struct_Machine
-	contextId *int64 // has a finalizer attached to remove the preimage resolver from the global map
-	frozen    bool   // does not allow anything that changes machine state, not cloned with the machine
+	contextId *int64
+	frozen    bool // does not allow anything that changes machine state, not cloned with the machine
 }
 
 // Assert that ArbitratorMachine implements MachineInterface
 var _ MachineInterface = (*ArbitratorMachine)(nil)
 
-var preimageResolvers containers.SyncMap[int64, GoPreimageResolver]
+var preimageResolvers containers.SyncMap[int64, goPreimageResolverWithRefCounter]
 var lastPreimageResolverId atomic.Int64 // atomic
 
+func dereferenceContextId(contextId *int64) {
+	if contextId != nil {
+		resolverWithRefCounter, ok := preimageResolvers.Load(*contextId)
+		if !ok {
+			panic(fmt.Sprintf("dereferenceContextId: resolver with ref counter not found, contextId: %v", *contextId))
+		}
+
+		refCount := resolverWithRefCounter.refCounter.Add(-1)
+		if refCount < 0 {
+			panic(fmt.Sprintf("dereferenceContextId: ref counter is negative, contextId: %v", *contextId))
+		} else if refCount == 0 {
+			preimageResolvers.Delete(*contextId)
+		}
+	}
+}
+
 // Any future calls to this machine will result in a panic
 func (m *ArbitratorMachine) Destroy() {
 	m.mutex.Lock()
@@ -71,11 +87,9 @@ func (m *ArbitratorMachine) Destroy() {
 		// We no longer need a finalizer
 		runtime.SetFinalizer(m, nil)
 	}
-	m.contextId = nil
-}
 
-func freeContextId(context *int64) {
-	preimageResolvers.Delete(*context)
+	dereferenceContextId(m.contextId)
+	m.contextId = nil
 }
 
 func machineFromPointer(ptr *C.struct_Machine) *ArbitratorMachine {
@@ -112,6 +126,16 @@ func (m *ArbitratorMachine) Clone() *ArbitratorMachine {
 	defer m.mutex.Unlock()
 	newMach := machineFromPointer(C.arbitrator_clone_machine(m.ptr))
 	newMach.contextId = m.contextId
+
+	if m.contextId != nil {
+		resolverWithRefCounter, ok := preimageResolvers.Load(*m.contextId)
+		if ok {
+			resolverWithRefCounter.refCounter.Add(1)
+		} else {
+			panic(fmt.Sprintf("Clone: resolver with ref counter not found, contextId: %v", *m.contextId))
+		}
+	}
+
 	return newMach
 }
 
@@ -350,19 +374,24 @@ func (m *ArbitratorMachine) AddDelayedInboxMessage(index uint64, data []byte) er
 }
 
 type GoPreimageResolver = func(arbutil.PreimageType, common.Hash) ([]byte, error)
+type goPreimageResolverWithRefCounter struct {
+	resolver   GoPreimageResolver
+	refCounter *atomic.Int64
+}
 
 //export preimageResolver
 func preimageResolver(context C.size_t, ty C.uint8_t, ptr unsafe.Pointer) C.ResolvedPreimage {
 	var hash common.Hash
 	input := (*[1 << 30]byte)(ptr)[:32]
 	copy(hash[:], input)
-	resolver, ok := preimageResolvers.Load(int64(context))
+	resolverWithRefCounter, ok := preimageResolvers.Load(int64(context))
 	if !ok {
+		log.Error("preimageResolver: resolver with ref counter not found", "context", int64(context))
 		return C.ResolvedPreimage{
 			len: -1,
 		}
 	}
-	preimage, err := resolver(arbutil.PreimageType(ty), hash)
+	preimage, err := resolverWithRefCounter.resolver(arbutil.PreimageType(ty), hash)
 	if err != nil {
 		log.Error("preimage resolution failed", "err", err)
 		return C.ResolvedPreimage{
@@ -382,10 +411,18 @@ func (m *ArbitratorMachine) SetPreimageResolver(resolver GoPreimageResolver) err
 	if m.frozen {
 		return errors.New("machine frozen")
 	}
+	dereferenceContextId(m.contextId)
+
 	id := lastPreimageResolverId.Add(1)
-	preimageResolvers.Store(id, resolver)
+	refCounter := atomic.Int64{}
+	refCounter.Store(1)
+	resolverWithRefCounter := goPreimageResolverWithRefCounter{
+		resolver:   resolver,
+		refCounter: &refCounter,
+	}
+	preimageResolvers.Store(id, resolverWithRefCounter)
+
 	m.contextId = &id
-	runtime.SetFinalizer(m.contextId, freeContextId)
 	C.arbitrator_set_context(m.ptr, u64(id))
 	return nil
 }

validator/server_arb/machine_test.go

@@ -0,0 +1,93 @@
+// Copyright 2021-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+package server_arb
+
+import (
+	"path"
+	"reflect"
+	"runtime"
+	"sort"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/offchainlabs/nitro/arbutil"
+	"github.com/offchainlabs/nitro/util/testhelpers"
+)
+
+func TestEntriesAreDeletedFromPreimageResolversGlobalMap(t *testing.T) {
+	resolver := func(arbutil.PreimageType, common.Hash) ([]byte, error) {
+		return nil, nil
+	}
+
+	sortedKeys := func() []int64 {
+		keys := preimageResolvers.Keys()
+		sort.Slice(keys, func(i, j int) bool {
+			return keys[i] < keys[j]
+		})
+		return keys
+	}
+
+	// clear global map before running test
+	preimageKeys := sortedKeys()
+	for _, key := range preimageKeys {
+		preimageResolvers.Delete(key)
+	}
+
+	_, filename, _, _ := runtime.Caller(0)
+	wasmDir := path.Join(path.Dir(filename), "../../arbitrator/prover/test-cases/")
+	wasmPath := path.Join(wasmDir, "global-state.wasm")
+	modulePaths := []string{path.Join(wasmDir, "global-state-wrapper.wasm")}
+
+	machine1, err := LoadSimpleMachine(wasmPath, modulePaths, true)
+	testhelpers.RequireImpl(t, err)
+	err = machine1.SetPreimageResolver(resolver)
+	testhelpers.RequireImpl(t, err)
+
+	machine2, err := LoadSimpleMachine(wasmPath, modulePaths, true)
+	testhelpers.RequireImpl(t, err)
+	err = machine2.SetPreimageResolver(resolver)
+	testhelpers.RequireImpl(t, err)
+
+	machine1Clone1 := machine1.Clone()
+	machine1Clone2 := machine1.Clone()
+
+	checkKeys := func(expectedKeys []int64, scenario string) {
+		keys := sortedKeys()
+		if !reflect.DeepEqual(keys, expectedKeys) {
+			t.Fatal("Unexpected preimageResolversKeys got", keys, "expected", expectedKeys, "scenario", scenario)
+		}
+	}
+
+	machine1ContextId := *machine1.contextId
+	machine2ContextId := *machine2.contextId
+
+	checkKeys([]int64{machine1ContextId, machine2ContextId}, "initial")
+
+	// the machine's contextId should change when setting preimage resolver for the second time,
+	// and the entry for the old context id should be deleted
+	err = machine2.SetPreimageResolver(resolver)
+	testhelpers.RequireImpl(t, err)
+	if machine2ContextId == *machine2.contextId {
+		t.Fatal("Context id didn't change after setting preimage resolver for the second time")
+	}
+	machine2ContextId = *machine2.contextId
+	checkKeys([]int64{machine1ContextId, machine2ContextId}, "after setting preimage resolver for machine2 for the second time")
+
+	machine1Clone1.Destroy()
+	checkKeys([]int64{machine1ContextId, machine2ContextId}, "after machine1Clone1 is destroyed")
+
+	machine1.Destroy()
+	checkKeys([]int64{machine1ContextId, machine2ContextId}, "after machine1 is destroyed")
+
+	// it is possible to destroy the same machine multiple times
+	machine1.Destroy()
+	checkKeys([]int64{machine1ContextId, machine2ContextId}, "after machine1 is destroyed again")
+
+	// entry for machine1ContextId should be deleted only after machine1 and all its clones are destroyed
+	machine1Clone2.Destroy()
+	checkKeys([]int64{machine2ContextId}, "after machine1Clone2 is destroyed")
+
+	machine2.Destroy()
+	checkKeys([]int64{}, "after machine2 is destroyed")
+}