sql,server: add transaction diagnostic polling logic

Adds logic to the TxnRegistry to poll for new transaction
diagnostic requests. This is necessary to ensure that all
nodes in a cluster are watching for transactions to capture
diagnostic bundles for. The diagnostic polling is controlled
by a new cluster setting, `sql.txn_diagnostics.poll_interval`,
which can be set to 0 to disable polling. Note that this
wont disable transaction diagnostics request handling all together,
the gateway node handling a transaction diagnostics request
will still have the request in its local registry.

The poll requests implementation is the same as the implementation
in statement_diagnostics.go for the statement diagnostics regsitry

Resolves: CRDB-54320
Epic: CRDB-53541
Release note: None

Author: kyle-a-wong

pkg/server/server_sql.go

@@ -172,6 +172,7 @@ type SQLServer struct {
 	statsRefresher                 *stats.Refresher
 	temporaryObjectCleaner         *sql.TemporaryObjectCleaner
 	stmtDiagnosticsRegistry        *stmtdiagnostics.Registry
+	txnDiagnosticsRegistry         *stmtdiagnostics.TxnRegistry
 	sqlLivenessSessionID           sqlliveness.SessionID
 	sqlLivenessProvider            sqlliveness.Provider
 	sqlInstanceReader              *instancestorage.Reader
@@ -1257,6 +1258,10 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 	)
 	execCfg.StmtDiagnosticsRecorder = stmtDiagnosticsRegistry
 
