Add `deadline_type` prometheus label to counters

This adds the `deadline_type` prometheus label to the counters from
`go-grpc-prometheus`. The possible values for `deadline_type` are,
`none`, `unknown`, `limited` and `regular`:

- `none`: is set when there is no deadline on the context.
- `unknown`: the default value, this is used when the `deadline_type`
was not included in the headers but a deadline was set
- `limited`: is set when the deadline was shorter than what is
configured. For example when the deadline would exceed the worker
timeout
- `regular`: is set when the deadline is what was configured for the
call in GitLab's settings

3 files changed, 139 insertions, 12 deletions

diff --git a/changelogs/unreleased/bvl-add-deadline-type-prom-label.yml b/changelogs/unreleased/bvl-add-deadline-type-prom-label.yml
new file mode 100644
index 000000000..633d6615e
--- /dev/null
+++ b/changelogs/unreleased/bvl-add-deadline-type-prom-label.yml
@@ -0,0 +1,5 @@
+---
+title: Add deadline_type prometheus label
+merge_request: 1737
+author:
+type: other
diff --git a/internal/middleware/metadatahandler/metadatahandler.go b/internal/middleware/metadatahandler/metadatahandler.go
index b4af28275..0581e82da 100644
--- a/internal/middleware/metadatahandler/metadatahandler.go
+++ b/internal/middleware/metadatahandler/metadatahandler.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	grpc_ctxtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/prometheus/client_golang/prometheus"
 	gitalyauth "gitlab.com/gitlab-org/gitaly/auth"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
@@ -18,16 +19,17 @@ var (
 			Namespace: "gitaly",
 			Subsystem: "service",
 			Name:      "client_requests_total",
-			Help:      "Counter of client requests received by client, call_site, auth version, and response code",
+			Help:      "Counter of client requests received by client, call_site, auth version, response code and deadline_type",
 		},
