Sharing a compilation cache

A four-line Wasm module calls a host function that hands back a counter. Two runtimes run it, each with its own counter, and both take the machine code from one cache. Compiled code, it turns out, does not belong to the runtime that compiled it. It belongs to the cache, and it can outlive every runtime, every goroutine, and the process itself.

Run it

cd examples/multiple-runtimes
go run counter.go

m1 counter=0
m2 counter=0
m1 counter=1
m2 counter=1

Two instances of the same module, interleaved, each counting from zero on its own. One compilation produced both.

Four lines of wat, and no state

(module $counter
  ;; import `next_i32` function from the `env` module.
  (import "env" "next_i32" (func (result i32)))
  ;; get returns the next counter value by calling the imported `next_i32` function.
  (func (export "get") (result i32) (call 0))
)

The module holds no state. It exports get, which forwards to an import the host supplies. All the state lives on the Go side, which is exactly what forces the shape of the rest of this example.

counter.wasm is committed next to the source. The build command is recorded in the Go file that embeds it:

// counterWasm was generated by the following:
//
//  cd testdata; wat2wasm --debug-names counter.wat
//
//go:embed testdata/counter.wasm
var counterWasm []byte

Two runtimes, one compile

A host module name is registered once per runtime. If two module instances each need their own env — their own counter — then they need two runtimes. That would normally mean compiling counter.wasm twice, since each runtime carries its own engine and its own compiled-code map. A CompilationCache passed to both configs collapses that back to one.

// main shows how to share the same compilation cache across the multiple runtimes.
func main() {
  // Choose the context to use for function calls.
  ctx := context.Background()

6 collapsed lines
  // Prepare a cache directory.
  cacheDir, err := os.MkdirTemp("", "example")
  if err != nil {
    log.Panicln(err)
  }
  defer os.RemoveAll(cacheDir)

  // Initializes the new compilation cache with the cache directory.
  // This allows the compilation caches to be shared even across multiple OS processes.
  cache, err := wazy.NewCompilationCacheWithDir(cacheDir)
  if err != nil {
    log.Panicln(err)
  }
  defer cache.Close(ctx)

  // Creates a shared runtime config to share the cache across multiple wazy.Runtime.
  runtimeConfig := wazy.NewRuntimeConfig().WithCompilationCache(cache)

  // Creates two wazy.Runtimes with the same compilation cache.
  runtimeFoo := wazy.NewRuntimeWithConfig(ctx, runtimeConfig)
  runtimeBar := wazy.NewRuntimeWithConfig(ctx, runtimeConfig)

  // Instantiate two modules on separate Runtimes with identical configuration, which allows each instance
  // has the isolated states of "env" module.
  m1 := instantiateWithEnv(ctx, runtimeFoo)
  m2 := instantiateWithEnv(ctx, runtimeBar)

  for i := 0; i < 2; i++ {
    fmt.Printf("m1 counter=%d\n", counterGet(ctx, m1))
    fmt.Printf("m2 counter=%d\n", counterGet(ctx, m2))
  }
}

What is shared here is the cache, not the config value — WithCompilationCache clones the config, so passing the same variable to both calls is a convenience, not the mechanism. The cache owns the engine: NewRuntimeWithConfig asks it for an engine of the kind it needs and gets the same one twice, so the second Instantiate finds the compiled module already sitting in it.

The temp directory is folded above because it is scaffolding for a runnable example. In a real program the directory is a fixed path you own — /var/cache/yourapp — and the point of it is the next process, not this one.

State that has to be per-instance

