client: Add context parameter and enable tracing support

Add a `context.Context` parameter to the client `NewAgentClient()` API and enable gRPC tracing if the specified context contains an opentracing span. Fixes #327. Signed-off-by: James O. D. Hunt <[email protected]>
kata-containers · Aug 20, 2018 · 6d26d61 · 6d26d61
1 parent cd8f37b
commit 6d26d61
Show file tree

Hide file tree

Showing 21 changed files with 2,010 additions and 13 deletions.
diff --git a/Gopkg.lock b/Gopkg.lock
diff --git a/protocols/client/client.go b/protocols/client/client.go
@@ -14,8 +14,10 @@ import (
 	"strings"
 	"time"
 
+	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
 	"github.com/hashicorp/yamux"
 	"github.com/mdlayher/vsock"
+	opentracing "github.com/opentracing/opentracing-go"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	grpcStatus "google.golang.org/grpc/status"
@@ -54,14 +56,28 @@ type dialer func(string, time.Duration) (net.Conn, error)
 //   - unix://<unix socket path>
 //   - vsock://<cid>:<port>
 //   - <unix socket path>
-func NewAgentClient(sock string, enableYamux bool) (*AgentClient, error) {
+func NewAgentClient(ctx context.Context, sock string, enableYamux bool) (*AgentClient, error) {
 	grpcAddr, parsedAddr, err := parse(sock)
 	if err != nil {
 		return nil, err
 	}
 	dialOpts := []grpc.DialOption{grpc.WithInsecure(), grpc.WithBlock()}
 	dialOpts = append(dialOpts, grpc.WithDialer(agentDialer(parsedAddr, enableYamux)))
-	ctx := context.Background()
+
+	var tracer opentracing.Tracer
+
+	span := opentracing.SpanFromContext(ctx)
+
+	// If the context contains a trace span, trace all client comms
+	if span != nil {
+		tracer = span.Tracer()
+
+		dialOpts = append(dialOpts,
+			grpc.WithUnaryInterceptor(otgrpc.OpenTracingClientInterceptor(tracer)))
+		dialOpts = append(dialOpts,
+			grpc.WithStreamInterceptor(otgrpc.OpenTracingStreamClientInterceptor(tracer)))
+	}
+
 	ctx, cancel := context.WithTimeout(ctx, defaultDialTimeout)
 	defer cancel()
 	conn, err := grpc.DialContext(ctx, grpcAddr, dialOpts...)

diff --git a/protocols/client/client_test.go b/protocols/client/client_test.go
@@ -93,7 +93,7 @@ func agentClientTest(t *testing.T, sock string, success, enableYamux bool, expec
 	defer func() {
 		defaultDialTimeout = dialTimeout
 	}()
-	cli, err := NewAgentClient(sock, enableYamux)
+	cli, err := NewAgentClient(context.Background(), sock, enableYamux)
 	if success {
 		assert.Nil(t, err, "Failed to create new agent client: %s", err)
 	} else if !success {

diff --git a/vendor/github.com/grpc-ecosystem/grpc-opentracing/LICENSE b/vendor/github.com/grpc-ecosystem/grpc-opentracing/LICENSE
diff --git a/vendor/github.com/grpc-ecosystem/grpc-opentracing/PATENTS b/vendor/github.com/grpc-ecosystem/grpc-opentracing/PATENTS