cache: put WASM compilation behind the action cache

Compiling a module to machine code is cacheable work like any other, so
model it as a CompileModule action keyed by the module checksum, wazero
version, GOOS, and GOARCH instead of handing wazero a private directory
outside the cache's discipline.

The action cache gains tree-shaped outputs for tools that read and
write output directories: PutTree stores every file under a directory
as a named output blob in the CAS, and GetTree materializes them back.
Compiled machine code is materialized into exec/<action-hash>, which
wazero's compilation cache reads on start; like the rest of the cache,
exec trees are safe to delete because the authoritative bytes live in
the CAS.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MPgq3rwip76D554BktbqR4

Author: claude

docs/reference/environment-variables.md

@@ -26,15 +26,16 @@ Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-
 The cache is designed after Bazel's local disk cache and has three parts:
 
 - `cas/` — a content-addressable store holding blobs (query analysis
-  results, WASM plugin binaries) keyed by the SHA-256 hash of their
-  contents. A remotely fetched plugin's address is exactly the checksum
-  declared in the configuration file, so it is loaded directly by that
-  address.
+  results, WASM plugin binaries, compiled WASM machine code) keyed by the
+  SHA-256 hash of their contents. A remotely fetched plugin's address is
+  exactly the checksum declared in the configuration file, so it is loaded
+  directly by that address.
 - `ac/` — an action cache mapping the digest of a unit of cacheable work and
