Add support for tracing function calls

describe.go

@@ -184,7 +184,7 @@ func describeModule(ctx context.Context, reg *timemachine.Registry, id string, c
 		return nil, err
 	}
 
-	runtime, err := timecraft.NewRuntime(ctx, config)
+	runtime, ctx, err := timecraft.NewRuntime(ctx, config)
 	if err != nil {
 		return nil, err
 	}

internal/timecraft/process.go

@@ -21,12 +21,14 @@ import (
 	"github.com/stealthrocket/timecraft/internal/object"
 	"github.com/stealthrocket/timecraft/internal/sandbox"
 	"github.com/stealthrocket/timecraft/internal/timemachine"
+	"github.com/stealthrocket/timecraft/internal/timemachine/funccall"
 	"github.com/stealthrocket/timecraft/internal/timemachine/wasicall"
 	"github.com/stealthrocket/wasi-go"
 	"github.com/stealthrocket/wasi-go/imports"
 	"github.com/stealthrocket/wasi-go/imports/wasi_snapshot_preview1"
 	"github.com/stealthrocket/wazergo"
 	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/experimental"
 )
 
 // ProcessManager runs WebAssembly modules.
@@ -316,6 +318,17 @@ func (pm *ProcessManager) Start(moduleSpec ModuleSpec, logSpec *LogSpec, parentI
 				panic(err) // caught/handled by wazero
 			}
 		})
+
+		funcListenerFactory := pm.ctx.Value(experimental.FunctionListenerFactoryKey{}).(*funccall.Factory)
+		funcListenerFactory.Enable(func(id int, ts time.Time, data []byte) {
+			b.Reset(logSpec.StartTime)
+			b.SetTimestamp(ts)
+			b.SetFunctionID(id)
+			b.SetFunctionCall(data)
+			if err := recordWriter.WriteRecord(&b); err != nil {
+				panic(err) // caught/handled by wazero
+			}
+		})
 	} else {
 		processID = uuid.New()
 	}

internal/timecraft/runtime.go

@@ -4,12 +4,14 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/stealthrocket/timecraft/internal/timemachine/funccall"
 	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/experimental"
 )
 
 // NewRuntime constructs a wazero.Runtime that's configured according
 // to the provided timecraft Config.
-func NewRuntime(ctx context.Context, config *Config) (wazero.Runtime, error) {
+func NewRuntime(ctx context.Context, config *Config) (wazero.Runtime, context.Context, error) {
 	runtimeConfig := wazero.NewRuntimeConfig()
 
 	var cache wazero.CompilationCache
@@ -18,25 +20,32 @@ func NewRuntime(ctx context.Context, config *Config) (wazero.Runtime, error) {
 		// user but still go ahead with the runtime instantiation.
 		path, err := cachePath.Resolve()
 		if err != nil {
-			return nil, fmt.Errorf("failed to resolve timecraft cache location: %w", err)
+			return nil, nil, fmt.Errorf("failed to resolve timecraft cache location: %w", err)
 		} else {
 			cache, err = createCacheDirectory(path)
 			if err != nil {
-				return nil, fmt.Errorf("failed to create timecraft cache directory: %w", err)
+				return nil, nil, fmt.Errorf("failed to create timecraft cache directory: %w", err)
 			} else {
 				runtimeConfig = runtimeConfig.WithCompilationCache(cache)
 			}
 		}
 	}
 
+	// Initialize wazero with a disabled FunctionListenerFactory implementation. We
+	// will later decide to enable this when processing module.
+	//
+	// We have to return this context and pass it down stream so we can have access
+	// to *funccall.Factory initialized with this runtime.
+	ctx = context.WithValue(ctx, experimental.FunctionListenerFactoryKey{}, &funccall.Factory{})
+
 	runtime := wazero.NewRuntimeWithConfig(ctx, runtimeConfig)
 	if cache != nil {
 		runtime = &runtimeWithCompilationCache{
 			Runtime: runtime,
 			cache:   cache,
 		}
 	}
-	return runtime, nil
+	return runtime, ctx, nil
 }
 
 type runtimeWithCompilationCache struct {

internal/timemachine/funccall/format.go

@@ -0,0 +1,54 @@
+package funccall
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+	"text/tabwriter"
+
+	"github.com/stealthrocket/timecraft/internal/stream"
+	"github.com/stealthrocket/timecraft/internal/timemachine"
+)
+
+func Format(o io.Writer, reader stream.Reader[timemachine.Record]) error {
+	tr := TraceReader{Records: reader}
+	traces := make([]Trace, 1<<10)
+	w := tabwriter.NewWriter(o, 0, 0, 1, ' ', tabwriter.AlignRight)
+	fmt.Fprintln(w, "duration\t symbol ")
+	fmt.Fprintln(w, "---\t --- ")
+	var s strings.Builder
+	write := func(n int) {
+		chunk := traces[:n]
+		for x := range chunk {
+			t := &chunk[x]
+			fmt.Fprintf(w, "%v\t %s\n", t.Duration, formatSymbol(&s, t))
+		}
+		w.Flush()
+	}
+
+	for {
+		n, err := tr.Read(traces)
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				write(n)
+				return nil
+			}
+			return err
+		}
+		write(n)
+	}
+}
+
+func formatSymbol(s *strings.Builder, t *Trace) string {
+	if t.Depth == 0 {
+		return t.Name
+	}
+	s.Reset()
+	s.Grow(t.Depth)
+	for i := 0; i < t.Depth; i++ {
+		s.WriteByte(' ')
+	}
+	s.WriteString(t.Name)
+	return s.String()
+}