// instantiateWithEnv returns a module instance.
func instantiateWithEnv(ctx context.Context, r wazy.Runtime) api.Module {
  // Instantiate a new "env" module which exports a stateful function.
  c := &counter{}
  envBuilder := r.NewHostModuleBuilder("env")
  envBuilder.NewFunctionBuilder().
    WithGoFunction(api.GoFunc(func(_ context.Context, stack []uint64) {
      stack[0] = uint64(c.getAndIncrement())
    }), nil, []api.ValueType{api.ValueTypeI32}).
    Export("next_i32")
  _, err := envBuilder.Instantiate(ctx)
  if err != nil {
    log.Panicln(err)
  }

  // Instantiate the module that imports "env".
  mod, err := r.Instantiate(ctx, counterWasm)
  if err != nil {
    log.Panicln(err)
  }
  return mod
}

Each call builds a fresh counter and closes over it. api.GoFunc takes the stack directly — result in stack[0], no reflection — so the closure is the only place the state lives.

Host modules are deliberately not cached. wazy keys a host module by the address of its wasm.Module value, so the two env modules here are two distinct cache identities that never collide, and neither is written to disk:

  // Uses the address of *wasm.Module as the module ID so that host functions can have each state per compilation.
  // Downside of this is that compilation cache on host functions (trampoline codes for Go functions and
  // Wasm codes for Wasm-implemented host functions) are not available and compiles each time. On the other hand,
  // compilation of host modules is not costly as it's merely small trampolines vs the real-world native Wasm binary.

So the expensive artifact — the guest — is compiled once and shared; the cheap trampolines are rebuilt per runtime, which is what keeps the state isolated.

What the key is made of

The disk entry is named by the hex of a sha256 over three things: the module identity, a six-byte tag naming the engine that produced it, and the raw CPU feature bits of the machine that compiled it.

func fileCacheKey(m *wasm.Module) (ret filecache.Key) {
  s := sha256.New()
  s.Write(m.ID[:])
  s.Write(magic)
  // Write the CPU features so that we can cache the compiled module for the same CPU.
  // This prevents the incompatible CPU features from being used.
  cpu := platform.CpuFeatures.Raw()
  // Reuse the `ret` buffer to write the first 8 bytes of the CPU features so that we can avoid the allocation.
  binary.LittleEndian.PutUint64(ret[:8], cpu)
  s.Write(ret[:8])
  // Finally, write the hash to the ret buffer.
  s.Sum(ret[:0])
  return
}

m.ID is itself a hash, over the binary plus every flag that changes what the compiler emits:

func (m *Module) AssignModuleID(wasm []byte, listeners []api.FunctionListener, withEnsureTermination bool, interruptCheckInterval uint64) {
  h := sha256.New()
  h.Write(wasm)
  // Use the pre-allocated space backed by m.ID below.

  // Write the existence of listeners to the checksum per function.
  for i, l := range listeners {
    binary.LittleEndian.PutUint32(m.ID[:], uint32(i))
    m.ID[4] = boolToByte(l != nil)
    h.Write(m.ID[:5])
  }
  // Write the flag of ensureTermination to the checksum.
  m.ID[0] = boolToByte(withEnsureTermination)
  h.Write(m.ID[:1])
  // Write the interrupt-check interval to the checksum.
  binary.LittleEndian.PutUint64(m.ID[:], interruptCheckInterval)
  h.Write(m.ID[:8])
  // Get checksum by passing the slice underlying m.ID.
  h.Sum(m.ID[:0])
}

Flip WithCloseOnContextDone, attach a function listener to one function, change the interrupt-check interval, edit a single byte of the guest, move to a CPU without AVX: every one of those lands on a different key. The interpreter never enters into it — it is a separate engine that caches only in memory and writes no file at all.

One input is missing from the key: CoreFeatures. The enabled feature set reaches code generation — CoreFeatureGC changes what the compiler emits — but it is hashed into neither m.ID nor the file key, so two processes pointing at one directory with different feature sets can read each other’s entries. Keep the feature set fixed for a given cache directory. The same constraint reappears inside a single process, below.

The wazy version is guarded twice. The directory itself is version-stamped:

  // Create a version-specific directory to avoid conflicts.
  dirname := path.Join(dir, "wazy-"+wazyVersion+"-"+goruntime.GOARCH+"-"+goruntime.GOOS)

