binarylog: generalize binarylog's MethodLogger preparing for new observability features (#5244)

internal/binarylog/binarylog.go

@@ -31,7 +31,7 @@ import (
 // Logger is the global binary logger. It can be used to get binary logger for
 // each method.
 type Logger interface {
-	getMethodLogger(methodName string) *MethodLogger
+	GetMethodLogger(methodName string) MethodLogger
 }
 
 // binLogger is the global binary logger for the binary. One of this should be
@@ -49,17 +49,24 @@ func SetLogger(l Logger) {
 	binLogger = l
 }
 
+// GetLogger gets the binarg logger.
+//
+// Only call this at init time.
+func GetLogger() Logger {
+	return binLogger
+}
+
 // GetMethodLogger returns the methodLogger for the given methodName.
 //
 // methodName should be in the format of "/service/method".
 //
 // Each methodLogger returned by this method is a new instance. This is to
 // generate sequence id within the call.
-func GetMethodLogger(methodName string) *MethodLogger {
+func GetMethodLogger(methodName string) MethodLogger {
 	if binLogger == nil {
 		return nil
 	}
-	return binLogger.getMethodLogger(methodName)
+	return binLogger.GetMethodLogger(methodName)
 }
 
 func init() {
@@ -148,7 +155,7 @@ func (l *logger) setBlacklist(method string) error {
 //
 // Each methodLogger returned by this method is a new instance. This is to
 // generate sequence id within the call.
-func (l *logger) getMethodLogger(methodName string) *MethodLogger {
+func (l *logger) GetMethodLogger(methodName string) MethodLogger {
 	s, m, err := grpcutil.ParseMethod(methodName)
 	if err != nil {
 		grpclogLogger.Infof("binarylogging: failed to parse %q: %v", methodName, err)

internal/binarylog/binarylog_test.go

@@ -93,7 +93,7 @@ func (s) TestGetMethodLogger(t *testing.T) {
 			t.Errorf("in: %q, failed to create logger from config string", tc.in)
 			continue
 		}
-		ml := l.getMethodLogger(tc.method)
+		ml := l.GetMethodLogger(tc.method).(*methodLogger)
 		if ml == nil {
 			t.Errorf("in: %q, method logger is nil, want non-nil", tc.in)
 			continue
@@ -149,7 +149,7 @@ func (s) TestGetMethodLoggerOff(t *testing.T) {
 			t.Errorf("in: %q, failed to create logger from config string", tc.in)
 			continue
 		}
-		ml := l.getMethodLogger(tc.method)
+		ml := l.GetMethodLogger(tc.method)
 		if ml != nil {
 			t.Errorf("in: %q, method logger is non-nil, want nil", tc.in)
 		}

internal/binarylog/method_logger.go

@@ -48,7 +48,11 @@ func (g *callIDGenerator) reset() {
 var idGen callIDGenerator
 
 // MethodLogger is the sub-logger for each method.
-type MethodLogger struct {
+type MethodLogger interface {
+	Log(LogEntryConfig)
+}
+
+type methodLogger struct {
 	headerMaxLen, messageMaxLen uint64
 
 	callID          uint64
@@ -57,8 +61,8 @@ type MethodLogger struct {
 	sink Sink // TODO(blog): make this plugable.
 }
 
-func newMethodLogger(h, m uint64) *MethodLogger {
-	return &MethodLogger{
+func newMethodLogger(h, m uint64) *methodLogger {
+	return &methodLogger{
 		headerMaxLen:  h,
 		messageMaxLen: m,
 
@@ -69,8 +73,10 @@ func newMethodLogger(h, m uint64) *MethodLogger {
 	}
 }
 
-// Log creates a proto binary log entry, and logs it to the sink.
-func (ml *MethodLogger) Log(c LogEntryConfig) {
+// Build is an internal only method for building the proto message out of the
+// input event. It's made public to enable other library to reuse as much logic
+// in methodLogger as possible.
+func (ml *methodLogger) Build(c LogEntryConfig) *pb.GrpcLogEntry {
 	m := c.toProto()
 	timestamp, _ := ptypes.TimestampProto(time.Now())
 	m.Timestamp = timestamp
@@ -85,11 +91,15 @@ func (ml *MethodLogger) Log(c LogEntryConfig) {
 	case *pb.GrpcLogEntry_Message:
 		m.PayloadTruncated = ml.truncateMessage(pay.Message)
 	}
-
-	ml.sink.Write(m)
+	return m
 }
 
-func (ml *MethodLogger) truncateMetadata(mdPb *pb.Metadata) (truncated bool) {
+// Log creates a proto binary log entry, and logs it to the sink.
+func (ml *methodLogger) Log(c LogEntryConfig) {
+	ml.sink.Write(ml.Build(c))
+}
+
+func (ml *methodLogger) truncateMetadata(mdPb *pb.Metadata) (truncated bool) {
 	if ml.headerMaxLen == maxUInt {
 		return false
 	}
@@ -119,7 +129,7 @@ func (ml *MethodLogger) truncateMetadata(mdPb *pb.Metadata) (truncated bool) {
 	return truncated
 }
 
-func (ml *MethodLogger) truncateMessage(msgPb *pb.Message) (truncated bool) {
+func (ml *methodLogger) truncateMessage(msgPb *pb.Message) (truncated bool) {
 	if ml.messageMaxLen == maxUInt {
 		return false
 	}

internal/binarylog/method_logger_test.go

@@ -350,7 +350,7 @@ func (s) TestLog(t *testing.T) {
 
 func (s) TestTruncateMetadataNotTruncated(t *testing.T) {
 	testCases := []struct {
-		ml   *MethodLogger
+		ml   *methodLogger
 		mpPb *pb.Metadata
 	}{
 		{
@@ -417,7 +417,7 @@ func (s) TestTruncateMetadataNotTruncated(t *testing.T) {
 
 func (s) TestTruncateMetadataTruncated(t *testing.T) {
 	testCases := []struct {
-		ml   *MethodLogger
+		ml   *methodLogger
 		mpPb *pb.Metadata
 
 		entryLen int
@@ -478,7 +478,7 @@ func (s) TestTruncateMetadataTruncated(t *testing.T) {
 
 func (s) TestTruncateMessageNotTruncated(t *testing.T) {
 	testCases := []struct {
-		ml    *MethodLogger
+		ml    *methodLogger
 		msgPb *pb.Message
 	}{
 		{
@@ -511,7 +511,7 @@ func (s) TestTruncateMessageNotTruncated(t *testing.T) {
 
 func (s) TestTruncateMessageTruncated(t *testing.T) {
 	testCases := []struct {
-		ml    *MethodLogger
+		ml    *methodLogger
 		msgPb *pb.Message
 
 		oldLength uint32

stream.go

@@ -462,7 +462,7 @@ type clientStream struct {
 
 	retryThrottler *retryThrottler // The throttler active when the RPC began.
 
-	binlog *binarylog.MethodLogger // Binary logger, can be nil.
+	binlog binarylog.MethodLogger // Binary logger, can be nil.
 	// serverHeaderBinlogged is a boolean for whether server header has been
 	// logged. Server header will be logged when the first time one of those
 	// happens: stream.Header(), stream.Recv().
@@ -1434,7 +1434,7 @@ type serverStream struct {
 
 	statsHandler stats.Handler
 
-	binlog *binarylog.MethodLogger
+	binlog binarylog.MethodLogger
 	// serverHeaderBinlogged indicates whether server header has been logged. It
 	// will happen when one of the following two happens: stream.SendHeader(),
 	// stream.Send().