UT for concurrent inject/extract metadata

Signed-off-by: Jin Dong <[email protected]>

go.mod

@@ -9,6 +9,7 @@ require (
 	go.opentelemetry.io/otel/metric v1.19.0
 	go.opentelemetry.io/otel/sdk v1.19.0
 	go.opentelemetry.io/otel/trace v1.19.0
+	golang.org/x/sync v0.10.0
 	google.golang.org/grpc v1.57.0
 	google.golang.org/protobuf v1.30.0
 )

go.sum

@@ -1051,6 +1051,8 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ=
+golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

interceptor_test.go

@@ -42,9 +42,12 @@ import (
 	"github.com/containerd/otelttrpc/internal"
 	"github.com/containerd/ttrpc"
 	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"go.opentelemetry.io/otel/trace"
+	"golang.org/x/sync/errgroup"
 )
 
 const serviceName = "testService"
@@ -88,9 +91,64 @@ func (s *testingServer) Test(ctx context.Context, req *internal.TestPayload) (*i
 	return tp, nil
 }
 
+func TestClientCallServerConcurrent(t *testing.T) {
+	var (
+		ctx            = ttrpc.WithMetadata(context.Background(), ttrpc.MD{"test-key": []string{"test-val"}})
+		exp, tp        = newTracerProvider()
+		server         = mustServer(t)(newServerWithTTRPCInterceptor(tp))
+		testImpl       = &testingServer{}
+		addr, listener = newTestListener(t)
+		payload        = &internal.TestPayload{
+			Foo: "bar",
+		}
+	)
+
+	concurrency := 30
+	testClients := make([]*testingClient, 0, concurrency)
+	for i := 0; i < concurrency; i++ {
+		client, cleanup := newTestClient(t, addr, tp)
+		testClients = append(testClients, newTestingClient(client))
+		defer cleanup()
+	}
+	defer listener.Close()
+	defer func() { _ = tp.Shutdown(ctx) }()
+
+	registerTestingService(server, testImpl)
+
+	go func() {
+		_ = server.Serve(ctx, listener)
+	}()
+	defer func() {
+		_ = server.Shutdown(ctx)
+	}()
+
+	var eg errgroup.Group
+	for _, testClient := range testClients {
+		// capture range variable
+		// TODO: we can remove this once we upgrade golang to >= 1.22
+		testClient := testClient
+		eg.Go(func() error {
+			if _, err := testClient.Test(ctx, payload); err != nil {
+				return err
+			}
+			return nil
+		})
+	}
+
+	if err := eg.Wait(); err != nil {
+		t.Fatal(err)
+	}
+
+	// get exported spans
+	snapshots := exp.GetSpans().Snapshots()
+	// we should capture `concurrency * 2` spans, one each from client and server side
+	// TODO: validate individual spans and their attributes
+	assert.Equal(t, concurrency*2, len(snapshots), "Number of spans mismatched")
+}
+
 func TestClientCallServer(t *testing.T) {
 	var (
-		ctx             = context.Background()
+		ctx             = ttrpc.WithMetadata(context.Background(), ttrpc.MD{"test-key": []string{"test-val"}})
 		exp, tp         = newTracerProvider()
 		server          = mustServer(t)(newServerWithTTRPCInterceptor(tp))
 		testImpl        = &testingServer{}
@@ -153,6 +211,8 @@ func newTracerProvider() (*tracetest.InMemoryExporter, *sdktrace.TracerProvider)
 	tp := sdktrace.NewTracerProvider(
 		sdktrace.WithSyncer(exp),
 	)
+
+	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}))
 	return exp, tp
 }