and each entry repeats the version in its header, which the reader compares before trusting a byte of the payload — a mismatch is reported as staleCache and forces a recompile. A crc32 checksum over the machine code catches corruption on top of that. Upgrading wazy therefore cannot serve you last version’s code; it starts a fresh directory beside the old one, and old directories are yours to delete.

Memory or disk

Two constructors, one interface:

wazy.NewCompilationCache() — in-memory only. Scope is the process. Worth it when several runtimes in one program compile the same modules and you want them to compile once, which is precisely this example’s problem minus the persistence.
wazy.NewCompilationCacheWithDir(dirname) — the same sharing, plus every entry written through to the file system. Scope is every process that points at that directory, now or later.

The disk form costs a serialize and a write on the miss path, and pays off on the second process. The doc on NewCompilationCacheWithDir names the case exactly: “Those running wazy as a CLI or frequently restarting a process using the same wasm should use this feature to reduce time waiting to compile the same module a second time.” A long-lived server that compiles at boot and serves for weeks gets close to nothing from it; a CLI, a serverless worker, or a test binary that starts a hundred times gets the whole compile back.

You can fill the directory before anything needs it. The wazy CLI’s compile subcommand is exactly a compile with no run:

wazy compile <options> <path to wasm file>

Give it -cachedir — “Writeable directory for native code compiled from wasm. Contents are re-used for the same version of wazy.” — and the first real invocation starts warm. Building a container image is the natural place: compile during the build, ship the cache directory in the layer.

Lifetime, engines, and the version trap

A cache changes who owns the engine. Without one, Runtime.Close frees the machine code with the runtime. With one, it does not:

  err := r.store.CloseWithExitCode(ctx, exitCode)
  if r.cache == nil {
    // Close the engine if the cache is not configured, which means that this engine is scoped in this runtime.
    if errCloseEngine := r.store.Engine.Close(); errCloseEngine != nil {
      return errCloseEngine
    }
  }

The engine now lives until cache.Close(ctx). That is the point — it is how the second runtime gets the first one’s work — but it means closing runtimes in a loop reclaims nothing, and a cache you never close is compiled code you never release.

One engine per kind, configured by whoever got there first. A cache holds one engine for the compiler and one for the interpreter, each built on a sync.Once. The features passed to that constructor come from the first runtime configured with the cache. Build runtimes with different CoreFeatures off one cache and the later ones do not reconfigure the engine they inherit — give them their own cache, or keep the config identical, as this example does.

Concurrent compiles are not deduplicated. The lock is held for the lookup and the store, not across the compilation itself, so N goroutines that hit the same cold module all compile it. The CompilationCache doc is blunt about the fix: “we always recommend to produce CompiledModule share it across multiple Goroutines to avoid trying to compile the same module simultaneously.” Call CompileModule once, then InstantiateModule from many. A CompiledModule is even portable across runtimes sharing a cache — wazy re-resolves the per-store type IDs on the way in.

Your tests are probably all caching as dev. The version comes from debug.BuildInfo, which does not report it for a package that is not main (golang/go#33976). Every entry a test writes is stamped dev, so the version guard stops distinguishing anything. The workaround is in the WithCompilationCache doc: pass -ldflags "-X github.com/samyfodil/wazy/internal/version.version=foo" when running tests.

The two-cache picture — this one for core modules, component.CompileCache for component graphs — is in the caching guide, along with a table of which to reach for. wazy compile and its flags are in the CLI reference.

Compilation cachingCore-module and component caches, and when each pays.

CLI referencewazy compile, wazy run, and the cachedir flag.

The next example, Running a binary from the CLI, is the other half of the pre-warming story: it is wazy run on a command module, and --cachedir applies there the same way. The rest of the runnable programs are in the examples index; the source for this one is examples/multiple-runtimes.