From 078d9d76e0f03cf20480f5e9afa76484f2701a30 Mon Sep 17 00:00:00 2001
From: kevinye202 <kevin.ye@robinhood.com>
Date: Thu, 12 Dec 2024 03:27:16 -0500
Subject: [PATCH] add grpc code to metrics replicator (#231)

* add grpc code to metrics replicator

* fix style
---
 .../replication/metrics_blob_replicator.go    |  2 +-
 .../metrics_blob_replicator_test.go           | 53 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/pkg/blobstore/replication/metrics_blob_replicator.go b/pkg/blobstore/replication/metrics_blob_replicator.go
index d59dd84b..a884f3bc 100644
--- a/pkg/blobstore/replication/metrics_blob_replicator.go
+++ b/pkg/blobstore/replication/metrics_blob_replicator.go
@@ -27,7 +27,7 @@ var (
 			Help:      "Amount of time spent per operation on blob replicator, in seconds.",
 			Buckets:   util.DecimalExponentialBuckets(-3, 6, 2),
 		},
-		[]string{"storage_type", "operation"})
+		[]string{"storage_type", "operation", "grpc_code"})
 
 	blobReplicatorOperationsBlobSizeBytes = prometheus.NewHistogramVec(
 		prometheus.HistogramOpts{
diff --git a/pkg/blobstore/replication/metrics_blob_replicator_test.go b/pkg/blobstore/replication/metrics_blob_replicator_test.go
index 394a4775..b26282d8 100644
--- a/pkg/blobstore/replication/metrics_blob_replicator_test.go
+++ b/pkg/blobstore/replication/metrics_blob_replicator_test.go
@@ -1,12 +1,20 @@
 package replication_test
 
 import (
+	"context"
 	"testing"
+	"time"
 
-	"github.com/stretchr/testify/require"
-
+	remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
 	"github.com/buildbarn/bb-storage/internal/mock"
+	"github.com/buildbarn/bb-storage/pkg/blobstore/buffer"
 	"github.com/buildbarn/bb-storage/pkg/blobstore/replication"
+	"github.com/buildbarn/bb-storage/pkg/digest"
+	"github.com/stretchr/testify/require"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
 	"go.uber.org/mock/gomock"
 )
 
@@ -23,3 +31,44 @@ func TestNewMetricsBlobReplicator(t *testing.T) {
 	metricsReplicator := replication.NewMetricsBlobReplicator(mockReplicator, mockClock, storageTypeName)
 	require.NotNil(t, metricsReplicator)
 }
+
+func TestMetricsBlobReplicatorLabelCardinality(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClock := mock.NewMockClock(ctrl)
+	mockReplicator := mock.NewMockBlobReplicator(ctrl)
+
+	// Set up fixed times for predictable testing
+	startTime := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC)
+	endTime := startTime.Add(1 * time.Second)
+	mockClock.EXPECT().Now().Return(startTime).AnyTimes()
+	mockClock.EXPECT().Now().Return(endTime).AnyTimes()
+
+	storageTypeName := "cas"
+	metricsReplicator := replication.NewMetricsBlobReplicator(
+		mockReplicator,
+		mockClock,
+		storageTypeName,
+	)
+
+	// Test ReplicateSingle
+	d := digest.MustNewDigest("hello", remoteexecution.DigestFunction_MD5, "8b1a9953c4611296a827abf8c47804d7", 5)
+	mockReplicator.EXPECT().
+		ReplicateSingle(gomock.Any(), d).
+		Return(buffer.NewBufferFromError(status.Error(codes.NotFound, "not found")))
+
+	// This should not panic due to label cardinality mismatch
+	b := metricsReplicator.ReplicateSingle(context.Background(), d)
+	b.Discard()
+
+	// Test ReplicateMultiple
+	digests := digest.NewSetBuilder().Add(d).Build()
+	mockReplicator.EXPECT().
+		ReplicateMultiple(gomock.Any(), digests).
+		Return(status.Error(codes.Internal, "internal error"))
+
+	// This should not panic due to label cardinality mismatch
+	err := metricsReplicator.ReplicateMultiple(context.Background(), digests)
+	require.Error(t, err)
+}