add payload with original type to payload stats

rpc_util.go

@@ -270,6 +270,7 @@ func encode(c Codec, msg interface{}, cp Compressor, cbuf *bytes.Buffer, outgoin
 			return nil, err
 		}
 		if outgoingPayloadStats != nil {
+			outgoingPayloadStats.Payload = msg
 			outgoingPayloadStats.Data = b
 			outgoingPayloadStats.Length = len(b)
 		}
@@ -350,6 +351,7 @@ func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{
 		return Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
 	}
 	if incomingPayloadStats != nil {
+		incomingPayloadStats.Payload = m
 		incomingPayloadStats.Data = d
 		incomingPayloadStats.Length = len(d)
 	}

server.go

@@ -681,6 +681,7 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
 				return err
 			}
 			if incomingPayloadStats != nil {
+				incomingPayloadStats.Payload = v
 				incomingPayloadStats.Data = req
 				incomingPayloadStats.Length = len(req)
 				stats.Handle(stream.Context(), incomingPayloadStats)

stats/stats.go

@@ -70,6 +70,8 @@ func (s *InitStats) isStats() {}
 type IncomingPayloadStats struct {
 	// IsClient indicates if this stats is a client stats.
 	IsClient bool
+	// Payload is the payload with original type.
+	Payload interface{}
 	// Data is the unencrypted message payload.
 	Data []byte
 	// Length is the length of uncompressed data.
@@ -106,6 +108,8 @@ func (s *IncomingTrailerStats) isStats() {}
 type OutgoingPayloadStats struct {
 	// IsClient indicates if this stats is a client stats.
 	IsClient bool
+	// Payload is the payload with original type.
+	Payload interface{}
 	// Data is the unencrypted message payload.
 	Data []byte
 	// Length is the length of uncompressed data.

stats/stats_test.go

@@ -37,6 +37,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"reflect"
 	"sync"
 	"testing"
 
@@ -349,10 +350,13 @@ func checkIncomingPayloadStats(t *testing.T, d *gotData, e *expectedData) {
 		t.Fatalf("st IsClient = true, want false")
 	}
 	b, err := proto.Marshal(e.incoming[e.expectedInIdx])
-	e.expectedInIdx++
 	if err != nil {
 		t.Fatalf("failed to marshal message: %v", err)
 	}
+	if reflect.TypeOf(st.Payload) != reflect.TypeOf(e.incoming[e.expectedInIdx]) {
+		t.Fatalf("st.Payload = %T, want %T", st.Payload, e.incoming[e.expectedInIdx])
+	}
+	e.expectedInIdx++
 	if string(st.Data) != string(b) {
 		t.Fatalf("st.Data = %v, want %v", st.Data, b)
 	}
@@ -420,10 +424,13 @@ func checkOutgoingPayloadStats(t *testing.T, d *gotData, e *expectedData) {
 		t.Fatalf("st IsClient = true, want false")
 	}
 	b, err := proto.Marshal(e.outgoing[e.expectedOutIdx])
-	e.expectedOutIdx++
 	if err != nil {
 		t.Fatalf("failed to marshal message: %v", err)
 	}
+	if reflect.TypeOf(st.Payload) != reflect.TypeOf(e.outgoing[e.expectedOutIdx]) {
+		t.Fatalf("st.Payload = %T, want %T", st.Payload, e.outgoing[e.expectedOutIdx])
+	}
+	e.expectedOutIdx++
 	if string(st.Data) != string(b) {
 		t.Fatalf("st.Data = %v, want %v", st.Data, b)
 	}