internal/timemachine/funccall/log.go

@@ -0,0 +1,86 @@
+package funccall
+
+import (
+	"context"
+	"encoding/binary"
+	"time"
+
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
+)
+
+// This number is arbitrary. For now we use the same segment as system calls ,
+// any number larger than the largest system call will suffice.
+const ID = 100
+
+type Factory struct {
+	enabled bool
+	start   stack[time.Time]
+	write   func(id int, ts time.Time, data []byte)
+}
+
+func (f *Factory) Enable(w func(id int, ts time.Time, data []byte)) {
+	f.enabled = true
+	f.write = w
+}
+
+var _ experimental.FunctionListenerFactory = (*Factory)(nil)
+
+func (f *Factory) NewFunctionListener(fnd api.FunctionDefinition) experimental.FunctionListener {
+	return &listener{
+		f: f,
+	}
+}
+
+type listener struct {
+	f      *Factory
+	buffer []byte
+}
+
+var _ experimental.FunctionListener = (*listener)(nil)
+
+func (l *listener) Before(ctx context.Context,
+	mod api.Module, def api.FunctionDefinition, params []uint64, _ experimental.StackIterator) {
+	if !l.f.enabled {
+		return
+	}
+	l.f.start.push(time.Now())
+}
+
+func (l *listener) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
+	if !l.f.enabled {
+		return
+	}
+	start := l.f.start.pop().(time.Time)
+	l.buffer = l.buffer[:0]
+
+	// 4 bytes - depth ( for indentation)
+	// 8 bytes - duration
+	// followed by the function name
+	l.buffer = binary.BigEndian.AppendUint32(l.buffer, uint32(l.f.start.size()))
+	l.buffer = binary.BigEndian.AppendUint64(l.buffer, uint64(time.Since(start)))
+	l.buffer = append(l.buffer, []byte(def.DebugName())...)
+	l.f.write(ID, start, l.buffer)
+}
+
+func (l *listener) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
+	if !l.f.enabled {
+		return
+	}
+	l.f.start.pop()
+}
+
+type stack[T any] struct {
+	slice []any
+}
+
+func (s *stack[T]) push(v T) { s.slice = append(s.slice, v) }
+
+func (s *stack[T]) pop() any {
+	i := len(s.slice) - 1
+	v := s.slice[i]
+	s.slice = s.slice[:i]
+	return v
+}
+
+func (s *stack[T]) size() int { return len(s.slice) }

internal/timemachine/funccall/read.go

@@ -0,0 +1,56 @@
+package funccall
+
+import (
+	"encoding/binary"
+	"errors"
+	"io"
+	"slices"
+	"time"
+
+	"github.com/stealthrocket/timecraft/internal/stream"
+	"github.com/stealthrocket/timecraft/internal/timemachine"
+)
+
+type Trace struct {
+	Start    time.Time
+	Duration time.Duration
+	Depth    int
+	Name     string
+}
+
+type TraceReader struct {
+	Records stream.Reader[timemachine.Record]
+	records []timemachine.Record
+}
+
+func (r *TraceReader) Read(traces []Trace) (n int, err error) {
+	r.records = slices.Grow(r.records, len(traces))[:len(traces)]
+	for {
+		rn, err := stream.ReadFull(r.Records, r.records)
+		for i := 0; i < rn; i++ {
+			record := &r.records[i]
+			if record.FunctionID == ID {
+				// read depth
+				data := record.FunctionCall
+				depth := binary.BigEndian.Uint32(data)
+				data = data[4:]
+				//read duration
+				duration := binary.BigEndian.Uint64(data)
+				name := data[8:]
+				traces[n] = Trace{
+					Start:    record.Time,
+					Duration: time.Duration(duration),
+					Depth:    int(depth),
+					Name:     string(name),
+				}
+				n++
+			}
+		}
+		if n > 0 || err != nil {
+			if errors.Is(err, io.ErrUnexpectedEOF) {
+				err = io.EOF
+			}
+			return n, err
+		}
+	}
+}

