feature/disable logging of gRPC responses (#160)

* added ability to obscure logs of http/grpc responses * updated docs * fix logging as required * fixed docs * fix logging as suggested * fix flag * fix logging finally * review comments * review comments Co-authored-by: Nikos Katirtzis <[email protected]> * review comments Co-authored-by: Nikos Katirtzis <[email protected]> * review comments Co-authored-by: Nikos Katirtzis <[email protected]> Co-authored-by: Nikos Katirtzis <[email protected]>
ExpediaGroup · Apr 21, 2021 · 3d3ffb2 · 3d3ffb2
1 parent cc68e69
commit 3d3ffb2
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 5 deletions.
diff --git a/go.mod b/go.mod
@@ -2,6 +2,7 @@ module mittens
 
 require (
 	github.com/fullstorydev/grpcurl v1.8.0
+	github.com/golang/protobuf v1.4.2
 	github.com/jhump/protoreflect v1.8.2
 	github.com/stretchr/testify v1.7.0
 	golang.org/x/net v0.0.0-20200625001655-4c5254603344

diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go
@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/fullstorydev/grpcurl"
+	"github.com/golang/protobuf/proto"
 	"github.com/jhump/protoreflect/grpcreflect"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
@@ -40,6 +41,12 @@ type Client struct {
 	descriptorSource grpcurl.DescriptorSource
 }
 
+// eventHandler is a custom event handler with the option to enable/disable logging of responses.
+type eventHandler struct {
+	grpcurl.InvocationEventHandler
+	logResponses bool
+}
+
 // NewClient returns a gRPC client.
 func NewClient(host string, insecure bool, timeoutSeconds int) Client {
 	return Client{host: host, timeoutSeconds: timeoutSeconds, insecure: insecure, connClose: func() error { return nil }}
@@ -74,7 +81,7 @@ func (c *Client) Connect(headers []string) error {
 
 // SendRequest sends a request to the gRPC server and wraps useful information into a Response object.
 // Note that the message cannot be null. Even if there is no message to be sent this needs to be set to an empty string.
-func (c *Client) SendRequest(serviceMethod string, message string, headers []string) response.Response {
+func (c *Client) SendRequest(serviceMethod string, message string, headers []string, logResponses bool) response.Response {
 	const respType = "grpc"
 	in := bytes.NewBufferString(message)
 
@@ -85,8 +92,13 @@ func (c *Client) SendRequest(serviceMethod string, message string, headers []str
 		// FIXME FATAL
 		return response.Response{Duration: time.Duration(0), Err: err, Type: respType}
 	}
-	loggingEventHandler := grpcurl.NewDefaultEventHandler(os.Stdout, c.descriptorSource, formatter, false)
 
+	delegate := &grpcurl.DefaultEventHandler{
+		Out:       os.Stdout,
+		Formatter: formatter,
+	}
+
+	loggingEventHandler := eventHandler{InvocationEventHandler: delegate, logResponses: logResponses}
 	startTime := time.Now()
 	err = grpcurl.InvokeRPC(context.Background(), c.descriptorSource, c.conn, serviceMethod, headers, loggingEventHandler, requestParser.Next)
 	endTime := time.Now()
@@ -97,6 +109,13 @@ func (c *Client) SendRequest(serviceMethod string, message string, headers []str
 	return response.Response{Duration: endTime.Sub(startTime), Err: nil, Type: respType}
 }
 
+// OnReceiveResponse overrides the default method and allows enabling/disabling logging of responses. 
+func (h eventHandler) OnReceiveResponse(msg proto.Message) {
+	if h.logResponses {
+		h.InvocationEventHandler.OnReceiveResponse(msg)
+	}
+}
+
 // Close calling close on a client that has not established connection does not return an error.
 func (c Client) Close() error {
 	log.Print("Closing gRPC client connection")

diff --git a/pkg/warmup/target.go b/pkg/warmup/target.go
@@ -83,7 +83,7 @@ func (t Target) WaitForReadinessProbe() error {
 					if connErr != nil {
 						log.Printf("gRPC readiness client connect error: %v", connErr)
 					}
-					err1 := t.readinessGrpcClient.SendRequest(request.ServiceMethod, "", nil)
+					err1 := t.readinessGrpcClient.SendRequest(request.ServiceMethod, "", nil, false)
 					if err1.Err != nil {
 						log.Printf("gRPC target not ready yet...")
 						continue

diff --git a/pkg/warmup/warmup.go b/pkg/warmup/warmup.go
@@ -100,13 +100,12 @@ func (w Warmup) GrpcWarmupWorker(wg *sync.WaitGroup, requests <-chan grpc.Reques
 	for request := range requests {
 		time.Sleep(time.Duration(requestDelayMilliseconds) * time.Millisecond)
 
-		resp := w.Target.grpcClient.SendRequest(request.ServiceMethod, request.Message, headers)
+		resp := w.Target.grpcClient.SendRequest(request.ServiceMethod, request.Message, headers, false)
 
 		if resp.Err != nil {
 			log.Printf("🔴 Error in request for %s: %v", request.ServiceMethod, resp.Err)
 		} else {
 			*requestsSentCounter++
-
 			log.Printf("🟢 %s response\t%d ms %s", resp.Type, resp.Duration/time.Millisecond, request.ServiceMethod)
 		}