Optimize WASM plugin execution with compilation caching (#363)

* optimize WASM plugin execution with compilation caching

* fix WASM plugin cache key to include sha256 and return errors from Close

Include opt.Sha256 in the cache key so that the same path with different
hash values does not bypass hash verification. Return errors from Close
methods using errors.Join instead of silently discarding them.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* address PR #363 review feedback

- Use content-derived sha256 as cache key instead of user-provided value
- Make wasmPluginCache thread-safe with sync.RWMutex
- Add defer g.Close(ctx) in CLI entry point
- Return SupportedFeatures even when result.Files is empty

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* remove counter field

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: goccy

cmd/grpc-federation-generator/main.go

@@ -43,6 +43,7 @@ func _main(ctx context.Context, args []string, opt *option) error {
 		protoPath = args[0]
 	}
 	g := generator.New(cfg)
+	defer g.Close(ctx)
 	var opts []generator.Option
 	if opt.WatchMode {
 		opts = append(opts, generator.WatchMode())

generator/generator.go

@@ -3,8 +3,6 @@ package generator
 import (
 	"bytes"
 	"context"
-	"crypto/sha256"
-	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -25,7 +23,6 @@ import (
 
 	"github.com/mercari/grpc-federation/compiler"
 	"github.com/mercari/grpc-federation/grpc/federation/generator"
-	"github.com/mercari/grpc-federation/grpc/federation/generator/plugin"
 	"github.com/mercari/grpc-federation/resolver"
 	"github.com/mercari/grpc-federation/source"
 	"github.com/mercari/grpc-federation/validator"
@@ -47,6 +44,7 @@ type Generator struct {
 	postProcessHandler        PostProcessHandler
 	buildCacheMap             BuildCacheMap
 	absPathToRelativePath     map[string]string
+	wasmPluginCache           *wasmPluginCache
 }
 
 type Option func(*Generator) error
@@ -135,9 +133,18 @@ func New(cfg Config) *Generator {
 		validator:             validator.New(),
 		importPaths:           cfg.Imports,
 		absPathToRelativePath: make(map[string]string),
+		wasmPluginCache:       newWasmPluginCache(),
 	}
 }
 
+// Close releases resources held by the Generator, including cached WASM plugin runtimes.
+func (g *Generator) Close(ctx context.Context) error {
+	if g.wasmPluginCache != nil {
+		return g.wasmPluginCache.Close(ctx)
+	}
+	return nil
+}
+
 func (g *Generator) SetPostProcessHandler(postProcessHandler func(ctx context.Context, path string, result Result) error) {
 	g.postProcessHandler = postProcessHandler
 }
@@ -182,7 +189,7 @@ func (g *Generator) Generate(ctx context.Context, protoPath string, opts ...Opti
 		}
 		results = append(results, result)
 	}
-	pluginResp, err := evalAllCodeGenerationPlugin(ctx, results, g.federationGeneratorOption)
+	pluginResp, err := evalAllCodeGenerationPlugin(ctx, results, g.federationGeneratorOption, g.wasmPluginCache)
 	if err != nil {
 		return err
 	}
@@ -359,7 +366,7 @@ func (g *Generator) setWatcher(w *Watcher) error {
 			results = append(results, result)
 		}
 		results = append(results, g.otherResults(path)...)
-		pluginResp, err := evalAllCodeGenerationPlugin(ctx, results, g.federationGeneratorOption)
+		pluginResp, err := evalAllCodeGenerationPlugin(ctx, results, g.federationGeneratorOption, g.wasmPluginCache)
 		if err != nil {
 			log.Printf("failed to run code generator plugin: %+v", err)
 		}
@@ -651,7 +658,7 @@ func (g *Generator) fileNameWithoutExt(name string) string {
 	return name[:len(name)-len(filepath.Ext(name))]
 }
 
-func CreateCodeGeneratorResponse(ctx context.Context, req *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse, error) {
+func CreateCodeGeneratorResponse(ctx context.Context, req *pluginpb.CodeGeneratorRequest) (_ *pluginpb.CodeGeneratorResponse, err error) {
 	opt, err := parseOptString(req.GetParameter())
 	if err != nil {
 		return nil, err
@@ -665,20 +672,33 @@ func CreateCodeGeneratorResponse(ctx context.Context, req *pluginpb.CodeGenerato
 		fmt.Fprint(os.Stderr, validator.Format(outs))
 	}
 
+	var resp pluginpb.CodeGeneratorResponse
+	// TODO: Since we don’t currently support editions, we will comment it out.
+	// Strictly speaking, proto3 optional is also not fully supported, but because it cannot be used together when other plugins support proto3 optional,
+	// we have enabled it for the time being.
+	resp.SupportedFeatures = proto.Uint64(uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL /*| pluginpb.CodeGeneratorResponse_FEATURE_SUPPORTS_EDITIONS*/))
+
+	if len(result.Files) == 0 {
+		return &resp, nil
+	}
+
+	cache := newWasmPluginCache()
+	defer func() {
+		if closeErr := cache.Close(ctx); closeErr != nil {
+			err = errors.Join(err, closeErr)
+		}
+	}()
 	pluginResp, err := evalAllCodeGenerationPlugin(ctx, []*ProtoFileResult{
 		{
 			Type:            ProtocAction,
 			ProtoPath:       "",
 			FederationFiles: result.Files,
 		},
-	}, opt)
+	}, opt, cache)
 	if err != nil {
 		return nil, err
 	}
-
-	var resp pluginpb.CodeGeneratorResponse
 	resp.File = append(resp.File, pluginResp.File...)
-	resp.SupportedFeatures = proto.Uint64(uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL /*| pluginpb.CodeGeneratorResponse_FEATURE_SUPPORTS_EDITIONS*/))
 	for _, file := range result.Files {
 		out, err := NewCodeGenerator().Generate(file)
 		if err != nil {
@@ -696,61 +716,38 @@ func CreateCodeGeneratorResponse(ctx context.Context, req *pluginpb.CodeGenerato
 	return &resp, nil
 }
 
-func evalAllCodeGenerationPlugin(ctx context.Context, results []*ProtoFileResult, opt *CodeGeneratorOption) (*pluginpb.CodeGeneratorResponse, error) {
+func evalAllCodeGenerationPlugin(ctx context.Context, results []*ProtoFileResult, opt *CodeGeneratorOption, cache *wasmPluginCache) (*pluginpb.CodeGeneratorResponse, error) {
 	if len(results) == 0 {
 		return &pluginpb.CodeGeneratorResponse{}, nil
 	}
 	if opt == nil || len(opt.Plugins) == 0 {
 		return &pluginpb.CodeGeneratorResponse{}, nil
 	}
+
 	var resp pluginpb.CodeGeneratorResponse
 	for _, result := range results {
-		pluginFiles := make([]*plugin.ProtoCodeGeneratorResponse_File, 0, len(result.Files))
-		for _, file := range result.Files {
-			fileBytes, err := proto.Marshal(file)
+		for _, p := range opt.Plugins {
+			wp, err := cache.getOrCreate(ctx, p)
 			if err != nil {
 				return nil, err
 			}
-			var pluginFile plugin.ProtoCodeGeneratorResponse_File
-			if err := proto.Unmarshal(fileBytes, &pluginFile); err != nil {
-				return nil, err
-			}
-			pluginFiles = append(pluginFiles, &pluginFile)
-		}
-		genReq := generator.CreateCodeGeneratorRequest(&generator.CodeGeneratorRequestConfig{
-			Type:                 generator.ActionType(result.Type),
-			ProtoPath:            result.ProtoPath,
-			Files:                pluginFiles,
-			GRPCFederationFiles:  result.FederationFiles,
-			OutputFilePathConfig: opt.Path,
-		})
-		encodedGenReq, err := proto.Marshal(genReq)
-		if err != nil {
-			return nil, err
-		}
-		for _, plugin := range opt.Plugins {
-			wasmFile, err := os.ReadFile(plugin.Path)
+			genReq := generator.CreateCodeGeneratorRequest(&generator.CodeGeneratorRequestConfig{
+				ProtoPath:            result.ProtoPath,
+				GRPCFederationFiles:  result.FederationFiles,
+				OutputFilePathConfig: opt.Path,
+			})
+			encodedGenReq, err := proto.Marshal(genReq)
 			if err != nil {
-				return nil, fmt.Errorf("grpc-federation: failed to read plugin file: %s: %w", plugin.Path, err)
-			}
-			if plugin.Sha256 != "" {
-				hash := sha256.Sum256(wasmFile)
-				gotHash := hex.EncodeToString(hash[:])
-				if plugin.Sha256 != gotHash {
-					return nil, fmt.Errorf(
-						`grpc-federation: expected plugin sha256 value is [%s] but got [%s]`,
-						plugin.Sha256,
-						gotHash,
-					)
-				}
+				return nil, err
 			}
-			pluginRes, err := evalCodeGeneratorPlugin(ctx, wasmFile, bytes.NewBuffer(encodedGenReq))
+			pluginRes, err := wp.Execute(ctx, bytes.NewBuffer(encodedGenReq))
 			if err != nil {
 				return nil, err
 			}
 			resp.File = append(resp.File, pluginRes.File...)
 		}
 	}
+
 	return &resp, nil
 }
 

generator/wasm.go

@@ -3,37 +3,63 @@ package generator
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"sync"
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/pluginpb"
 )
 
-func evalCodeGeneratorPlugin(ctx context.Context, pluginFile []byte, req io.Reader) (*pluginpb.CodeGeneratorResponse, error) {
+// wasmPlugin holds a compiled WASM module and its runtime, allowing
+// the module to be instantiated multiple times without recompilation.
+// Execute is safe for concurrent use.
+type wasmPlugin struct {
+	runtime  wazero.Runtime
+	compiled wazero.CompiledModule
+}
+
+func newWasmPlugin(ctx context.Context, wasmBytes []byte) (*wasmPlugin, error) {
 	runtimeCfg := wazero.NewRuntimeConfigInterpreter()
 	if cache := getCompilationCache(); cache != nil {
 		runtimeCfg = runtimeCfg.WithCompilationCache(cache)
 	}
-
 	r := wazero.NewRuntimeWithConfig(ctx, runtimeCfg)
-
 	wasi_snapshot_preview1.MustInstantiate(ctx, r)
-	buf := bytes.NewBuffer([]byte{})
 
+	compiled, err := r.CompileModule(ctx, wasmBytes)
+	if err != nil {
+		r.Close(ctx)
+		return nil, fmt.Errorf("grpc-federation: failed to compile code-generator plugin: %w", err)
+	}
+	return &wasmPlugin{
+		runtime:  r,
+		compiled: compiled,
+	}, nil
+}
+
+func (p *wasmPlugin) Execute(ctx context.Context, req io.Reader) (*pluginpb.CodeGeneratorResponse, error) {
+	buf := new(bytes.Buffer)
 	modCfg := wazero.NewModuleConfig().
 		WithFSConfig(wazero.NewFSConfig().WithDirMount(".", "/")).
 		WithStdin(req).
 		WithStdout(buf).
 		WithStderr(os.Stderr).
 		WithArgs("wasi")
-	if _, err := r.InstantiateWithConfig(ctx, pluginFile, modCfg); err != nil {
+
+	mod, err := p.runtime.InstantiateModule(ctx, p.compiled, modCfg)
+	if err != nil {
 		return nil, fmt.Errorf("grpc-federation: failed to instantiate code-generator plugin: %w", err)
 	}
+	mod.Close(ctx)
+
 	var res pluginpb.CodeGeneratorResponse
 	resBytes := buf.Bytes()
 	if len(resBytes) != 0 {
@@ -44,6 +70,70 @@ func evalCodeGeneratorPlugin(ctx context.Context, pluginFile []byte, req io.Read
 	return &res, nil
 }
 
+func (p *wasmPlugin) Close(ctx context.Context) error {
+	return p.runtime.Close(ctx)
+}
+
+// wasmPluginCache caches compiled WASM plugins so that the expensive
+// compilation step is performed only once per plugin path.
+type wasmPluginCache struct {
+	mu      sync.RWMutex
+	plugins map[string]*wasmPlugin
+}
+
+func newWasmPluginCache() *wasmPluginCache {
+	return &wasmPluginCache{plugins: make(map[string]*wasmPlugin)}
+}
+
+func (c *wasmPluginCache) getOrCreate(ctx context.Context, opt *WasmPluginOption) (*wasmPlugin, error) {
+	wasmFile, err := os.ReadFile(opt.Path)
+	if err != nil {
+		return nil, fmt.Errorf("grpc-federation: failed to read plugin file: %s: %w", opt.Path, err)
+	}
+	hash := sha256.Sum256(wasmFile)
+	gotHash := hex.EncodeToString(hash[:])
+	if opt.Sha256 != "" && opt.Sha256 != gotHash {
+		return nil, fmt.Errorf(
+			`grpc-federation: expected plugin sha256 value is [%s] but got [%s]`,
+			opt.Sha256,
+			gotHash,
+		)
+	}
+	cacheKey := opt.Path + ":" + gotHash
+
+	c.mu.RLock()
+	if wp, ok := c.plugins[cacheKey]; ok {
+		c.mu.RUnlock()
+		return wp, nil
+	}
+	c.mu.RUnlock()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// Double-check after acquiring write lock.
+	if wp, ok := c.plugins[cacheKey]; ok {
+		return wp, nil
+	}
+	wp, err := newWasmPlugin(ctx, wasmFile)
+	if err != nil {
+		return nil, err
+	}
+	c.plugins[cacheKey] = wp
+	return wp, nil
+}
+
+func (c *wasmPluginCache) Close(ctx context.Context) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	errs := make([]error, 0, len(c.plugins))
+	for _, wp := range c.plugins {
+		errs = append(errs, wp.Close(ctx))
+	}
+	return errors.Join(errs...)
+}
+
 func getCompilationCache() wazero.CompilationCache {
 	tmpDir := os.TempDir()
 	if tmpDir == "" {