profile.go

@@ -99,7 +99,7 @@ func profile(ctx context.Context, args []string) error {
 	if err != nil {
 		return err
 	}
-	runtime, err := timecraft.NewRuntime(ctx, config)
+	runtime, ctx, err := timecraft.NewRuntime(ctx, config)
 	if err != nil {
 		return err
 	}

replay.go

@@ -48,7 +48,7 @@ func replay(ctx context.Context, args []string) error {
 	if err != nil {
 		return err
 	}
-	runtime, err := timecraft.NewRuntime(ctx, config)
+	runtime, ctx, err := timecraft.NewRuntime(ctx, config)
 	if err != nil {
 		return err
 	}

run.go

@@ -94,7 +94,7 @@ func run(ctx context.Context, args []string) error {
 	if err != nil {
 		return err
 	}
-	runtime, err := timecraft.NewRuntime(ctx, config)
+	runtime, ctx, err := timecraft.NewRuntime(ctx, config)
 	if err != nil {
 		return err
 	}

trace.go

@@ -16,6 +16,7 @@ import (
 	"github.com/stealthrocket/timecraft/internal/stream"
 	"github.com/stealthrocket/timecraft/internal/timecraft"
 	"github.com/stealthrocket/timecraft/internal/timemachine"
+	"github.com/stealthrocket/timecraft/internal/timemachine/funccall"
 )
 
 const traceUsage = `
@@ -33,7 +34,7 @@ Options:
 type layer struct {
 	typ   string
 	alt   []string
-	trace func(io.Writer, outputFormat, string, stream.Reader[tracing.Event]) error
+	trace func(io.Writer, outputFormat, string, stream.Reader[timemachine.Record]) error
 }
 
 var layers = [...]layer{
@@ -47,6 +48,11 @@ var layers = [...]layer{
 		alt:   []string{"req", "reqs", "requests"},
 		trace: traceRequest,
 	},
+	{
+		typ:   "functions",
+		alt:   []string{"f", "func", "funcs"},
+		trace: traceFunctions,
+	},
 }
 
 func findLayer(typ string) (*layer, error) {
@@ -81,6 +87,7 @@ func trace(ctx context.Context, args []string) error {
 
    $ timecraft trace network <process id> ...
    $ timecraft trace request <process id> ...
+   $ timecraft trace function <process id> ...
 `)
 		return exitCode(2)
 	}
@@ -129,12 +136,11 @@ func trace(ctx context.Context, args []string) error {
 		format = "%+v"
 	}
 
-	return layer.trace(os.Stdout, output, format, &tracing.EventReader{
-		Records: timemachine.NewLogRecordReader(logReader),
-	})
+	return layer.trace(os.Stdout, output, format, timemachine.NewLogRecordReader(logReader))
 }
 
-func traceNetwork(w io.Writer, output outputFormat, format string, events stream.Reader[tracing.Event]) error {
+func traceNetwork(w io.Writer, output outputFormat, format string, records stream.Reader[timemachine.Record]) error {
+	events := &tracing.EventReader{Records: records}
 	var writer stream.WriteCloser[tracing.Event]
 	switch output {
 	case "json":
@@ -152,7 +158,8 @@ func traceNetwork(w io.Writer, output outputFormat, format string, events stream
 	return err
 }
 
-func traceRequest(w io.Writer, output outputFormat, format string, events stream.Reader[tracing.Event]) error {
+func traceRequest(w io.Writer, output outputFormat, format string, records stream.Reader[timemachine.Record]) error {
+	events := &tracing.EventReader{Records: records}
 	var writer stream.WriteCloser[tracing.Exchange]
 	switch output {
 	case "json":
@@ -177,3 +184,11 @@ func traceRequest(w io.Writer, output outputFormat, format string, events stream
 	})
 	return err
 }
+
+func traceFunctions(w io.Writer, output outputFormat, format string, records stream.Reader[timemachine.Record]) error {
+	switch output {
+	case "json", "yaml":
+		return fmt.Errorf("%s output is not implemented for function tracing", output)
+	}
+	return funccall.Format(w, records)
+}