-		[]string{"client_name", "call_site", "auth_version", "grpc_code"},
+		[]string{"client_name", "call_site", "auth_version", "grpc_code", "deadline_type"},
 	)
 )
 
 type metadataTags struct {
-	clientName  string
-	callSite    string
-	authVersion string
+	clientName   string
+	callSite     string
+	authVersion  string
+	deadlineType string
 }
 
 func init() {
@@ -43,6 +45,9 @@ const ClientNameKey = "grpc.meta.client_name"
 // AuthVersionKey is the key used in ctx_tags to store the auth version
 const AuthVersionKey = "grpc.meta.auth_version"
 
+// DeadlineTypeKey is the key used in ctx_tags to store the deadline type
+const DeadlineTypeKey = "grpc.meta.deadline_type"
+
 const correlationIDKey = "correlation_id"
 
 // Unknown client and feature. Matches the prometheus grpc unknown value
@@ -62,9 +67,10 @@ func getFromMD(md metadata.MD, header string) string {
 // using `unknown` if a value is not set
 func addMetadataTags(ctx context.Context) metadataTags {
 	metaTags := metadataTags{
-		clientName:  unknownValue,
-		callSite:    unknownValue,
-		authVersion: unknownValue,
+		clientName:   unknownValue,
+		callSite:     unknownValue,
+		authVersion:  unknownValue,
+		deadlineType: unknownValue,
 	}
 
 	md, ok := metadata.FromIncomingContext(ctx)
@@ -86,6 +92,17 @@ func addMetadataTags(ctx context.Context) metadataTags {
 		tags.Set(ClientNameKey, metadata)
 	}
 
+	metadata = getFromMD(md, "deadline_type")
+	_, deadlineSet := ctx.Deadline()
+	if !deadlineSet {
+		metaTags.deadlineType = "none"
+	} else if metadata != "" {
+		metaTags.deadlineType = metadata
+	}
+
+	// Set the deadline type in the logs
+	tags.Set(DeadlineTypeKey, metaTags.deadlineType)
+
 	authInfo, _ := gitalyauth.ExtractAuthInfo(ctx)
 	if authInfo != nil {
 		metaTags.authVersion = authInfo.Version
@@ -101,14 +118,19 @@ func addMetadataTags(ctx context.Context) metadataTags {
 	return metaTags
 }
 
+func reportWithPrometheusLabels(metaTags metadataTags, err error) {
+	grpcCode := helper.GrpcCode(err)
+	requests.WithLabelValues(metaTags.clientName, metaTags.callSite, metaTags.authVersion, grpcCode.String(), metaTags.deadlineType).Inc()
+	grpc_prometheus.WithConstLabels(prometheus.Labels{"deadline_type": metaTags.deadlineType})
+}
+
 // UnaryInterceptor returns a Unary Interceptor
 func UnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 	metaTags := addMetadataTags(ctx)
 
 	res, err := handler(ctx, req)
 
-	grpcCode := helper.GrpcCode(err)
-	requests.WithLabelValues(metaTags.clientName, metaTags.callSite, metaTags.authVersion, grpcCode.String()).Inc()
+	reportWithPrometheusLabels(metaTags, err)
 
 	return res, err
 }
@@ -120,8 +142,7 @@ func StreamInterceptor(srv interface{}, stream grpc.ServerStream, info *grpc.Str
 
 	err := handler(srv, stream)
 
-	grpcCode := helper.GrpcCode(err)
-	requests.WithLabelValues(metaTags.clientName, metaTags.callSite, metaTags.authVersion, grpcCode.String()).Inc()
+	reportWithPrometheusLabels(metaTags, err)
 
 	return err
 }
diff --git a/internal/middleware/metadatahandler/metadatahandler_test.go b/internal/middleware/metadatahandler/metadatahandler_test.go
new file mode 100644
index 000000000..82f266f38
--- /dev/null
+++ b/internal/middleware/metadatahandler/metadatahandler_test.go
@@ -0,0 +1,101 @@
+package metadatahandler
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestAddMetadataTags(t *testing.T) {
+	baseContext, cancel := testhelper.Context()
+	defer cancel()
+
+	testCases := []struct {
+		desc             string
+		metadata         metadata.MD
+		deadline         bool
+		expectedMetatags metadataTags
+	}{
+		{
+			desc:     "empty metadata",
+			metadata: metadata.Pairs(),
+			deadline: false,
+			expectedMetatags: metadataTags{
+				clientName:   unknownValue,
+				callSite:     unknownValue,
+				authVersion:  unknownValue,
+				deadlineType: "none",
+			},
+		},
+		{
+			desc:     "context containing metadata",
+			metadata: metadata.Pairs("call_site", "testsite"),
+			deadline: false,
+			expectedMetatags: metadataTags{
+				clientName:   unknownValue,
+				callSite:     "testsite",
+				authVersion:  unknownValue,
+				deadlineType: "none",
+			},
+		},
+		{
+			desc:     "context containing metadata and a deadline",
+			metadata: metadata.Pairs("call_site", "testsite"),
+			deadline: true,
+			expectedMetatags: metadataTags{
+				clientName:   unknownValue,
+				callSite:     "testsite",
+				authVersion:  unknownValue,
+				deadlineType: unknownValue,
+			},
+		},
+		{
+			desc:     "context containing metadata and a deadline type",
+			metadata: metadata.Pairs("deadline_type", "regular"),
+			deadline: true,
+			expectedMetatags: metadataTags{
+				clientName:   unknownValue,
+				callSite:     unknownValue,
+				authVersion:  unknownValue,
+				deadlineType: "regular",
+			},
+		},
+		{
+			desc:     "a context without deadline but with deadline type",
+			metadata: metadata.Pairs("deadline_type", "regular"),
+			deadline: false,
+			expectedMetatags: metadataTags{
+				clientName:   unknownValue,
+				callSite:     unknownValue,
+				authVersion:  unknownValue,
+				deadlineType: "none",
+			},
+		},
+		{
+			desc:     "with a context containing metadata",
+			metadata: metadata.Pairs("deadline_type", "regular", "client_name", "rails"),
+			deadline: true,
+			expectedMetatags: metadataTags{
+				clientName:   "rails",
+				callSite:     unknownValue,
+				authVersion:  unknownValue,
+				deadlineType: "regular",
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.desc, func(t *testing.T) {
+			ctx := metadata.NewIncomingContext(baseContext, testCase.metadata)
+			if testCase.deadline {
+				ctx, cancel = context.WithDeadline(ctx, time.Now().Add(50*time.Millisecond))
+				defer cancel()
+			}
+			require.Equal(t, testCase.expectedMetatags, addMetadataTags(ctx))
+		})
+	}
+}