WIP: implement trace2 tracingqmnguyen0711/export-pack-objects-metrics

4 files changed, 140 insertions, 8 deletions

diff --git a/internal/command/command.go b/internal/command/command.go
index 3bae16807..6bad6efc4 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/trace2"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/trace2/hooks"
 	"io"
 	"os"
 	"os/exec"
@@ -116,6 +118,11 @@ var (
 	// envInjector is responsible for injecting environment variables required for tracing into
 	// the child process.
 	envInjector = labkittracing.NewEnvInjector()
+
+	// trace2Hooks
+	trace2Hooks = []trace2.Hook{
+		hooks.TracingHook{},
+	}
 )
 
 const (
@@ -143,7 +150,8 @@ type Command struct {
 
 	finalizer func(*Command)
 
-	span opentracing.Span
+	trace2Manager *trace2.Manager
+	span          opentracing.Span
 
 	metricsCmd    string
 	metricsSubCmd string
@@ -228,8 +236,14 @@ func New(ctx context.Context, nameAndArgs []string, opts ...Option) (*Command, e
 	cmd.Env = AllowedEnvironment(os.Environ())
 	// Append environment variables explicitly requested by the caller.
 	cmd.Env = append(cmd.Env, cfg.environment...)
-	// And finally inject environment variables required for tracing into the command.
+	// Inject environment variables required for tracing into the command.
 	cmd.Env = envInjector(ctx, cmd.Env)
+	// And finally inject trace2 environment variable
+	command.trace2Manager = trace2.NewManager(ctx, trace2Hooks, cmd.Path, cmd.Args)
+	cmd.Env, err = command.trace2Manager.Inject(cmd.Env)
+	if err != nil {
+		return nil, fmt.Errorf("trace2 inject env: %w", err)
+	}
 
 	// Start the command in its own process group (nice for signalling)
 	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
@@ -382,6 +396,9 @@ func (c *Command) wait() {
 
 	inFlightCommandGauge.Dec()
 
+	// Handle trace2 hooks if needed
+	c.trace2Manager.Finish()
+
 	c.logProcessComplete()
 
 	// This is a bit out-of-place here given that the `commandcounter.Increment()` call is in
diff --git a/internal/trace2/hooks/tracing.go b/internal/trace2/hooks/tracing.go
new file mode 100644
index 000000000..51a02eda5
--- /dev/null
+++ b/internal/trace2/hooks/tracing.go
@@ -0,0 +1,32 @@
+package hooks
+
+import (
+	"context"
+	"fmt"
+	"github.com/opentracing/opentracing-go"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/trace2"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/tracing"
+)
+
+type TracingHook struct{}
+
+func (t TracingHook) Activate(context.Context, string, []string) (trace2.TraceHandler, bool) {
+	return func(rootCtx context.Context, trace *trace2.Trace) {
+		trace.Walk(rootCtx, func(ctx context.Context, trace *trace2.Trace) context.Context {
+			if trace.IsRoot() {
+				return rootCtx
+			}
+
+			spanName := fmt.Sprintf("git:%s", trace.Name)
+			span, parentCtx := tracing.StartSpanIfHasParent(ctx, spanName, nil, opentracing.StartTime(trace.StartTime))
+			span.SetTag("thread", trace.Thread)
+			span.SetTag("childID", trace.ChildID)
+			for key, value := range trace.Metadata {
+				span.SetTag(key, value)
+			}
+			span.FinishWithOptions(opentracing.FinishOptions{FinishTime: trace.FinishTime})
+
+			return parentCtx
+		})
+	}, true
+}
diff --git a/internal/trace2/manager.go b/internal/trace2/manager.go
new file mode 100644
index 000000000..9c0223de3
--- /dev/null
+++ b/internal/trace2/manager.go
@@ -0,0 +1,84 @@
+package trace2
+
+import (
+	"context"
+	"fmt"
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
+	"github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/labkit/correlation"
+	"os"
+)
+
+type TraceHandler func(context.Context, *Trace)
+
+type Hook interface {
+	Activate(ctx context.Context, path string, args []string) (TraceHandler, bool)
+}
+
+func NewManager(ctx context.Context, hooks []Hook, path string, args []string) *Manager {
+	manager := &Manager{
+		ctx:    ctx,
+		logger: ctxlogrus.Extract(ctx),
+	}
+	for _, hook := range hooks {
+		if handler, ok := hook.Activate(ctx, path, args); ok {
+			manager.handlers = append(manager.handlers, handler)
+		}
+	}
+	return manager
+}
+
+type Manager struct {
+	ctx      context.Context
+	logger   *logrus.Entry
+	handlers []TraceHandler
+	fd       *os.File
+}
+
+func (m *Manager) Inject(env []string) ([]string, error) {
+	if len(m.handlers) <= 0 {
+		return env, nil
+	}
+
+	fd, err := os.CreateTemp("", "gitaly-trace2")
+	if err != nil {
+		return env, fmt.Errorf("trace2 create tempfile: %w", err)
+	}
+	m.fd = fd
+
+	env = append(
+		env,
+		"GIT_TRACE2_BRIEF=true",
+		fmt.Sprintf("GIT_TRACE2_EVENT=%s", fd.Name()),
+		fmt.Sprintf("GIT_TRACE2_PARENT_SID=%s", correlation.ExtractFromContextOrGenerate(m.ctx)),
+	)
+	return env, nil
+}
+
+func (m *Manager) Finish() {
+	if len(m.handlers) <= 0 {
+		return
+	}
+
+	defer func() {
+		err := m.fd.Close()
+		if err != nil {
+			m.logger.Errorf("trace2: fail to close tempfile: %s", err)
+		}
+		err = os.Remove(m.fd.Name())
+		if err != nil {
+			m.logger.Errorf("trace2: fail to remove tempfile: %s", err)
+		}
+	}()
+
+	parser := &Parser{}
+	trace, err := parser.Parse(m.ctx, m.fd)
+	if err != nil {
+		m.logger.Errorf("trace2: fail to parse events: %s", err)
+		return
+	}
+
+	for _, handler := range m.handlers {
+		handler(m.ctx, trace)
+	}
+}
diff --git a/internal/tracing/tracing.go b/internal/tracing/tracing.go
index 4ba7c2091..9b4d74c2d 100644
--- a/internal/tracing/tracing.go
+++ b/internal/tracing/tracing.go
@@ -11,19 +11,19 @@ type Tags map[string]any
 
 // StartSpan creates a new span with name and options (mostly tags). This function is a wrapper for
 // underlying tracing libraries. This method should only be used at the entrypoint of the program.
-func StartSpan(ctx context.Context, spanName string, tags Tags) (opentracing.Span, context.Context) {
-	return opentracing.StartSpanFromContext(ctx, spanName, tagsToOpentracingTags(tags)...)
+func StartSpan(ctx context.Context, spanName string, tags Tags, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
+	return opentracing.StartSpanFromContext(ctx, spanName, tagsToOpentracingTags(opts, tags)...)
 }
 
 // StartSpanIfHasParent creates a new span if the context already has an existing span. This function
 // adds a simple validation to prevent orphan spans outside interested code paths. It returns a dummy
 // span, which acts as normal span, but does absolutely nothing and is not recorded later.
-func StartSpanIfHasParent(ctx context.Context, spanName string, tags Tags) (opentracing.Span, context.Context) {
+func StartSpanIfHasParent(ctx context.Context, spanName string, tags Tags, opts ...opentracing.StartSpanOption) (opentracing.Span, context.Context) {
 	parent := opentracing.SpanFromContext(ctx)
 	if parent == nil {
 		return &NoopSpan{}, ctx
 	}
-	return opentracing.StartSpanFromContext(ctx, spanName, tagsToOpentracingTags(tags)...)
+	return opentracing.StartSpanFromContext(ctx, spanName, tagsToOpentracingTags(opts, tags)...)
 }
 
 // DiscardSpanInContext discards the current active span from the context. This function is helpful
@@ -36,8 +36,7 @@ func DiscardSpanInContext(ctx context.Context) context.Context {
 	return opentracing.ContextWithSpan(ctx, nil)
 }
 
-func tagsToOpentracingTags(tags Tags) []opentracing.StartSpanOption {
-	var opts []opentracing.StartSpanOption
+func tagsToOpentracingTags(opts []opentracing.StartSpanOption, tags Tags) []opentracing.StartSpanOption {
 	for key, value := range tags {
 		opts = append(opts, opentracing.Tag{Key: key, Value: value})
 	}