tracing,tracingservice: adds a trace service to pull clusterwide trace spans

pkg/BUILD.bazel

@@ -358,6 +358,8 @@ ALL_TESTS = [
     "//pkg/util/timetz:timetz_test",
     "//pkg/util/timeutil/pgdate:pgdate_test",
     "//pkg/util/timeutil:timeutil_test",
+    "//pkg/util/tracing/collector:collector_test",
+    "//pkg/util/tracing/service:service_test",
     "//pkg/util/tracing:tracing_test",
     "//pkg/util/treeprinter:treeprinter_test",
     "//pkg/util/uint128:uint128_test",

pkg/kv/kvserver/client_raft_test.go

@@ -4321,7 +4321,7 @@ func TestStoreRangeWaitForApplication(t *testing.T) {
 
 	var targets []target
 	for _, s := range tc.Servers {
-		conn, err := s.NodeDialer().Dial(ctx, s.NodeID(), rpc.DefaultClass)
+		conn, err := s.NodeDialer().(*nodedialer.Dialer).Dial(ctx, s.NodeID(), rpc.DefaultClass)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -4448,7 +4448,7 @@ func TestStoreWaitForReplicaInit(t *testing.T) {
 	defer tc.Stopper().Stop(ctx)
 	store := tc.GetFirstStoreFromServer(t, 0)
 
-	conn, err := tc.Servers[0].NodeDialer().Dial(ctx, store.Ident.NodeID, rpc.DefaultClass)
+	conn, err := tc.Servers[0].NodeDialer().(*nodedialer.Dialer).Dial(ctx, store.Ident.NodeID, rpc.DefaultClass)
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/server/BUILD.bazel

@@ -173,7 +173,9 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
+        "//pkg/util/tracing/service",
         "//pkg/util/tracing/tracingpb",
+        "//pkg/util/tracing/tracingservicepb:tracingservicepb_go_proto",
         "//pkg/util/uuid",
         "@com_github_cenkalti_backoff//:backoff",
         "@com_github_cockroachdb_apd_v2//:apd",

pkg/server/server_sql.go

@@ -82,6 +82,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/service"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingservicepb"
 	"github.com/cockroachdb/errors"
 	"github.com/marusama/semaphore"
 	"google.golang.org/grpc"
@@ -101,6 +103,7 @@ type SQLServer struct {
 	internalExecutor *sql.InternalExecutor
 	leaseMgr         *lease.Manager
 	blobService      *blobs.Service
+	tracingService   *service.Service
 	tenantConnect    kvtenant.Connector
 	// sessionRegistry can be queried for info on running SQL sessions. It is
 	// shared between the sql.Server and the statusServer.
@@ -261,6 +264,10 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 	}
 	blobspb.RegisterBlobServer(cfg.grpcServer, blobService)
 
+	// Create trace service for inter-node sharing of inflight trace spans.
+	tracingService := service.New(cfg.Settings.Tracer)
+	tracingservicepb.RegisterTracingServer(cfg.grpcServer, tracingService)
+
 	jobRegistry := cfg.circularJobRegistry
 	{
 		regLiveness := cfg.nodeLiveness
@@ -732,6 +739,7 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		internalExecutor:        cfg.circularInternalExecutor,
 		leaseMgr:                leaseMgr,
 		blobService:             blobService,
+		tracingService:          tracingService,
 		tenantConnect:           cfg.tenantConnect,
 		sessionRegistry:         cfg.sessionRegistry,
 		jobRegistry:             jobRegistry,

pkg/server/testserver.go

@@ -376,6 +376,14 @@ func (ts *TestServer) NodeLiveness() interface{} {
 	return nil
 }
 
+// NodeDialer returns the NodeDialer used by the TestServer.
+func (ts *TestServer) NodeDialer() interface{} {
+	if ts != nil {
+		return ts.nodeDialer
+	}
+	return nil
+}
+
 // HeartbeatNodeLiveness heartbeats the server's NodeLiveness record.
 func (ts *TestServer) HeartbeatNodeLiveness() error {
 	if ts == nil {
@@ -437,14 +445,6 @@ func (ts *TestServer) RaftTransport() *kvserver.RaftTransport {
 	return nil
 }
 
-// NodeDialer returns the NodeDialer used by the TestServer.
-func (ts *TestServer) NodeDialer() *nodedialer.Dialer {
-	if ts != nil {
-		return ts.nodeDialer
-	}
-	return nil
-}
-
 // Start starts the TestServer by bootstrapping an in-memory store
 // (defaults to maximum of 100M). The server is started, launching the
 // node RPC server and all HTTP endpoints. Use the value of

pkg/testutils/serverutils/test_server_shim.go

@@ -136,6 +136,10 @@ type TestServerInterface interface {
 	// HeartbeatNodeLiveness heartbeats the server's NodeLiveness record.
 	HeartbeatNodeLiveness() error
 
+	// NodeDialer exposes the NodeDialer instance used by the TestServer as an
+	// interface{}.
+	NodeDialer() interface{}
+
 	// SetDistSQLSpanResolver changes the SpanResolver used for DistSQL inside the
 	// server's executor. The argument must be a physicalplan.SpanResolver
 	// instance.

pkg/testutils/testcluster/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//pkg/kv/kvserver/liveness/livenesspb",
         "//pkg/roachpb",
         "//pkg/rpc",
+        "//pkg/rpc/nodedialer",
         "//pkg/server",
         "//pkg/server/serverpb",
         "//pkg/storage",

pkg/testutils/testcluster/testcluster.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
@@ -1439,7 +1440,7 @@ func (tc *TestCluster) RestartServerWithInspect(idx int, inspect func(s *server.
 						}
 						for i := 0; i < rpc.NumConnectionClasses; i++ {
 							class := rpc.ConnectionClass(i)
-							if _, err := s.NodeDialer().Dial(ctx, srv.NodeID(), class); err != nil {
+							if _, err := s.NodeDialer().(*nodedialer.Dialer).Dial(ctx, srv.NodeID(), class); err != nil {
 								return errors.Wrapf(err, "connecting n%d->n%d (class %v)", s.NodeID(), srv.NodeID(), class)
 							}
 						}

pkg/util/tracing/collector/BUILD.bazel

@@ -0,0 +1,50 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "collector",
+    srcs = [
+        "collector.go",
+        "nodes.go",
+    ],
+    importpath = "github.com/cockroachdb/cockroach/pkg/util/tracing/collector",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/kv/kvserver/liveness/livenesspb",
+        "//pkg/roachpb",
+        "//pkg/rpc",
+        "//pkg/rpc/nodedialer",
+        "//pkg/util/ctxgroup",
+        "//pkg/util/log",
+        "//pkg/util/quotapool",
+        "//pkg/util/syncutil",
+        "//pkg/util/tracing",
+        "//pkg/util/tracing/tracingpb",
+        "//pkg/util/tracing/tracingservicepb:tracingservicepb_go_proto",
+        "@com_github_cockroachdb_redact//:redact",
+    ],
+)
+
+go_test(
+    name = "collector_test",
+    srcs = [
+        "collector_test.go",
+        "main_test.go",
+    ],
+    deps = [
+        ":collector",
+        "//pkg/base",
+        "//pkg/ccl/utilccl",
+        "//pkg/kv/kvserver/liveness",
+        "//pkg/rpc/nodedialer",
+        "//pkg/security",
+        "//pkg/security/securitytest",
+        "//pkg/server",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/testcluster",
+        "//pkg/util/leaktest",
+        "//pkg/util/randutil",
+        "//pkg/util/tracing",
+        "@com_github_gogo_protobuf//types",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/util/tracing/collector/collector.go

@@ -0,0 +1,111 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package collector
+
+import (
+	"context"
+	"sort"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
+	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingservicepb"
+)
+
+// NodeLiveness is the subset of the interface satisfied by CRDB's node liveness
+// component that the tracing service relies upon.
+type NodeLiveness interface {
+	GetLivenessesFromKV(context.Context) ([]livenesspb.Liveness, error)
+	IsLive(roachpb.NodeID) (bool, error)
+}
+
+// TraceCollector can be used to extract recordings from inflight spans for a
+// given traceID, from all nodes of the cluster.
+type TraceCollector struct {
+	tracer       *tracing.Tracer
+	dialer       *nodedialer.Dialer
+	nodeliveness NodeLiveness
+}
+
+// New returns a TraceCollector.
+func New(
+	dialer *nodedialer.Dialer, nodeliveness NodeLiveness, tracer *tracing.Tracer,
+) *TraceCollector {
+	return &TraceCollector{
+		dialer:       dialer,
+		nodeliveness: nodeliveness,
+		tracer:       tracer,
+	}
+}
+
+// GetSpanRecordingsFromCluster returns the inflight span recordings from all
+// nodes in the cluster.
+// This method does not distinguish between requests for local and remote
+// inflight spans, and relies on gRPC short circuiting local requests.
+func (t *TraceCollector) GetSpanRecordingsFromCluster(
+	ctx context.Context, traceID uint64,
+) ([]tracingpb.RecordedSpan, error) {
+	var res []tracingpb.RecordedSpan
+	ns, err := nodesFromNodeLiveness(ctx, t.nodeliveness)
+	if err != nil {
+		return res, err
+	}
+
+	// Collect spans from all clients.
+	// We'll want to rate limit outgoing RPCs (limit pulled out of thin air).
+	var mu syncutil.Mutex
+	qp := quotapool.NewIntPool("every-node", 25)
+	log.Infof(ctx, "executing GetSpanRecordings on nodes %s", ns)
+	grp := ctxgroup.WithContext(ctx)
+
+	for _, node := range ns {
+		id := node.id // copy out of the loop variable
+		alloc, err := qp.Acquire(ctx, 1)
+		if err != nil {
+			return res, err
+		}
+
+		grp.GoCtx(func(ctx context.Context) error {
+			defer alloc.Release()
+
+			conn, err := t.dialer.Dial(ctx, id, rpc.DefaultClass)
+			if err != nil {
+				return err
+			}
+			traceClient := tracingservicepb.NewTracingClient(conn)
+			resp, err := traceClient.GetSpanRecordings(ctx,
+				&tracingservicepb.SpanRecordingRequest{TraceID: traceID})
+			if err != nil {
+				return err
+			}
+			mu.Lock()
+			res = append(res, resp.SpanRecordings...)
+			mu.Unlock()
+			return nil
+		})
+	}
+	if err := grp.Wait(); err != nil {
+		return res, err
+	}
+
+	sort.SliceStable(res, func(i, j int) bool {
+		return res[i].StartTime.Before(res[j].StartTime)
+	})
+
+	return res, nil
+}