add pprof options to root command + always close compilation cache

Signed-off-by: Achille Roussel <[email protected]>

config.go

@@ -253,14 +253,15 @@ func readConfig(r io.Reader) (*configuration, error) {
 func (c *configuration) newRuntime(ctx context.Context) wazero.Runtime {
  config := wazero.NewRuntimeConfig()
 
+ var cache wazero.CompilationCache
  if cachePath, ok := c.Cache.Location.Value(); ok {
  // The cache is an optimization, so if we encounter errors we notify the
  // user but still go ahead with the runtime instantiation.
  path, err := cachePath.Resolve()
  if err != nil {
  fmt.Fprintf(os.Stderr, "ERR: resolving timecraft cache location: %s\n", err)
  } else {
- cache, err := createCacheDirectory(path)
+ cache, err = createCacheDirectory(path)
  if err != nil {
  fmt.Fprintf(os.Stderr, "ERR: creating timecraft cache directory: %s\n", err)
  } else {
@@ -269,7 +270,26 @@ func (c *configuration) newRuntime(ctx context.Context) wazero.Runtime {
  }
  }
 
- return wazero.NewRuntimeWithConfig(ctx, config)
+ runtime := wazero.NewRuntimeWithConfig(ctx, config)
+ if cache != nil {
+ runtime = &runtimeWithCompilationCache{
+ Runtime: runtime,
+ cache: cache,
+ }
+ }
+ return runtime
+}
+
+type runtimeWithCompilationCache struct {
+ wazero.Runtime
+ cache wazero.CompilationCache
+}
+
+func (r *runtimeWithCompilationCache) Close(ctx context.Context) error {
+ if r.cache != nil {
+ defer r.cache.Close(ctx)
+ }
+ return r.Runtime.Close(ctx)
 }
 
 func (c *configuration) createRegistry() (*timemachine.Registry, error) {

root.go

@@ -22,8 +22,11 @@ import (
  "fmt"
  _ "net/http/pprof"
  "os"
+ "runtime"
+ "runtime/pprof"
  "strings"
 
+ "github.com/stealthrocket/timecraft/internal/print/human"
  "golang.org/x/exp/maps"
  "golang.org/x/exp/slices"
 )
@@ -47,14 +50,49 @@ For a list of commands available, run 'timecraft help'.`
 
 // root is the timecraft entrypoint.
 func root(ctx context.Context, args ...string) int {
+ var (
+ // Secret options, we don't document them since they are only used for
+ // development. Since they are not part of the public interface we may
+ // remove or change the syntax at any time.
+ cpuProfile human.Path
+ memProfile human.Path
+ )
+
  flagSet := newFlagSet("timecraft", helpUsage)
+ customVar(flagSet, &cpuProfile, "cpuprofile")
+ customVar(flagSet, &memProfile, "memprofile")
  _ = flagSet.Parse(args)
 
  if args = flagSet.Args(); len(args) == 0 {
  fmt.Println(rootUsage)
  return 0
  }
 
+ if cpuProfile != "" {
+ path, _ := cpuProfile.Resolve()
+ f, err := os.Create(path)
+ if err != nil {
+ fmt.Fprintf(os.Stderr, "WARN: could not create CPU profile: %s\n", err)
+ } else {
+ defer f.Close()
+ _ = pprof.StartCPUProfile(f)
+ defer pprof.StopCPUProfile()
+ }
+ }
+
+ if memProfile != "" {
+ path, _ := memProfile.Resolve()
+ defer func() {
+ f, err := os.Create(path)
+ if err != nil {
+ fmt.Fprintf(os.Stderr, "WARN: could not create memory profile: %s\n", err)
+ }
+ defer f.Close()
+ runtime.GC()
+ _ = pprof.WriteHeapProfile(f)
+ }()
+ }
+
  var err error
  cmd, args := args[0], args[1:]
  switch cmd {