wazevo: concurrent local wasm function compilation (#2381)

Signed-off-by: David Desmarais-Michaud <desmarais.michaud.david@gmail.com>
Co-authored-by: Nuno Cruces <ncruces@users.noreply.github.com>

cmd/wazero/wazero.go

@@ -91,6 +91,7 @@ func doCompile(args []string, stdErr io.Writer) int {
 	}
 
 	cacheDir := cacheDirFlag(flags)
+	workers := workersFlag(flags)
 
 	_ = flags.Parse(args)
 
@@ -130,6 +131,14 @@ func doCompile(args []string, stdErr io.Writer) int {
 	}
 
 	ctx := context.Background()
+
+	if maxProcs := runtime.GOMAXPROCS(0); *workers > maxProcs {
+		fmt.Fprintf(stdErr, "warning: desired number of compilation workers (%d) greater than GOMAXPROCS (%d) defaulting to GOMAXPROCS\n", *workers, maxProcs)
+		*workers = maxProcs
+	}
+
+	ctx = experimental.WithCompilationWorkers(ctx, *workers)
+
 	rt := wazero.NewRuntimeWithConfig(ctx, c)
 	defer rt.Close(ctx)
 
@@ -207,6 +216,7 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int {
 	}
 
 	cacheDir := cacheDirFlag(flags)
+	workers := workersFlag(flags)
 
 	_ = flags.Parse(args)
 
@@ -281,6 +291,18 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int {
 		rtc = rtc.WithCompilationCache(cache)
 	}
 
+	// Save the current state of the context for compilation.
+	// The -timeout flag is documented as only affecting the runtime of the module
+	// To honor this we must not pass the regular ctx as it will be potentially reassigned from this point onwards.
+	compilationCtx := ctx
+
+	if maxProcs := runtime.GOMAXPROCS(0); *workers > maxProcs {
+		fmt.Fprintf(stdErr, "warning: desired number of compilation workers (%d) greater than GOMAXPROCS (%d) defaulting to GOMAXPROCS\n", *workers, maxProcs)
+		*workers = maxProcs
+	}
+
+	compilationCtx = experimental.WithCompilationWorkers(compilationCtx, *workers)
+
 	if timeout > 0 {
 		newCtx, cancel := context.WithTimeout(ctx, timeout)
 		ctx = newCtx
@@ -317,7 +339,7 @@ func doRun(args []string, stdOut io.Writer, stdErr logging.Writer) int {
 		conf = conf.WithEnv(env[i], env[i+1])
 	}
 
-	guest, err := rt.CompileModule(ctx, wasm)
+	guest, err := rt.CompileModule(compilationCtx, wasm)
 	if err != nil {
 		fmt.Fprintf(stdErr, "error compiling wasm binary: %v\n", err)
 		return 1
@@ -463,6 +485,11 @@ func cacheDirFlag(flags *flag.FlagSet) *string {
 		"Contents are re-used for the same version of wazero.")
 }
 
+func workersFlag(flags *flag.FlagSet) *int {
+	return flags.Int("workers", 1, "(experimental) Number of desired compilation workers. "+
+		"Increasing this value may improve compilation speed at the cost of higher memory usage.")
+}
+
 func maybeUseCacheDir(cacheDir *string, stdErr io.Writer) (int, wazero.CompilationCache) {
 	if dir := *cacheDir; dir != "" {
 		if cache, err := wazero.NewCompilationCacheWithDir(dir); err != nil {

cmd/wazero/wazero_test.go

@@ -10,6 +10,7 @@ import (
 	"os/exec"
 	"path"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"testing"
 
@@ -112,6 +113,10 @@ func TestCompile(t *testing.T) {
 				require.True(t, len(entries) > 0)
 			},
 		},
+		{
+			name:       "workers equal max procs",
+			wazeroOpts: []string{fmt.Sprintf("--workers=%d", runtime.GOMAXPROCS(0))},
+		},
 		{
 			name:       "enable cpu profiling",
 			wazeroOpts: []string{"-cpuprofile=" + cpuProfile},

experimental/compilationworkers.go

@@ -0,0 +1,19 @@
+package experimental
+
+import (
+	"context"
+
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
+)
+
+// WithCompilationWorkers sets the desired number of compilation workers.
+func WithCompilationWorkers(ctx context.Context, workers int) context.Context {
+	return context.WithValue(ctx, expctxkeys.CompilationWorkers{}, workers)
+}
+
+// GetCompilationWorkers returns the desired number of compilation workers.
+// The minimum value returned is 1.
+func GetCompilationWorkers(ctx context.Context) int {
+	workers, _ := ctx.Value(expctxkeys.CompilationWorkers{}).(int)
+	return max(workers, 1)
+}

internal/engine/wazevo/engine.go

@@ -8,6 +8,7 @@ import (
 	"runtime"
 	"sort"
 	"sync"
+	"sync/atomic"
 	"unsafe"
 
 	"github.com/tetratelabs/wazero/api"
@@ -206,6 +207,10 @@ func (exec *executables) compileEntryPreambles(m *wasm.Module, machine backend.M
 }
 
 func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listeners []experimental.FunctionListener, ensureTermination bool) (*compiledModule, error) {
+	if module.IsHostModule {
+		return e.compileHostModule(ctx, module, listeners)
+	}
+
 	withListener := len(listeners) > 0
 	cm := &compiledModule{
 		offsets: wazevoapi.NewModuleContextOffsetData(module, withListener), parent: e, module: module,
@@ -213,116 +218,137 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 		executables:       &executables{},
 	}
 
-	if module.IsHostModule {
-		return e.compileHostModule(ctx, module, listeners)
-	}
-
 	importedFns, localFns := int(module.ImportFunctionCount), len(module.FunctionSection)
 	if localFns == 0 {
 		return cm, nil
 	}
 
-	rels := make([]backend.RelocationInfo, 0)
-	refToBinaryOffset := make([]int, importedFns+localFns)
-
-	if wazevoapi.DeterministicCompilationVerifierEnabled {
-		// The compilation must be deterministic regardless of the order of functions being compiled.
-		wazevoapi.DeterministicCompilationVerifierRandomizeIndexes(ctx)
+	machine := newMachine()
+	relocator, err := newEngineRelocator(machine, importedFns, localFns)
+	if err != nil {
+		return nil, err
 	}
 
 	needSourceInfo := module.DWARFLines != nil
 
-	// Creates new compiler instances which are reused for each function.
 	ssaBuilder := ssa.NewBuilder()
-	fe := frontend.NewFrontendCompiler(module, ssaBuilder, &cm.offsets, ensureTermination, withListener, needSourceInfo)
-	machine := newMachine()
 	be := backend.NewCompiler(ctx, machine, ssaBuilder)
-
 	cm.executables.compileEntryPreambles(module, machine, be)
-
-	totalSize := 0 // Total binary size of the executable.
 	cm.functionOffsets = make([]int, localFns)
-	bodies := make([][]byte, localFns)
 
-	// Trampoline relocation related variables.
-	trampolineInterval, callTrampolineIslandSize, err := machine.CallTrampolineIslandInfo(localFns)
-	if err != nil {
-		return nil, err
+	var indexes []int
+	if wazevoapi.DeterministicCompilationVerifierEnabled {
+		// The compilation must be deterministic regardless of the order of functions being compiled.
+		indexes = wazevoapi.DeterministicCompilationVerifierRandomizeIndexes(ctx)
 	}
-	needCallTrampoline := callTrampolineIslandSize > 0
-	var callTrampolineIslandOffsets []int // Holds the offsets of trampoline islands.
 
-	for i := range module.CodeSection {
-		if wazevoapi.DeterministicCompilationVerifierEnabled {
-			i = wazevoapi.DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex(ctx, i)
-		}
+	if workers := experimental.GetCompilationWorkers(ctx); workers <= 1 {
+		// Compile with a single goroutine.
+		fe := frontend.NewFrontendCompiler(module, ssaBuilder, &cm.offsets, ensureTermination, withListener, needSourceInfo)
 
-		fidx := wasm.Index(i + importedFns)
-
-		if wazevoapi.NeedFunctionNameInContext {
-			def := module.FunctionDefinition(fidx)
-			name := def.DebugName()
-			if len(def.ExportNames()) > 0 {
-				name = def.ExportNames()[0]
+		for i := range module.CodeSection {
+			if wazevoapi.DeterministicCompilationVerifierEnabled {
+				i = indexes[i]
 			}
-			ctx = wazevoapi.SetCurrentFunctionName(ctx, i, fmt.Sprintf("[%d/%d]%s", i, len(module.CodeSection)-1, name))
-		}
 
-		needListener := len(listeners) > 0 && listeners[i] != nil
-		body, relsPerFunc, err := e.compileLocalWasmFunction(ctx, module, wasm.Index(i), fe, ssaBuilder, be, needListener)
-		if err != nil {
-			return nil, fmt.Errorf("compile function %d/%d: %v", i, len(module.CodeSection)-1, err)
-		}
+			fidx := wasm.Index(i + importedFns)
+			fctx := functionContext(ctx, module, i, fidx)
 
-		// Align 16-bytes boundary.
-		totalSize = (totalSize + 15) &^ 15
-		cm.functionOffsets[i] = totalSize
-
-		if needSourceInfo {
-			// At the beginning of the function, we add the offset of the function body so that
-			// we can resolve the source location of the call site of before listener call.
-			cm.sourceMap.executableOffsets = append(cm.sourceMap.executableOffsets, uintptr(totalSize))
-			cm.sourceMap.wasmBinaryOffsets = append(cm.sourceMap.wasmBinaryOffsets, module.CodeSection[i].BodyOffsetInCodeSection)
-
-			for _, info := range be.SourceOffsetInfo() {
-				cm.sourceMap.executableOffsets = append(cm.sourceMap.executableOffsets, uintptr(totalSize)+uintptr(info.ExecutableOffset))
-				cm.sourceMap.wasmBinaryOffsets = append(cm.sourceMap.wasmBinaryOffsets, uint64(info.SourceOffset))
+			needListener := len(listeners) > i && listeners[i] != nil
+			body, relsPerFunc, err := e.compileLocalWasmFunction(fctx, module, wasm.Index(i), fe, ssaBuilder, be, needListener)
+			if err != nil {
+				return nil, fmt.Errorf("compile function %d/%d: %v", i, len(module.CodeSection)-1, err)
 			}
+
+			relocator.appendFunction(fctx, module, cm, i, fidx, body, relsPerFunc, be.SourceOffsetInfo())
+		}
+	} else {
+		// Compile with N worker goroutines.
+		// Collect compiled functions across workers in a slice,
+		// to be added to the relocator in-order and resolved serially at the end.
+		// This uses more memory and CPU (across cores), but can be significantly faster.
+		type compiledFunc struct {
+			fctx        context.Context
+			fnum        int
+			fidx        wasm.Index
+			body        []byte
+			relsPerFunc []backend.RelocationInfo
+			offsPerFunc []backend.SourceOffsetInfo
 		}
 
-		fref := frontend.FunctionIndexToFuncRef(fidx)
-		refToBinaryOffset[fref] = totalSize
+		compiledFuncs := make([]compiledFunc, len(module.CodeSection))
+		ctx, cancel := context.WithCancelCause(ctx)
+		defer cancel(nil)
 
-		// At this point, relocation offsets are relative to the start of the function body,
-		// so we adjust it to the start of the executable.
-		for _, r := range relsPerFunc {
-			r.Offset += int64(totalSize)
-			rels = append(rels, r)
+		var count atomic.Uint32
+		var wg sync.WaitGroup
+		wg.Add(workers)
+
+		for range workers {
+			go func() {
+				defer wg.Done()
+
+				// Creates new compiler instances which are reused for each function.
+				machine := newMachine()
+				ssaBuilder := ssa.NewBuilder()
+				be := backend.NewCompiler(ctx, machine, ssaBuilder)
+				fe := frontend.NewFrontendCompiler(module, ssaBuilder, &cm.offsets, ensureTermination, withListener, needSourceInfo)
+
+				for {
+					if err := ctx.Err(); err != nil {
+						// Compilation canceled!
+						return
+					}
+
+					i := int(count.Add(1)) - 1
+					if i >= len(module.CodeSection) {
+						return
+					}
+
+					if wazevoapi.DeterministicCompilationVerifierEnabled {
+						i = indexes[i]
+					}
+
+					fidx := wasm.Index(i + importedFns)
+					fctx := functionContext(ctx, module, i, fidx)
+
+					needListener := len(listeners) > i && listeners[i] != nil
+					body, relsPerFunc, err := e.compileLocalWasmFunction(fctx, module, wasm.Index(i), fe, ssaBuilder, be, needListener)
+					if err != nil {
+						cancel(fmt.Errorf("compile function %d/%d: %v", i, len(module.CodeSection)-1, err))
+						return
+					}
+
+					compiledFuncs[i] = compiledFunc{
+						fctx, i, fidx, body,
+						// These slices are internal to the backend compiler and since we are going to buffer them instead
+						// of process them immediately we need to copy the memory.
+						cloneSlice(relsPerFunc),
+						cloneSlice(be.SourceOffsetInfo()),
+					}
+				}
+			}()
 		}
 
-		bodies[i] = body
-		totalSize += len(body)
-		if wazevoapi.PrintMachineCodeHexPerFunction {
-			fmt.Printf("[[[machine code for %s]]]\n%s\n\n", wazevoapi.GetCurrentFunctionName(ctx), hex.EncodeToString(body))
+		wg.Wait()
+		if err := context.Cause(ctx); err != nil {
+			return nil, err
 		}
 
-		if needCallTrampoline {
-			// If the total size exceeds the trampoline interval, we need to add a trampoline island.
-			if totalSize/trampolineInterval > len(callTrampolineIslandOffsets) {
-				callTrampolineIslandOffsets = append(callTrampolineIslandOffsets, totalSize)
-				totalSize += callTrampolineIslandSize
-			}
+		for i := range compiledFuncs {
+			fn := &compiledFuncs[i]
+			relocator.appendFunction(fn.fctx, module, cm, fn.fnum, fn.fidx, fn.body, fn.relsPerFunc, fn.offsPerFunc)
 		}
 	}
 
 	// Allocate executable memory and then copy the generated machine code.
-	executable, err := platform.MmapCodeSegment(totalSize)
+	executable, err := platform.MmapCodeSegment(relocator.totalSize)
 	if err != nil {
 		panic(err)
 	}
 	cm.executable = executable
 
-	for i, b := range bodies {
+	for i, b := range relocator.bodies {
 		offset := cm.functionOffsets[i]
 		copy(executable[offset:], b)
 	}
@@ -337,10 +363,7 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 		}
 	}
 
-	// Resolve relocations for local function calls.
-	if len(rels) > 0 {
-		machine.ResolveRelocations(refToBinaryOffset, importedFns, executable, rels, callTrampolineIslandOffsets)
-	}
+	relocator.resolveRelocations(machine, executable, importedFns)
 
 	if err = platform.MprotectRX(executable); err != nil {
 		return nil, err
@@ -350,6 +373,96 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 	return cm, nil
 }
 
+func functionContext(ctx context.Context, module *wasm.Module, fnum int, fidx wasm.Index) context.Context {
+	if wazevoapi.NeedFunctionNameInContext {
+		def := module.FunctionDefinition(fidx)
+		name := def.DebugName()
+		if len(def.ExportNames()) > 0 {
+			name = def.ExportNames()[0]
+		}
+		ctx = wazevoapi.SetCurrentFunctionName(ctx, fnum, fmt.Sprintf("[%d/%d]%s", fnum, len(module.CodeSection)-1, name))
+	}
+	return ctx
+}
+
+type engineRelocator struct {
+	bodies                      [][]byte
+	refToBinaryOffset           []int
+	rels                        []backend.RelocationInfo
+	totalSize                   int // Total binary size of the executable.
+	trampolineInterval          int
+	callTrampolineIslandSize    int
+	callTrampolineIslandOffsets []int // Holds the offsets of trampoline islands.
+}
+
+func newEngineRelocator(
+	machine backend.Machine,
+	importedFns, localFns int,
+) (r engineRelocator, err error) {
+	// Trampoline relocation related variables.
+	r.trampolineInterval, r.callTrampolineIslandSize, err = machine.CallTrampolineIslandInfo(localFns)
+	r.refToBinaryOffset = make([]int, importedFns+localFns)
+	return
+}
+
+func (r *engineRelocator) resolveRelocations(machine backend.Machine, executable []byte, importedFns int) {
+	// Resolve relocations for local function calls.
+	if len(r.rels) > 0 {
+		machine.ResolveRelocations(r.refToBinaryOffset, importedFns, executable, r.rels, r.callTrampolineIslandOffsets)
+	}
+}
+
+func (r *engineRelocator) appendFunction(
+	ctx context.Context,
+	module *wasm.Module,
+	cm *compiledModule,
+	fnum int, fidx wasm.Index,
+	body []byte,
+	relsPerFunc []backend.RelocationInfo,
+	offsPerFunc []backend.SourceOffsetInfo,
+) {
+	// Align 16-bytes boundary.
+	r.totalSize = (r.totalSize + 15) &^ 15
+	cm.functionOffsets[fnum] = r.totalSize
+
+	needSourceInfo := module.DWARFLines != nil
+	if needSourceInfo {
+		// At the beginning of the function, we add the offset of the function body so that
+		// we can resolve the source location of the call site of before listener call.
+		cm.sourceMap.executableOffsets = append(cm.sourceMap.executableOffsets, uintptr(r.totalSize))
+		cm.sourceMap.wasmBinaryOffsets = append(cm.sourceMap.wasmBinaryOffsets, module.CodeSection[fnum].BodyOffsetInCodeSection)
+
+		for _, info := range offsPerFunc {
+			cm.sourceMap.executableOffsets = append(cm.sourceMap.executableOffsets, uintptr(r.totalSize)+uintptr(info.ExecutableOffset))
+			cm.sourceMap.wasmBinaryOffsets = append(cm.sourceMap.wasmBinaryOffsets, uint64(info.SourceOffset))
+		}
+	}
+
+	fref := frontend.FunctionIndexToFuncRef(fidx)
+	r.refToBinaryOffset[fref] = r.totalSize
+
+	// At this point, relocation offsets are relative to the start of the function body,
+	// so we adjust it to the start of the executable.
+	for _, rel := range relsPerFunc {
+		rel.Offset += int64(r.totalSize)
+		r.rels = append(r.rels, rel)
+	}
+
+	r.totalSize += len(body)
+	r.bodies = append(r.bodies, body)
+	if wazevoapi.PrintMachineCodeHexPerFunction {
+		fmt.Printf("[[[machine code for %s]]]\n%s\n\n", wazevoapi.GetCurrentFunctionName(ctx), hex.EncodeToString(body))
+	}
+
+	if r.callTrampolineIslandSize > 0 {
+		// If the total size exceeds the trampoline interval, we need to add a trampoline island.
+		if r.totalSize/r.trampolineInterval > len(r.callTrampolineIslandOffsets) {
+			r.callTrampolineIslandOffsets = append(r.callTrampolineIslandOffsets, r.totalSize)
+			r.totalSize += r.callTrampolineIslandSize
+		}
+	}
+}
+
 func (e *engine) compileLocalWasmFunction(
 	ctx context.Context,
 	module *wasm.Module,
@@ -396,9 +509,7 @@ func (e *engine) compileLocalWasmFunction(
 	}
 
 	// TODO: optimize as zero copy.
-	copied := make([]byte, len(original))
-	copy(copied, original)
-	return copied, rels, nil
+	return cloneSlice(original), rels, nil
 }
 
 func (e *engine) compileHostModule(ctx context.Context, module *wasm.Module, listeners []experimental.FunctionListener) (*compiledModule, error) {
@@ -470,9 +581,7 @@ func (e *engine) compileHostModule(ctx context.Context, module *wasm.Module, lis
 		}
 
 		// TODO: optimize as zero copy.
-		copied := make([]byte, len(body))
-		copy(copied, body)
-		bodies[i] = copied
+		bodies[i] = cloneSlice(body)
 		totalSize += len(body)
 	}
 
@@ -835,3 +944,10 @@ func (cm *compiledModule) getSourceOffset(pc uintptr) uint64 {
 	}
 	return cm.sourceMap.wasmBinaryOffsets[index]
 }
+
+func cloneSlice[S ~[]E, E any](s S) S {
+	if s == nil {
+		return nil
+	}
+	return append(S{}, s...)
+}

internal/engine/wazevo/engine_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"unsafe"
 
+	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 	"github.com/tetratelabs/wazero/internal/wasm"
@@ -45,33 +46,37 @@ func (f fakeFinalizer) setFinalizer(obj interface{}, finalizer interface{}) {
 }
 
 func TestEngine_CompileModule(t *testing.T) {
-	ctx := context.Background()
-	e := NewEngine(ctx, 0, nil).(*engine)
-	ff := fakeFinalizer{}
-	e.setFinalizer = ff.setFinalizer
+	for _, concurrency := range []int{1, 4} {
+		t.Run(fmt.Sprintf("concurrency_%d", concurrency), func(t *testing.T) {
+			ctx := experimental.WithCompilationWorkers(context.Background(), concurrency)
+			e := NewEngine(ctx, 0, nil).(*engine)
+			ff := fakeFinalizer{}
+			e.setFinalizer = ff.setFinalizer
 
-	okModule := &wasm.Module{
-		TypeSection:     []wasm.FunctionType{{}},
-		FunctionSection: []wasm.Index{0, 0, 0, 0},
-		CodeSection: []wasm.Code{
-			{Body: []byte{wasm.OpcodeEnd}},
-			{Body: []byte{wasm.OpcodeEnd}},
-			{Body: []byte{wasm.OpcodeEnd}},
-			{Body: []byte{wasm.OpcodeEnd}},
-		},
-		ID: wasm.ModuleID{},
-	}
+			okModule := &wasm.Module{
+				TypeSection:     []wasm.FunctionType{{}},
+				FunctionSection: []wasm.Index{0, 0, 0, 0},
+				CodeSection: []wasm.Code{
+					{Body: []byte{wasm.OpcodeEnd}},
+					{Body: []byte{wasm.OpcodeEnd}},
+					{Body: []byte{wasm.OpcodeEnd}},
+					{Body: []byte{wasm.OpcodeEnd}},
+				},
+				ID: wasm.ModuleID{},
+			}
 
-	err := e.CompileModule(ctx, okModule, nil, false)
-	require.NoError(t, err)
+			err := e.CompileModule(ctx, okModule, nil, false)
+			require.NoError(t, err)
 
-	// Compiling same module shouldn't be compiled again, but instead should be cached.
-	err = e.CompileModule(ctx, okModule, nil, false)
-	require.NoError(t, err)
+			// Compiling same module shouldn't be compiled again, but instead should be cached.
+			err = e.CompileModule(ctx, okModule, nil, false)
+			require.NoError(t, err)
 
-	// Pretend the finalizer executed, by invoking them one-by-one.
-	for k, v := range ff {
-		v(k)
+			// Pretend the finalizer executed, by invoking them one-by-one.
+			for k, v := range ff {
+				v(k)
+			}
+		})
 	}
 }
 

internal/engine/wazevo/wazevoapi/debug_options.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math/rand"
 	"os"
+	"sync"
 	"time"
 )
 
@@ -91,7 +92,7 @@ type (
 		initialCompilationDone bool
 		maybeRandomizedIndexes []int
 		r                      *rand.Rand
-		values                 map[string]string
+		values                 sync.Map
 	}
 	verifierStateContextKey struct{}
 	currentFunctionNameKey  struct{}
@@ -106,31 +107,24 @@ func NewDeterministicCompilationVerifierContext(ctx context.Context, localFuncti
 	}
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
 	return context.WithValue(ctx, verifierStateContextKey{}, &verifierState{
-		r: r, maybeRandomizedIndexes: maybeRandomizedIndexes, values: map[string]string{},
+		r: r, maybeRandomizedIndexes: maybeRandomizedIndexes, values: sync.Map{},
 	})
 }
 
 // DeterministicCompilationVerifierRandomizeIndexes randomizes the indexes for the deterministic compilation verifier.
-// To get the randomized index, use DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex.
-func DeterministicCompilationVerifierRandomizeIndexes(ctx context.Context) {
+// Returns a slice that maps an index to the randomized index.
+func DeterministicCompilationVerifierRandomizeIndexes(ctx context.Context) []int {
 	state := ctx.Value(verifierStateContextKey{}).(*verifierState)
 	if !state.initialCompilationDone {
 		// If this is the first attempt, we use the index as-is order.
 		state.initialCompilationDone = true
-		return
+		return state.maybeRandomizedIndexes
 	}
 	r := state.r
 	r.Shuffle(len(state.maybeRandomizedIndexes), func(i, j int) {
 		state.maybeRandomizedIndexes[i], state.maybeRandomizedIndexes[j] = state.maybeRandomizedIndexes[j], state.maybeRandomizedIndexes[i]
 	})
-}
-
-// DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex returns the randomized index for the given `index`
-// which is assigned by DeterministicCompilationVerifierRandomizeIndexes.
-func DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex(ctx context.Context, index int) int {
-	state := ctx.Value(verifierStateContextKey{}).(*verifierState)
-	ret := state.maybeRandomizedIndexes[index]
-	return ret
+	return state.maybeRandomizedIndexes
 }
 
 // VerifyOrSetDeterministicCompilationContextValue verifies that the `newValue` is the same as the previous value for the given `scope`
@@ -141,9 +135,8 @@ func VerifyOrSetDeterministicCompilationContextValue(ctx context.Context, scope
 	fn := ctx.Value(currentFunctionNameKey{}).(string)
 	key := fn + ": " + scope
 	verifierCtx := ctx.Value(verifierStateContextKey{}).(*verifierState)
-	oldValue, ok := verifierCtx.values[key]
-	if !ok {
-		verifierCtx.values[key] = newValue
+	oldValue, loaded := verifierCtx.values.LoadOrStore(key, newValue)
+	if !loaded {
 		return
 	}
 	if oldValue != newValue {

internal/expctxkeys/compilationworkers.go

@@ -0,0 +1,6 @@
+package expctxkeys
+
+// CompilationWorkers is a context.Context Value key.
+// Its associated value should be an int representing the number of workers
+// we want to spawn to compile a given wasm input.
+type CompilationWorkers struct{}