-  its inputs (for example, analyzing a query against a schema) to the CAS
-  digests of its outputs.
-- `wazero/` — compiled WASM machine code, managed by the
-  [wazero](https://wazero.io) runtime in its own format.
+  its inputs (analyzing a query against a schema, compiling a WASM module to
+  machine code) to the CAS digests of its outputs.
+- `exec/` — per-action directories where cached output trees are
+  materialized for tools that read them from disk, such as the
+  [wazero](https://wazero.io) runtime's compilation cache.
 
 The entire directory is safe to delete at any time; sqlc will rebuild it as
 needed.

internal/cache/action.go

@@ -3,6 +3,7 @@ package cache
 import (
 	"encoding/json"
 	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
 )
@@ -63,6 +64,81 @@ func (a *ActionCache) Get(action Digest) (*ActionResult, error) {
 	return &result, nil
 }
 
+// PutTree stores every file under dir in the CAS and records them as the
+// action's outputs, named by their paths relative to dir. Use this for
+// actions whose tool writes an output directory, like WASM compilation.
+func (a *ActionCache) PutTree(action Digest, dir string) error {
+	outputs := map[string]Digest{}
+	err := filepath.WalkDir(dir, func(path string, entry fs.DirEntry, err error) error {
+		if err != nil || entry.IsDir() {
+			return err
+		}
+		rel, err := filepath.Rel(dir, path)
+		if err != nil {
+			return err
+		}
+		data, err := os.ReadFile(path)
+		if err != nil {
+			return err
+		}
+		d, err := a.cas.Put(data)
+		if err != nil {
+			return err
+		}
+		outputs[filepath.ToSlash(rel)] = d
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("cache: %w", err)
+	}
+	if len(outputs) == 0 {
+		return fmt.Errorf("cache: no outputs found under %s", dir)
+	}
+	return a.Put(action, &ActionResult{Outputs: outputs})
+}
+
+// GetTree materializes a cached action's outputs as files under dir, or
+// returns ErrNotFound on a miss. Files already present with the right size
+// are left in place; missing ones are staged and renamed so concurrent
+// processes never observe partial files.
+func (a *ActionCache) GetTree(action Digest, dir string) error {
+	result, err := a.Get(action)
+	if err != nil {
+		return err
+	}
+	for rel, d := range result.Outputs {
+		path := filepath.Join(dir, filepath.FromSlash(rel))
+		if fi, err := os.Stat(path); err == nil && fi.Size() == d.SizeBytes {
+			continue
+		}
+		data, err := a.cas.Get(d)
+		if err != nil {
+			return err
+		}
+		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
+			return fmt.Errorf("cache: %w", err)
+		}
+		f, err := os.CreateTemp(a.cas.tmp, d.Hash[:8]+"-*")
+		if err != nil {
+			return fmt.Errorf("cache: %w", err)
+		}
+		if _, err := f.Write(data); err != nil {
+			f.Close()
+			os.Remove(f.Name())
+			return fmt.Errorf("cache: %w", err)
+		}
+		if err := f.Close(); err != nil {
+			os.Remove(f.Name())
+			return fmt.Errorf("cache: %w", err)
+		}
+		if err := os.Rename(f.Name(), path); err != nil {
+			os.Remove(f.Name())
+			return fmt.Errorf("cache: %w", err)
+		}
+	}
+	return nil
+}
+
 // Put records the result of an action. All outputs must already be in the
 // CAS; writes are staged and renamed so concurrent processes never observe a
 // partial entry.

internal/cache/cache.go

@@ -26,6 +26,7 @@ import (
 
 // Cache bundles the CAS and the action cache that shares it.
 type Cache struct {
+	root    string
 	CAS     *CAS
 	Actions *ActionCache
 }
@@ -61,20 +62,19 @@ func OpenAt(root string) (*Cache, error) {
 		return nil, err
 	}
 	return &Cache{
+		root:    root,
 		CAS:     cas,
 		Actions: newActionCache(root, cas),
 	}, nil
 }
 
-// WazeroDir returns the directory for wazero's compilation cache. Compiled
-// module machine code is managed by wazero in its own format, so it lives
-// beside the CAS rather than inside it.
-func WazeroDir() (string, error) {
-	root, err := Dir()
-	if err != nil {
-		return "", err
-	}
-	dir := filepath.Join(root, "wazero")
+// ExecDir returns a stable directory for materializing the output tree of
+// the given action, for tools that need their outputs on disk (like wazero's
+// compilation cache). It lives at exec/<action-hash> under the cache root
+// and, like everything else in the cache, is safe to delete at any time: the
+// authoritative copy of its contents is the CAS.
+func (c *Cache) ExecDir(action Digest) (string, error) {
+	dir := filepath.Join(c.root, "exec", action.Hash)
 	if err := os.MkdirAll(dir, 0755); err != nil {
 		return "", fmt.Errorf("failed to create %s directory: %w", dir, err)
 	}

internal/cache/cache_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"os"
+	"path/filepath"
 	"testing"
 )
 
@@ -197,6 +198,63 @@ func TestActionCacheMissingOutputIsMiss(t *testing.T) {
 	}
 }
 
+func TestActionCacheTreeRoundTrip(t *testing.T) {
+	c := testCache(t)
+
+	// Simulate a tool writing an output directory, like wazero's
+	// compilation cache.
+	src := t.TempDir()
+	files := map[string]string{
+		"wazero-v1-amd64-linux/compiled": "machine code",
+		"manifest":                       "meta",
+	}
+	for rel, contents := range files {
+		path := filepath.Join(src, filepath.FromSlash(rel))
+		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.WriteFile(path, []byte(contents), 0644); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	action := NewAction("CompileModule").AddInput("wasm", []byte("checksum")).Digest()
+	if err := c.Actions.GetTree(action, t.TempDir()); !errors.Is(err, ErrNotFound) {
+		t.Errorf("want ErrNotFound before PutTree, got %v", err)
+	}
+	if err := c.Actions.PutTree(action, src); err != nil {
+		t.Fatal(err)
+	}
+
+	// Materialize into a fresh directory and compare contents.
+	dst := t.TempDir()
+	if err := c.Actions.GetTree(action, dst); err != nil {
+		t.Fatal(err)
+	}
+	for rel, contents := range files {
+		got, err := os.ReadFile(filepath.Join(dst, filepath.FromSlash(rel)))
+		if err != nil {
+			t.Fatal(err)
+		}
+		if string(got) != contents {
+			t.Errorf("%s: got %q, want %q", rel, got, contents)
+		}
+	}
+
+	// Materializing again over the same directory is a no-op.
+	if err := c.Actions.GetTree(action, dst); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestActionCachePutTreeRejectsEmptyDir(t *testing.T) {
+	c := testCache(t)
+	action := NewAction("CompileModule").Digest()
+	if err := c.Actions.PutTree(action, t.TempDir()); err == nil {
+		t.Error("PutTree must reject a directory with no files")
+	}
+}
+
 func TestActionCachePutRejectsMissingOutput(t *testing.T) {
 	c := testCache(t)
 	action := NewAction("Test").Digest()

internal/ext/wasm/wasm.go

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"runtime"
+	"runtime/debug"
 	"strings"
 
 	"github.com/tetratelabs/wazero"
@@ -134,11 +135,30 @@ func (r *Runner) loadAndCompileWASM(ctx context.Context, store *cache.Cache, exp
 		return nil, err
 	}
 
-	wazeroDir, err := cache.WazeroDir()
+	// Compiling the module to machine code is itself a cacheable action,
+	// keyed by the module's checksum and everything else that determines the
+	// generated code. Compiled artifacts are materialized into an exec
+	// directory for wazero's compilation cache to find; the authoritative
+	// copies live in the CAS.
+	compileAction := cache.NewAction("CompileModule").
+		AddInput("wasm", []byte(expected)).
+		AddInput("wazero", []byte(wazeroVersion())).
+		AddInput("goos", []byte(runtime.GOOS)).
+		AddInput("goarch", []byte(runtime.GOARCH)).
+		Digest()
+
+	execDir, err := store.ExecDir(compileAction)
 	if err != nil {
 		return nil, err
 	}
-	wazeroCache, err := wazero.NewCompilationCacheWithDir(wazeroDir)
+	compiled := true
+	if err := store.Actions.GetTree(compileAction, execDir); errors.Is(err, cache.ErrNotFound) {
+		compiled = false
+	} else if err != nil {
+		return nil, err
+	}
+
+	wazeroCache, err := wazero.NewCompilationCacheWithDir(execDir)
 	if err != nil {
 		return nil, fmt.Errorf("wazero.NewCompilationCacheWithDir: %w", err)
 	}
@@ -151,15 +171,36 @@ func (r *Runner) loadAndCompileWASM(ctx context.Context, store *cache.Cache, exp
 	}
 
 	// Compile the Wasm binary once so that we can skip the entire compilation
-	// time during instantiation.
+	// time during instantiation. On an action cache hit this loads the
+	// materialized machine code instead of compiling.
 	code, err := rt.CompileModule(ctx, wmod)
 	if err != nil {
 		return nil, fmt.Errorf("compile module: %w", err)
 	}
 
+	if !compiled {
+		if err := store.Actions.PutTree(compileAction, execDir); err != nil {
+			slog.Warn("caching compiled module failed", "err", err)
+		}
+	}
+
 	return &runtimeAndCode{rt: rt, code: code}, nil
 }
 
+// wazeroVersion returns the version of the wazero dependency, an input to
+// the CompileModule action: its generated machine code changes between
+// wazero releases.
+func wazeroVersion() string {
+	if bi, ok := debug.ReadBuildInfo(); ok {
+		for _, dep := range bi.Deps {
+			if dep.Path == "github.com/tetratelabs/wazero" {
+				return dep.Version
+			}
+		}
+	}
+	return info.Version
+}
+
 // removePGCatalog removes the pg_catalog schema from the request. There is a
 // mysterious (reason unknown) bug with wasm plugins when a large amount of
 // tables (like there are in the catalog) are sent.