chore(spanner): Add changes in Spanner executor for testing end to end tracing

spanner/client.go

@@ -63,6 +63,10 @@ const (
 
 	requestsCompressionHeader = "x-response-encoding"
 
+	// endToEndTracingHeader is the name of the metadata header if client
+	// has opted-in for the creation of trace spans on the Spanner layer.
+	endToEndTracingHeader = "x-goog-spanner-end-to-end-tracing"
+
 	// numChannels is the default value for NumChannels of client.
 	numChannels = 4
 )
@@ -331,6 +335,12 @@ type ClientConfig struct {
 	DirectedReadOptions *sppb.DirectedReadOptions
 
 	OpenTelemetryMeterProvider metric.MeterProvider
+
+	// EnableEndToEndTracing indicates whether end to end tracing is enabled or not.
+	// If it is enabled, trace spans will be created at Spanner layer.
+	//
+	// Default: false
+	EnableEndToEndTracing bool
 }
 
 type openTelemetryConfig struct {
@@ -444,6 +454,9 @@ func newClientWithConfig(ctx context.Context, database string, config ClientConf
 	if config.Compression == gzip.Name {
 		md.Append(requestsCompressionHeader, gzip.Name)
 	}
+	if config.EnableEndToEndTracing {
+		md.Append(endToEndTracingHeader, "true")
+	}
 
 	// Create a session client.
 	sc := newSessionClient(pool, database, config.UserAgent, sessionLabels, config.DatabaseRole, config.DisableRouteToLeader, md, config.BatchTimeout, config.Logger, config.CallOptions)

spanner/client_test.go

@@ -4059,6 +4059,52 @@ func TestClient_WithGRPCConnectionPoolAndNumChannels_Misconfigured(t *testing.T)
 	}
 }
 
