client: Improve tracing tests

Gitaly allows to set up global tracing, which may create spans for
several events. Each span denotes the start and end point, what is
happening and some arbitrary baggage which is user-defined.

We have a test which exercises this code and asserts that spans are
correctly created and registered across RPC calls, but the test is flaky
and causes lots of CI failures. The failure mode is kind of weird, in
that spans sometimes have a different order than the order we expect.

It is unclear what exactly is causing the races, and personally I have
not been able to reproduce it locally. This commit thus refactors both
testcases and makes them more explicit about what kind of data they
expect, which hopefully helps in nailing down what's happening.

One thing that _may_ be a potential source of the bug is that we now
flush the tracer earlier. Previously, we'd only close the tracer after
the testcase ran. But given that the close may also flush any pending
events, this could potentially be a source of uncertainty.

1 files changed, 109 insertions, 66 deletions

diff --git a/client/dial_test.go b/client/dial_test.go
index 61b225e8b..b1ca61c93 100644
--- a/client/dial_test.go
+++ b/client/dial_test.go
@@ -246,115 +246,158 @@ func TestDial_Correlation(t *testing.T) {
 }
 
 func TestDial_Tracing(t *testing.T) {
-	t.Run("unary", func(t *testing.T) {
-		serverSocketPath := testhelper.GetTemporaryGitalySocketFileName(t)
+	serverSocketPath := testhelper.GetTemporaryGitalySocketFileName(t)
 
-		listener, err := net.Listen("unix", serverSocketPath)
-		require.NoError(t, err)
+	listener, err := net.Listen("unix", serverSocketPath)
+	require.NoError(t, err)
 
-		grpcServer := grpc.NewServer(grpc.UnaryInterceptor(grpctracing.UnaryServerTracingInterceptor()))
-		svc := &testSvc{
-			PingMethod: func(ctx context.Context, r *proxytestdata.PingRequest) (*proxytestdata.PingResponse, error) {
-				span, _ := opentracing.StartSpanFromContext(ctx, "health")
-				defer span.Finish()
-				span.LogKV("was", "called")
-				return &proxytestdata.PingResponse{}, nil
-			},
-		}
-		proxytestdata.RegisterTestServiceServer(grpcServer, svc)
+	// This is our test service. All it does is to create additional spans
+	// which should in the end be visible when collecting all registered
+	// spans.
+	grpcServer := grpc.NewServer(
+		grpc.UnaryInterceptor(grpctracing.UnaryServerTracingInterceptor()),
+		grpc.StreamInterceptor(grpctracing.StreamServerTracingInterceptor()),
+	)
+	svc := &testSvc{
+		PingMethod: func(ctx context.Context, r *proxytestdata.PingRequest) (*proxytestdata.PingResponse, error) {
+			span, _ := opentracing.StartSpanFromContext(ctx, "nested-span")
+			defer span.Finish()
+			span.LogKV("was", "called")
+			return &proxytestdata.PingResponse{}, nil
+		},
+		PingStreamMethod: func(stream proxytestdata.TestService_PingStreamServer) error {
+			span, _ := opentracing.StartSpanFromContext(stream.Context(), "nested-span")
+			defer span.Finish()
+			span.LogKV("was", "called")
+			_, err := stream.Recv()
+			require.NoError(t, err)
+			return stream.Send(&proxytestdata.PingResponse{})
+		},
+	}
+	proxytestdata.RegisterTestServiceServer(grpcServer, svc)
 
-		go func() { assert.NoError(t, grpcServer.Serve(listener)) }()
-		defer grpcServer.Stop()
+	go func() { require.NoError(t, grpcServer.Serve(listener)) }()
+	defer grpcServer.Stop()
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
 
+	t.Run("unary", func(t *testing.T) {
 		reporter := jaeger.NewInMemoryReporter()
-		tracer, closer := jaeger.NewTracer("", jaeger.NewConstSampler(true), reporter)
-		defer closer.Close()
+		tracer, tracerCloser := jaeger.NewTracer("", jaeger.NewConstSampler(true), reporter)
 
 		defer func(old opentracing.Tracer) { opentracing.SetGlobalTracer(old) }(opentracing.GlobalTracer())
 		opentracing.SetGlobalTracer(tracer)
 
-		span := tracer.StartSpan("unary-check")
-		span = span.SetBaggageItem("service", "stub")
-
-		ctx, cancel := testhelper.Context()
-		defer cancel()
-
+		// This needs to be run after setting up the global tracer as it will cause us to
+		// create the span when executing the RPC call further down below.
 		cc, err := DialContext(ctx, "unix://"+serverSocketPath, nil)
 		require.NoError(t, err)
 		defer cc.Close()
 
-		client := proxytestdata.NewTestServiceClient(cc)
+		// We set up a "main" span here, which is going to be what the
+		// other spans inherit from. In order to check whether baggage
+		// works correctly, we also set up a "stub" baggage item which
+		// should be inherited to child contexts.
+		span := tracer.StartSpan("unary-check")
+		span = span.SetBaggageItem("service", "stub")
+		ctx := opentracing.ContextWithSpan(ctx, span)
 
-		ctx = opentracing.ContextWithSpan(ctx, span)
-		_, err = client.Ping(ctx, &proxytestdata.PingRequest{})
+		// We're now invoking the unary RPC with the span injected into
+		// the context. This should create a span that's nested into
+		// the "stream-check" span.
+		_, err = proxytestdata.NewTestServiceClient(cc).Ping(ctx, &proxytestdata.PingRequest{})
 		require.NoError(t, err)
 
 		span.Finish()
+		tracerCloser.Close()
 
 		spans := reporter.GetSpans()
 		require.Len(t, spans, 3)
-		require.Equal(t, "stub", spans[1].BaggageItem("service"))
-		require.Equal(t, "stub", spans[2].BaggageItem("service"))
-	})
 
-	t.Run("stream", func(t *testing.T) {
-		serverSocketPath := testhelper.GetTemporaryGitalySocketFileName(t)
-
-		listener, err := net.Listen("unix", serverSocketPath)
-		require.NoError(t, err)
-
-		grpcServer := grpc.NewServer(grpc.StreamInterceptor(grpctracing.StreamServerTracingInterceptor()))
-		svc := &testSvc{
-			PingStreamMethod: func(stream proxytestdata.TestService_PingStreamServer) error {
-				span, _ := opentracing.StartSpanFromContext(stream.Context(), "health")
-				defer span.Finish()
-				span.LogKV("was", "called")
-				_, err := stream.Recv()
-				assert.NoError(t, err)
-				return stream.Send(&proxytestdata.PingResponse{})
-			},
+		for i, expectedSpan := range []struct {
+			baggage   string
+			operation string
+		}{
+			// This is the first span we expect, which is the
+			// "health" span which we've manually created inside of
+			// PingMethod.
+			{baggage: "", operation: "nested-span"},
+			// This span is the RPC call to TestService/Ping. It
+			// inherits the "unary-check" we set up and thus has
+			// baggage.
+			{baggage: "stub", operation: "/mwitkow.testproto.TestService/Ping"},
+			// And this finally is the outermost span which we
+			// manually set up before the RPC call.
+			{baggage: "stub", operation: "unary-check"},
+		} {
+			assert.IsType(t, spans[i], &jaeger.Span{})
+			span := spans[i].(*jaeger.Span)
+
+			assert.Equal(t, expectedSpan.baggage, span.BaggageItem("service"), "wrong baggage item for span %d", i)
+			assert.Equal(t, expectedSpan.operation, span.OperationName(), "wrong operation name for span %d", i)
 		}
-		proxytestdata.RegisterTestServiceServer(grpcServer, svc)
-
-		go func() { assert.NoError(t, grpcServer.Serve(listener)) }()
-		defer grpcServer.Stop()
+	})
 
+	t.Run("stream", func(t *testing.T) {
 		reporter := jaeger.NewInMemoryReporter()
-		tracer, closer := jaeger.NewTracer("", jaeger.NewConstSampler(true), reporter)
-		defer closer.Close()
+		tracer, tracerCloser := jaeger.NewTracer("", jaeger.NewConstSampler(true), reporter)
 
 		defer func(old opentracing.Tracer) { opentracing.SetGlobalTracer(old) }(opentracing.GlobalTracer())
 		opentracing.SetGlobalTracer(tracer)
 
-		span := tracer.StartSpan("stream-check")
-		span = span.SetBaggageItem("service", "stub")
-
-		ctx, cancel := testhelper.Context()
-		defer cancel()
-
+		// This needs to be run after setting up the global tracer as it will cause us to
+		// create the span when executing the RPC call further down below.
 		cc, err := DialContext(ctx, "unix://"+serverSocketPath, nil)
 		require.NoError(t, err)
 		defer cc.Close()
 
-		client := proxytestdata.NewTestServiceClient(cc)
+		// We set up a "main" span here, which is going to be what the other spans inherit
+		// from. In order to check whether baggage works correctly, we also set up a "stub"
+		// baggage item which should be inherited to child contexts.
+		span := tracer.StartSpan("stream-check")
+		span = span.SetBaggageItem("service", "stub")
+		ctx := opentracing.ContextWithSpan(ctx, span)
 
-		ctx = opentracing.ContextWithSpan(ctx, span)
-		stream, err := client.PingStream(ctx)
+		// We're now invoking the streaming RPC with the span injected into the context.
+		// This should create a span that's nested into the "stream-check" span.
+		stream, err := proxytestdata.NewTestServiceClient(cc).PingStream(ctx)
 		require.NoError(t, err)
-
 		require.NoError(t, stream.Send(&proxytestdata.PingRequest{}))
 		require.NoError(t, stream.CloseSend())
-
 		_, err = stream.Recv()
 		require.NoError(t, err)
 
 		span.Finish()
+		tracerCloser.Close()
 
 		spans := reporter.GetSpans()
 		require.Len(t, spans, 3)
-		require.Equal(t, "stub", spans[0].BaggageItem("service"))
-		require.Equal(t, "", spans[1].BaggageItem("service"))
-		require.Equal(t, "stub", spans[2].BaggageItem("service"))
+
+		for i, expectedSpan := range []struct {
+			baggage   string
+			operation string
+		}{
+			// This span is the RPC call to TestService/Ping.
+			{baggage: "stub", operation: "/mwitkow.testproto.TestService/PingStream"},
+			// This is the second span we expect, which is the "nested-span" span which
+			// we've manually created inside of PingMethod. This is different than for
+			// unary RPCs: given that one can send multiple messages to the RPC, we may
+			// see multiple such "nested-span"s being created. And the PingStream span
+			// will only be finalized last.
+			{baggage: "", operation: "nested-span"},
+			// And this finally is the outermost span which we
+			// manually set up before the RPC call.
+			{baggage: "stub", operation: "stream-check"},
+		} {
+			if !assert.IsType(t, spans[i], &jaeger.Span{}) {
+				continue
+			}
+
+			span := spans[i].(*jaeger.Span)
+			assert.Equal(t, expectedSpan.baggage, span.BaggageItem("service"), "wrong baggage item for span %d", i)
+			assert.Equal(t, expectedSpan.operation, span.OperationName(), "wrong operation name for span %d", i)
+		}
 	})
 }