Merge pull request #211 from bergwolf/grpc-req

agent: add server interceptor to log grpc requests

agent.go

@@ -22,11 +22,13 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
 	pb "github.com/kata-containers/agent/protocols/grpc"
 	"github.com/opencontainers/runc/libcontainer"
 	"github.com/opencontainers/runc/libcontainer/configs"
 	_ "github.com/opencontainers/runc/libcontainer/nsenter"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/net/context"
 	"golang.org/x/sys/unix"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -60,17 +62,18 @@ type container struct {
 type sandbox struct {
 	sync.RWMutex
 
-	id          string
-	running     bool
-	noPivotRoot bool
-	containers  map[string]*container
-	channel     channel
-	network     network
-	wg          sync.WaitGroup
-	sharedPidNs namespace
-	mounts      []string
-	subreaper   reaper
-	server      *grpc.Server
+	id              string
+	running         bool
+	noPivotRoot     bool
+	enableGrpcTrace bool
+	containers      map[string]*container
+	channel         channel
+	network         network
+	wg              sync.WaitGroup
+	sharedPidNs     namespace
+	mounts          []string
+	subreaper       reaper
+	server          *grpc.Server
 }
 
 type namespace struct {
@@ -425,7 +428,7 @@ func (s *sandbox) initLogger() error {
 	if err := config.getConfig(kernelCmdlineFile); err != nil {
 		agentLog.WithError(err).Warn("Failed to get config from kernel cmdline")
 	}
-	config.applyConfig()
+	config.applyConfig(s)
 
 	return announce()
 }
@@ -441,13 +444,44 @@ func (s *sandbox) initChannel() error {
 	return nil
 }
 
+func grpcTracer(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
+	message := req.(proto.Message)
+	agentLog.WithFields(logrus.Fields{
+		"request": info.FullMethod,
+		"req":     message.String()}).Debug("new request")
+
+	start := time.Now()
+	resp, err = handler(ctx, req)
+	elapsed := time.Now().Sub(start)
+
+	message = resp.(proto.Message)
+	logger := agentLog.WithFields(logrus.Fields{
+		"request":  info.FullMethod,
+		"duration": elapsed.String(),
+		"resp":     message.String()})
+	if err != nil {
+		logger = logger.WithError(err)
+	}
+	logger.Debug("request end")
+
+	return resp, err
+}
+
 func (s *sandbox) startGRPC() {
 	grpcImpl := &agentGRPC{
 		sandbox: s,
 		version: version,
 	}
 
-	grpcServer := grpc.NewServer()
+	var grpcServer *grpc.Server
+	if s.enableGrpcTrace {
+		agentLog.Info("Enable grpc tracing")
+		opt := grpc.UnaryInterceptor(grpcTracer)
+		grpcServer = grpc.NewServer(opt)
+	} else {
+		grpcServer = grpc.NewServer()
+	}
+
 	pb.RegisterAgentServiceServer(grpcServer, grpcImpl)
 	pb.RegisterHealthServer(grpcServer, grpcImpl)
 	s.server = grpcServer

agent_test.go

@@ -11,8 +11,12 @@ import (
 	"reflect"
 	"testing"
 
+	"google.golang.org/grpc"
+
+	pb "github.com/kata-containers/agent/protocols/grpc"
 	"github.com/opencontainers/runc/libcontainer"
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/net/context"
 )
 
 const (
@@ -268,3 +272,28 @@ func TestStartStopGRPCServer(t *testing.T) {
 	s.stopGRPC()
 	assert.Nil(t, s.server, "failed stopping grpc server")
 }
+
+func TestSettingGrpcTracer(t *testing.T) {
+	_, out, err := os.Pipe()
+	assert.Nil(t, err, "%v", err)
+
+	s := &sandbox{
+		containers:      make(map[string]*container),
+		channel:         &serialChannel{serialConn: out},
+		enableGrpcTrace: true,
+	}
+
+	s.startGRPC()
+	assert.NotNil(t, s.server, "failed starting grpc server")
+
+	s.stopGRPC()
+	assert.Nil(t, s.server, "failed stopping grpc server")
+}
+
+func TestGrpcTracer(t *testing.T) {
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return &pb.HealthCheckResponse{}, nil
+	}
+	_, err := grpcTracer(context.Background(), &pb.CheckRequest{}, &grpc.UnaryServerInfo{}, handler)
+	assert.Nil(t, err, "failed to trace grpc request: %v", err)
+}

config.go

@@ -55,8 +55,11 @@ func (c *agentConfig) getConfig(cmdLineFile string) error {
 	return nil
 }
 
-func (c *agentConfig) applyConfig() {
+func (c *agentConfig) applyConfig(s *sandbox) {
 	agentLog.Logger.SetLevel(c.logLevel)
+	if c.logLevel == logrus.DebugLevel {
+		s.enableGrpcTrace = true
+	}
 }
 
 //Parse a string that represents a kernel cmdline option

config_test.go

@@ -140,3 +140,27 @@ func TestGetConfig(t *testing.T) {
 		"Log levels should be identical: got %+v, expecting %+v",
 		a.logLevel, logrus.InfoLevel)
 }
+
+func TestSetGrpcTrace(t *testing.T) {
+	assert := assert.New(t)
+
+	a := &agentConfig{}
+
+	tmpFile, err := ioutil.TempFile("", "test")
+	assert.NoError(err, "%v", err)
+	fileName := tmpFile.Name()
+
+	tmpFile.Write([]byte(logLevelFlag + "=debug"))
+	tmpFile.Close()
+
+	defer os.Remove(fileName)
+
+	err = a.getConfig(fileName)
+	assert.NoError(err, "%v", err)
+
+	s := &sandbox{}
+
+	a.applyConfig(s)
+
+	assert.True(s.enableGrpcTrace, "grpc trace should be enabled")
+}