+func TestClient_WithEndToEndTracingHeader(t *testing.T) {
+	t.Parallel()
+
+	server, opts, serverTeardown := NewMockedSpannerInMemTestServer(t)
+	defer serverTeardown()
+
+	wantEndToEndTracing := true
+	config := ClientConfig{EnableEndToEndTracing: wantEndToEndTracing}
+	client, err := makeClientWithConfig(context.Background(), "projects/p/instances/i/databases/d", config, server.ServerAddress, opts...)
+	if err != nil {
+		t.Fatalf("failed to get a client: %v", err)
+	}
+	gotEndToEndTracing := false
+	for _, val := range client.sc.md.Get(endToEndTracingHeader) {
+		if val == "true" {
+			gotEndToEndTracing = true
+		}
+	}
+	if gotEndToEndTracing != wantEndToEndTracing {
+		t.Fatalf("mismatch in client configuration for property EnableEndToEndTracing: got %v, want %v", gotEndToEndTracing, wantEndToEndTracing)
+	}
+}
+
+func TestClient_WithoutEndToEndTracingHeader(t *testing.T) {
+	t.Parallel()
+
+	server, opts, serverTeardown := NewMockedSpannerInMemTestServer(t)
+	defer serverTeardown()
+
+	wantEndToEndTracing := false
+	config := ClientConfig{EnableEndToEndTracing: wantEndToEndTracing}
+	client, err := makeClientWithConfig(context.Background(), "projects/p/instances/i/databases/d", config, server.ServerAddress, opts...)
+	if err != nil {
+		t.Fatalf("failed to get a client: %v", err)
+	}
+	gotEndToEndTracing := false
+	for _, val := range client.sc.md.Get(endToEndTracingHeader) {
+		if val == "true" {
+			gotEndToEndTracing = true
+		}
+	}
+	if gotEndToEndTracing != wantEndToEndTracing {
+		t.Fatalf("mismatch in client configuration for property EnableEndToEndTracing: got %v, want %v", gotEndToEndTracing, wantEndToEndTracing)
+	}
+}
+
 func TestClient_WithCustomBatchTimeout(t *testing.T) {
 	t.Parallel()
 

spanner/test/cloudexecutor/executor/actions/batch.go

@@ -23,6 +23,7 @@ import (
 	"cloud.google.com/go/spanner/executor/apiv1/executorpb"
 	"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/outputstream"
 	"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility"
+	trace "cloud.google.com/go/trace/apiv1"
 	"google.golang.org/api/option"
 	spb "google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc/codes"
@@ -31,10 +32,11 @@ import (
 
 // StartBatchTxnHandler holds the necessary components and options required for start batch transaction action.
 type StartBatchTxnHandler struct {
-	Action        *executorpb.StartBatchTransactionAction
-	FlowContext   *ExecutionFlowContext
-	OutcomeSender *outputstream.OutcomeSender
-	Options       []option.ClientOption
+	Action             *executorpb.StartBatchTransactionAction
+	FlowContext        *ExecutionFlowContext
+	OutcomeSender      *outputstream.OutcomeSender
+	Options            []option.ClientOption
+	TraceClientOptions []option.ClientOption
 }
 
 // ExecuteAction that starts a batch transaction
@@ -50,10 +52,16 @@ func (h *StartBatchTxnHandler) ExecuteAction(ctx context.Context) error {
 		return h.OutcomeSender.FinishWithError(spanner.ToSpannerError(status.Error(codes.InvalidArgument, "database path must be set for this action")))
 	}
 
-	client, err := spanner.NewClient(ctx, h.FlowContext.Database, h.Options...)
+	client, err := spanner.NewClientWithConfig(ctx, h.FlowContext.Database, spanner.ClientConfig{SessionPoolConfig: spanner.DefaultSessionPoolConfig, DisableRouteToLeader: false, EnableEndToEndTracing: true}, h.Options...)
+	if err != nil {
+		return h.OutcomeSender.FinishWithError(err)
+	}
+	// Create a trace client to read the traces from Cloud Trace.
+	traceClient, err := trace.NewClient(ctx, h.TraceClientOptions...)
 	if err != nil {
 		return h.OutcomeSender.FinishWithError(err)
 	}
+	h.FlowContext.TraceClient = traceClient
 	var txn *spanner.BatchReadOnlyTransaction
 	if h.Action.GetBatchTxnTime() != nil {
 		timestamp := time.Unix(h.Action.GetBatchTxnTime().Seconds, int64(h.Action.GetBatchTxnTime().Nanos))

spanner/test/cloudexecutor/executor/actions/execution_flow_context.go

@@ -24,6 +24,7 @@ import (
 	"cloud.google.com/go/spanner"
 	"cloud.google.com/go/spanner/executor/apiv1/executorpb"
 	"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility"
+	trace "cloud.google.com/go/trace/apiv1"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -49,6 +50,7 @@ type ExecutionFlowContext struct {
 	roTxn                    *spanner.ReadOnlyTransaction           // Current read-only transaction
 	batchTxn                 *spanner.BatchReadOnlyTransaction      // Current batch read-only transaction
 	DbClient                 *spanner.Client                        // Current database client
+	TraceClient              *trace.Client                          // Current trace client
 	tableMetadata            *utility.TableMetadataHelper           // If in a txn (except batch), this has metadata info about table columns
 	numPendingReads          int64                                  // Number of pending read/query actions.
 	readAborted              bool                                   // Indicate whether there's a read/query action got aborted and the transaction need to be reset.

spanner/test/cloudexecutor/executor/actions/transaction.go

@@ -24,6 +24,7 @@ import (
 	"cloud.google.com/go/spanner/executor/apiv1/executorpb"
 	"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/outputstream"
 	"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility"
+	trace "cloud.google.com/go/trace/apiv1"
 	"google.golang.org/api/option"
 	spb "google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc/codes"
@@ -33,10 +34,11 @@ import (
 
 // StartTxnHandler holds the necessary components and options required for start transaction action.
 type StartTxnHandler struct {
-	Action        *executorpb.StartTransactionAction
-	FlowContext   *ExecutionFlowContext
-	OutcomeSender *outputstream.OutcomeSender
-	Options       []option.ClientOption
+	Action             *executorpb.StartTransactionAction
+	FlowContext        *ExecutionFlowContext
+	OutcomeSender      *outputstream.OutcomeSender
+	Options            []option.ClientOption
+	TraceClientOptions []option.ClientOption
 }
 
 // ExecuteAction that starts a read-write or read-only transaction.
@@ -54,11 +56,17 @@ func (h *StartTxnHandler) ExecuteAction(ctx context.Context) error {
 	h.FlowContext.tableMetadata = metadata
 
 	// TODO(harsha) where do I close the client? defer client.Close()
-	client, err := spanner.NewClient(ctx, h.FlowContext.Database, h.Options...)
+	client, err := spanner.NewClientWithConfig(ctx, h.FlowContext.Database, spanner.ClientConfig{SessionPoolConfig: spanner.DefaultSessionPoolConfig, DisableRouteToLeader: false, EnableEndToEndTracing: true}, h.Options...)
 	if err != nil {
 		return h.OutcomeSender.FinishWithError(err)
 	}
 	h.FlowContext.DbClient = client
+	// Create a trace client to read the traces.
+	traceClient, err := trace.NewClient(ctx, h.TraceClientOptions...)
+	if err != nil {
+		return h.OutcomeSender.FinishWithError(err)
+	}
+	h.FlowContext.TraceClient = traceClient
 	if h.FlowContext.isTransactionActiveLocked() {
 		return h.OutcomeSender.FinishWithError(spanner.ToSpannerError(status.Error(codes.InvalidArgument, "already in a transaction")))
 	}

spanner/test/cloudexecutor/executor/executor_proxy_server_impl.go

@@ -25,24 +25,36 @@ import (
 	"google.golang.org/api/option"
 )
 
+const MAX_CLOUD_TRACE_CHECK_LIMIT = 10
+
 // CloudProxyServer holds the cloud executor server.
 type CloudProxyServer struct {
-	serverContext context.Context
-	options       []option.ClientOption
+	// members below should be set by the caller
+	serverContext      context.Context
+	options            []option.ClientOption
+	traceClientOptions []option.ClientOption
+	// members below represent internal state
+	cloudTraceCheckCount int
 }
 
 // NewCloudProxyServer initializes and returns a new CloudProxyServer instance.
-func NewCloudProxyServer(ctx context.Context, opts []option.ClientOption) (*CloudProxyServer, error) {
-	return &CloudProxyServer{serverContext: ctx, options: opts}, nil
+func NewCloudProxyServer(ctx context.Context, opts []option.ClientOption, traceClientOpts []option.ClientOption) (*CloudProxyServer, error) {
+	return &CloudProxyServer{serverContext: ctx, options: opts, traceClientOptions: traceClientOpts, cloudTraceCheckCount: 0}, nil
 }
 
 // ExecuteActionAsync is implementation of ExecuteActionAsync in SpannerExecutorProxyServer. It's a
 // streaming method in which client and server exchange SpannerActions and SpannerActionOutcomes.
 func (s *CloudProxyServer) ExecuteActionAsync(inputStream executorpb.SpannerExecutorProxy_ExecuteActionAsyncServer) error {
 	handler := &inputstream.CloudStreamHandler{
-		Stream:        inputStream,
-		ServerContext: s.serverContext,
-		Options:       s.options,
+		Stream:                 inputStream,
+		ServerContext:          s.serverContext,
+		Options:                s.options,
+		TraceClientOptions:     s.traceClientOptions,
+		CloudTraceCheckAllowed: (s.cloudTraceCheckCount < MAX_CLOUD_TRACE_CHECK_LIMIT),
+	}
+	if err := handler.Execute(); err != nil {
+		return err
 	}
-	return handler.Execute()
+	s.cloudTraceCheckCount += handler.GetCompletedCloudTraceCheckCount()
+	return nil
 }