+	txnDiagnosticsRegistry := stmtdiagnostics.NewTxnRegistry(cfg.internalDB,
+		cfg.Settings, stmtDiagnosticsRegistry, timeutil.DefaultTimeSource{})
+	execCfg.TxnDiagnosticsRecorder = txnDiagnosticsRegistry
+
 	var upgradeMgr *upgrademanager.Manager
 	{
 		var c upgrade.Cluster
@@ -1434,6 +1439,7 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		statsRefresher:                 statsRefresher,
 		temporaryObjectCleaner:         temporaryObjectCleaner,
 		stmtDiagnosticsRegistry:        stmtDiagnosticsRegistry,
+		txnDiagnosticsRegistry:         txnDiagnosticsRegistry,
 		sqlLivenessProvider:            cfg.sqlLivenessProvider,
 		sqlInstanceStorage:             cfg.sqlInstanceStorage,
 		sqlInstanceReader:              cfg.sqlInstanceReader,
@@ -1764,7 +1770,9 @@ func (s *SQLServer) preStart(
 	if err := s.statsRefresher.Start(ctx, stopper, stats.DefaultRefreshInterval); err != nil {
 		return err
 	}
-	s.stmtDiagnosticsRegistry.Start(ctx, stopper)
+
+	stmtdiagnostics.StartPolling(ctx, s.txnDiagnosticsRegistry, s.stmtDiagnosticsRegistry, stopper)
+
 	if err := s.execCfg.TableStatsCache.Start(ctx, s.execCfg.Codec, s.execCfg.RangeFeedFactory); err != nil {
 		return err
 	}

pkg/sql/exec_util.go

@@ -1671,6 +1671,9 @@ type ExecutorConfig struct {
 	// StmtDiagnosticsRecorder deals with recording statement diagnostics.
 	StmtDiagnosticsRecorder *stmtdiagnostics.Registry
 
+	// TxnDiagnosticsRecorder deals with recording transaction diagnostics.
+	TxnDiagnosticsRecorder *stmtdiagnostics.TxnRegistry
+
 	ExternalIODirConfig base.ExternalIODirConfig
 
 	GCJobNotifier *gcjobnotifier.Notifier

pkg/sql/stmtdiagnostics/BUILD.bazel

@@ -66,6 +66,7 @@ go_test(
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_lib_pq//:pq",
+        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/sql/stmtdiagnostics/statement_diagnostics.go

@@ -28,9 +28,10 @@ import (
 
 var pollingInterval = settings.RegisterDurationSetting(
 	settings.SystemVisible,
-	"sql.stmt_diagnostics.poll_interval",
-	"rate at which the stmtdiagnostics.Registry polls for requests, set to zero to disable",
+	"sql.diagnostics.poll_interval",
+	"rate at which the stmtdiagnostics registries polls for requests, set to zero to disable",
 	10*time.Second,
+	settings.WithRetiredName("sql.stmt_diagnostics.poll_interval"),
 )
 
 var bundleChunkSize = settings.RegisterByteSizeSetting(
@@ -148,68 +149,6 @@ func NewRegistry(db isql.DB, st *cluster.Settings) *Registry {
 	return r
 }
 
-// Start will start the polling loop for the Registry.
-func (r *Registry) Start(ctx context.Context, stopper *stop.Stopper) {
-	// The registry has the same lifetime as the server, so the cancellation
-	// function can be ignored and it'll be called by the stopper.
-	ctx, _ = stopper.WithCancelOnQuiesce(ctx) // nolint:quiesce
-
-	// Since background statement diagnostics collection is not under user
-	// control, exclude it from cost accounting and control.
-	ctx = multitenant.WithTenantCostControlExemption(ctx)
-
-	// NB: The only error that should occur here would be if the server were
-	// shutting down so let's swallow it.
-	_ = stopper.RunAsyncTask(ctx, "stmt-diag-poll", r.poll)
-}
-
-func (r *Registry) poll(ctx context.Context) {
-	var (
-		timer               timeutil.Timer
-		lastPoll            time.Time
-		deadline            time.Time
-		pollIntervalChanged = make(chan struct{}, 1)
-		maybeResetTimer     = func() {
-			if interval := pollingInterval.Get(&r.st.SV); interval == 0 {
-				// Setting the interval to zero stops the polling.
-				timer.Stop()
-			} else {
-				newDeadline := lastPoll.Add(interval)
-				if deadline.IsZero() || !deadline.Equal(newDeadline) {
-					deadline = newDeadline
-					timer.Reset(timeutil.Until(deadline))
-				}
-			}
-		}
-		poll = func() {
-			if err := r.pollRequests(ctx); err != nil {
-				if ctx.Err() != nil {
-					return
-				}
-				log.Dev.Warningf(ctx, "error polling for statement diagnostics requests: %s", err)
-			}
-			lastPoll = timeutil.Now()
-		}
-	)
-	pollingInterval.SetOnChange(&r.st.SV, func(ctx context.Context) {
-		select {
-		case pollIntervalChanged <- struct{}{}:
-		default:
-		}
-	})
-	for {
-		maybeResetTimer()
-		select {
-		case <-pollIntervalChanged:
-			continue // go back around and maybe reset the timer
-		case <-timer.C:
-		case <-ctx.Done():
-			return
-		}
-		poll()
-	}
-}
-
 type StmtDiagnostic struct {
 	requestID       RequestID
 	req             Request
@@ -733,9 +672,9 @@ func (r *Registry) innerInsertStatementDiagnostics(
 	return diagID, nil
 }
 
-// pollRequests reads the pending rows from system.statement_diagnostics_requests and
+// pollStmtRequests reads the pending rows from system.statement_diagnostics_requests and
 // updates r.mu.requests accordingly.
-func (r *Registry) pollRequests(ctx context.Context) error {
+func (r *Registry) pollStmtRequests(ctx context.Context) error {
 	var rows []tree.Datums
 
 	// Loop until we run the query without straddling an epoch increment.
@@ -823,3 +762,71 @@ func (r *Registry) pollRequests(ctx context.Context) error {
 	}
 	return nil
 }
+
+// StartPolling starts a background task that polls for new statement and
+// transaction requests and updates the corresponding registries.
+func StartPolling(ctx context.Context, tr *TxnRegistry, sr *Registry, stopper *stop.Stopper) {
+	// The registry has the same lifetime as the server, so the cancellation
+	// function can be ignored and it'll be called by the stopper.
+	ctx, _ = stopper.WithCancelOnQuiesce(ctx) // nolint:quiesce
+
+	// Since background diagnostics collection is not under user
+	// control, exclude it from cost accounting and control.
+	ctx = multitenant.WithTenantCostControlExemption(ctx)
+
+	// NB: The only error that should occur here would be if the server were
+	// shutting down so let's swallow it.
+	_ = stopper.RunAsyncTask(ctx, "stmt-txn-diag-poll", func(ctx context.Context) {
+		var (
+			timer               timeutil.Timer
+			lastPoll            time.Time
+			deadline            time.Time
+			pollIntervalChanged = make(chan struct{}, 1)
+			maybeResetTimer     = func() {
+				if interval := pollingInterval.Get(&sr.st.SV); interval == 0 {
+					// Setting the interval to zero stops the polling.
+					timer.Stop()
+				} else {
+					newDeadline := lastPoll.Add(interval)
+					if deadline.IsZero() || !deadline.Equal(newDeadline) {
+						deadline = newDeadline
+						timer.Reset(timeutil.Until(deadline))
+					}
+				}
+			}
+			poll = func() {
+				if err := tr.pollTxnRequests(ctx); err != nil {
+					if ctx.Err() != nil {
+						return
+					}
+					log.Ops.Warningf(ctx, "error polling for transaction diagnostics requests: %s", err)
+				}
+				if err := sr.pollStmtRequests(ctx); err != nil {
+					if ctx.Err() != nil {
+						return
+					}
+					log.Ops.Warningf(ctx, "error polling for statement diagnostics requests: %s", err)
+				}
+				lastPoll = timeutil.Now()
+			}
+		)
+
+		pollingInterval.SetOnChange(&sr.st.SV, func(ctx context.Context) {
+			select {
+			case pollIntervalChanged <- struct{}{}:
+			default:
+			}
+		})
+		for {
+			maybeResetTimer()
+			select {
+			case <-pollIntervalChanged:
+				continue // go back around and maybe reset the timer
+			case <-timer.C:
+			case <-ctx.Done():
+				return
+			}
+			poll()
+		}
+	})
+}

pkg/sql/stmtdiagnostics/statement_diagnostics_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -692,3 +693,81 @@ func TestChangePollInterval(t *testing.T) {
 	stmtdiagnostics.PollingInterval.Override(ctx, &settings.SV, 200*time.Microsecond)
 	waitForScans(10) // ensure several scans occur
 }
+
+func TestTxnRegistry_InsertTxnRequest_Polling(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	settings := cluster.MakeTestingClusterSettings()
+	// Set to 1s so that we can quickly pick up the request.
+	stmtdiagnostics.PollingInterval.Override(ctx, &settings.SV, time.Second)
+
+	tc := serverutils.StartCluster(t, 3, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			Settings: settings,
+		},
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	s0 := tc.Server(0).ApplicationLayer()
+	registry := s0.ExecutorConfig().(sql.ExecutorConfig).TxnDiagnosticsRecorder
+	registry2 := tc.Server(1).ApplicationLayer().ExecutorConfig().(sql.ExecutorConfig).TxnDiagnosticsRecorder
+	registry3 := tc.Server(2).ApplicationLayer().ExecutorConfig().(sql.ExecutorConfig).TxnDiagnosticsRecorder
+
+	id, err := registry.InsertTxnRequestInternal(
+		ctx,
+		1111,
+		[]uint64{1111, 2222, 3333},
+		"testuser",
+		0.5,
+		time.Millisecond*100,
+		0,
+		false,
+	)
+	require.NoError(t, err)
+	require.NotEqual(t, stmtdiagnostics.RequestID(0), id)
+
+	var expectedRequest, req stmtdiagnostics.TxnRequest
+	var ok bool
+	expectedRequest, ok = registry.GetRequest(id)
+	require.True(t, ok)
+	testutils.SucceedsSoon(t, func() error {
+
+		req, ok = registry2.GetRequest(id)
+		if !ok {
+			return errors.New("request not found on server 2")
+		}
+		if !assert.Equal(t, expectedRequest, req) {
+			return errors.Newf("request on server2 doesnt match expected request. Expected: %+v, got: %+v", expectedRequest, req)
+		}
+
+		req, ok = registry3.GetRequest(id)
+		if !ok {
+			return errors.New("request not found on server 3")
+		}
+		require.Equal(t, expectedRequest, req)
+		if !assert.Equal(t, expectedRequest, req) {
+			return errors.Newf("request on server3 doesnt match expected request. Expected: %+v, got: %+v", expectedRequest, req)
+		}
+
+		return nil
+	})
+
+	// mark the request as complete and ensure that it is removed from all 3 nodes
+	runner := sqlutils.MakeSQLRunner(tc.ServerConn(0))
+	runner.Exec(t, "UPDATE system.transaction_diagnostics_requests "+
+		"SET completed = true, transaction_diagnostics_id = 12345 WHERE id = $1", id)
+	testutils.SucceedsSoon(t, func() error {
+		if _, ok = registry.GetRequest(id); ok {
+			return errors.New("request still found on server 1")
+		}
+		if _, ok = registry2.GetRequest(id); ok {
+			return errors.New("request still found on server 2")
+		}
+		if _, ok = registry3.GetRequest(id); ok {
+			return errors.New("request still found on server 3")
+		}
+		return nil
+	})
+}