Merge branch 'pks-catfile-batched-tracing' into 'master'

catfile: Refactor tracing to allow for batching

See merge request gitlab-org/gitaly!4031

3 files changed, 69 insertions, 21 deletions

diff --git a/internal/git/catfile/object_info_reader.go b/internal/git/catfile/object_info_reader.go
index a61a7e29d..59d28a396 100644
--- a/internal/git/catfile/object_info_reader.go
+++ b/internal/git/catfile/object_info_reader.go
@@ -166,19 +166,16 @@ func (o *objectInfoReader) isDirty() bool {
 }
 
 func (o *objectInfoReader) Info(ctx context.Context, revision git.Revision) (*ObjectInfo, error) {
-	finish := startSpan(o.creationCtx, ctx, "Batch.Info", revision)
+	trace, finish := startTrace(ctx, o.creationCtx, o.counter, "catfile.Info")
 	defer finish()
 
-	if o.counter != nil {
-		o.counter.WithLabelValues("info").Inc()
-	}
-
 	o.Lock()
 	defer o.Unlock()
 
 	if _, err := fmt.Fprintln(o.cmd, revision.String()); err != nil {
 		return nil, err
 	}
+	trace.recordRequest("info")
 
 	return ParseObjectInfo(o.stdout)
 }
diff --git a/internal/git/catfile/object_reader.go b/internal/git/catfile/object_reader.go
index 51d15f269..af7d36a7f 100644
--- a/internal/git/catfile/object_reader.go
+++ b/internal/git/catfile/object_reader.go
@@ -130,7 +130,7 @@ func (o *objectReader) Object(
 	ctx context.Context,
 	revision git.Revision,
 ) (*Object, error) {
-	finish := startSpan(o.creationCtx, ctx, "Batch.Object", revision)
+	trace, finish := startTrace(ctx, o.creationCtx, o.counter, "catfile.Object")
 	defer finish()
 
 	o.Lock()
@@ -156,10 +156,7 @@ func (o *objectReader) Object(
 	if err != nil {
 		return nil, err
 	}
-
-	if o.counter != nil {
-		o.counter.WithLabelValues(oi.Type).Inc()
-	}
+	trace.recordRequest(oi.Type)
 
 	o.n = oi.Size + 1
 
diff --git a/internal/git/catfile/tracing.go b/internal/git/catfile/tracing.go
index 77a51304e..6055a9d4c 100644
--- a/internal/git/catfile/tracing.go
+++ b/internal/git/catfile/tracing.go
@@ -4,24 +4,78 @@ import (
 	"context"
 
 	"github.com/opentracing/opentracing-go"
-	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
+	"github.com/prometheus/client_golang/prometheus"
 	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
-func revisionTag(revision git.Revision) opentracing.Tag {
-	return opentracing.Tag{Key: "revision", Value: revision}
+type trace struct {
+	rpcSpan   opentracing.Span
+	cacheSpan opentracing.Span
+	counter   *prometheus.CounterVec
+
+	requests map[string]int
 }
 
-func correlationIDTag(ctx context.Context) opentracing.Tag {
-	return opentracing.Tag{Key: "correlation_id", Value: correlation.ExtractFromContext(ctx)}
+// startTrace starts a new tracing span and updates metrics according to how many requests have been
+// done during that trace. This must be called with two contexts: first the per-RPC context, which
+// is the context of the current RPC call. And then the cache context, which is the decorrelated
+// context for cached catfile processes. Spans are then created for both contexts.
+func startTrace(
+	rpcCtx context.Context,
+	cacheCtx context.Context,
+	counter *prometheus.CounterVec,
+	methodName string,
+) (*trace, func()) {
+	rpcSpan, _ := opentracing.StartSpanFromContext(rpcCtx, methodName)
+
+	// The per-RPC and cached context will be the same in case the process for which we're
+	// creating the tracing span for is not cached, and we shouldn't create the same span twice.
+	var cacheSpan opentracing.Span
+	if rpcCtx != cacheCtx {
+		cacheSpan, _ = opentracing.StartSpanFromContext(cacheCtx, methodName, opentracing.Tag{
+			Key: "correlation_id", Value: correlation.ExtractFromContext(rpcCtx),
+		})
+	}
+
+	trace := &trace{
+		rpcSpan:   rpcSpan,
+		cacheSpan: cacheSpan,
+		counter:   counter,
+		requests: map[string]int{
+			"blob":   0,
+			"commit": 0,
+			"tree":   0,
+			"tag":    0,
+			"info":   0,
+		},
+	}
+
+	return trace, trace.finish
 }
 
-func startSpan(innerCtx context.Context, outerCtx context.Context, methodName string, revision git.Revision) func() {
-	innerSpan, ctx := opentracing.StartSpanFromContext(innerCtx, methodName, revisionTag(revision))
-	outerSpan, _ := opentracing.StartSpanFromContext(outerCtx, methodName, revisionTag(revision), correlationIDTag(ctx))
+func (t *trace) recordRequest(requestType string) {
+	t.requests[requestType]++
+}
+
+func (t *trace) finish() {
+	for requestType, requestCount := range t.requests {
+		if requestCount == 0 {
+			continue
+		}
+
+		tag := opentracing.Tag{Key: requestType, Value: requestCount}
+		tag.Set(t.rpcSpan)
+		if t.cacheSpan != nil {
+			tag.Set(t.cacheSpan)
+		}
+
+		if t.counter != nil {
+			t.counter.WithLabelValues(requestType).Add(float64(requestCount))
+		}
+	}
 
-	return func() {
-		innerSpan.Finish()
-		outerSpan.Finish()
+	t.rpcSpan.Finish()
+	if t.cacheSpan != nil {
+		t.cacheSpan.Finish()
 	}
 }