Make gRPC logging optional via a custom interface

middleware/grpc_logging.go

@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"errors"
 	"time"
 
 	"golang.org/x/net/context"
@@ -16,6 +17,11 @@ const (
 	errorKey = "err"
 )
 
+// If an error implements this interface, it will get called and GRPCServerLog will do nothing.
+type CustomLog interface {
+	LogOperation(ctx context.Context, _ logging.Interface, method string, duration time.Duration)
+}
+
 // GRPCServerLog logs grpc requests, errors, and latency.
 type GRPCServerLog struct {
 	Log logging.Interface
@@ -31,6 +37,11 @@ func (s GRPCServerLog) UnaryServerInterceptor(ctx context.Context, req interface
 	if err == nil && s.DisableRequestSuccessLog {
 		return resp, nil
 	}
+	var customLog CustomLog
+	if errors.As(err, &customLog) {
+		customLog.LogOperation(ctx, s.Log, info.FullMethod, time.Since(begin))
+		return resp, err
+	}
 
 	entry := user.LogWith(ctx, s.Log).WithFields(logging.Fields{"method": info.FullMethod, "duration": time.Since(begin)})
 	if err != nil {

middleware/grpc_logging_test.go

@@ -1,11 +1,15 @@
 package middleware
 
 import (
+	"bytes"
 	"context"
+	"errors"
 	"testing"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 
 	"github.com/weaveworks/common/logging"
@@ -28,3 +32,52 @@ func BenchmarkGRPCServerLog_UnaryServerInterceptor_NoError(b *testing.B) {
 		_, _ = l.UnaryServerInterceptor(ctx, nil, info, handler)
 	}
 }
+
+type doNotLogError struct{ Err error }
+
+func (i doNotLogError) Error() string { return i.Err.Error() }
+func (i doNotLogError) Unwrap() error { return i.Err }
+func (i doNotLogError) LogOperation(_ context.Context, _ logging.Interface, _ string, _ time.Duration) {
+	// no-op
+}
+
+func TestGrpcLogging(t *testing.T) {
+	ctx := context.Background()
+	info := &grpc.UnaryServerInfo{FullMethod: "Test"}
+	for _, tc := range []struct {
+		err         error
+		logContains []string
+	}{{
+		err:         context.Canceled,
+		logContains: []string{"level=debug", "context canceled"},
+	}, {
+		err:         errors.New("yolo"),
+		logContains: []string{"level=warn", "err=yolo"},
+	}, {
+		err:         nil,
+		logContains: []string{"level=debug", "method=Test"},
+	}, {
+		err:         doNotLogError{Err: errors.New("yolo")},
+		logContains: nil,
+	}} {
+		t.Run("", func(t *testing.T) {
+			buf := bytes.NewBuffer(nil)
+			logger := logging.GoKit(log.NewLogfmtLogger(buf))
+			l := GRPCServerLog{Log: logger, WithRequest: true, DisableRequestSuccessLog: false}
+
+			handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+				return nil, tc.err
+			}
+
+			_, err := l.UnaryServerInterceptor(ctx, nil, info, handler)
+			require.ErrorIs(t, tc.err, err)
+
+			if len(tc.logContains) == 0 {
+				require.Empty(t, buf)
+			}
+			for _, content := range tc.logContains {
+				require.Contains(t, buf.String(), content)
+			}
+		})
+	}
+}