github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package stmtdiagnostics_test
    12  
    13  import (
    14  	"context"
    15  	gosql "database/sql"
    16  	"fmt"
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/base"
    21  	"github.com/cockroachdb/cockroach/pkg/keys"
    22  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
    23  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    24  	"github.com/cockroachdb/cockroach/pkg/sql"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils"
    26  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    27  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    28  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    29  	"github.com/cockroachdb/cockroach/pkg/util/uuid"
    30  	"github.com/cockroachdb/errors"
    31  	"github.com/stretchr/testify/require"
    32  )
    33  
    34  func TestDiagnosticsRequest(t *testing.T) {
    35  	defer leaktest.AfterTest(t)()
    36  	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{})
    37  	ctx := context.Background()
    38  	defer s.Stopper().Stop(ctx)
    39  	_, err := db.Exec("CREATE TABLE test (x int PRIMARY KEY)")
    40  	require.NoError(t, err)
    41  
    42  	// Ask to trace a particular query.
    43  	registry := s.ExecutorConfig().(sql.ExecutorConfig).StmtDiagnosticsRecorder
    44  	reqID, err := registry.InsertRequestInternal(ctx, "INSERT INTO test VALUES (_)")
    45  	require.NoError(t, err)
    46  	reqRow := db.QueryRow(
    47  		"SELECT completed, statement_diagnostics_id FROM system.statement_diagnostics_requests WHERE ID = $1", reqID)
    48  	var completed bool
    49  	var traceID gosql.NullInt64
    50  	require.NoError(t, reqRow.Scan(&completed, &traceID))
    51  	require.False(t, completed)
    52  	require.False(t, traceID.Valid) // traceID should be NULL
    53  
    54  	// Run the query.
    55  	_, err = db.Exec("INSERT INTO test VALUES (1)")
    56  	require.NoError(t, err)
    57  
    58  	// Check that the row from statement_diagnostics_request was marked as completed.
    59  	checkCompleted := func(reqID int64) {
    60  		traceRow := db.QueryRow(
    61  			"SELECT completed, statement_diagnostics_id FROM system.statement_diagnostics_requests WHERE ID = $1", reqID)
    62  		require.NoError(t, traceRow.Scan(&completed, &traceID))
    63  		require.True(t, completed)
    64  		require.True(t, traceID.Valid)
    65  	}
    66  
    67  	checkCompleted(reqID)
    68  
    69  	// Check the trace.
    70  	row := db.QueryRow("SELECT jsonb_pretty(trace) FROM system.statement_diagnostics WHERE ID = $1", traceID.Int64)
    71  	var json string
    72  	require.NoError(t, row.Scan(&json))
    73  	require.Contains(t, json, "traced statement")
    74  	require.Contains(t, json, "statement execution committed the txn")
    75  
    76  	// Verify that we can handle multiple requests at the same time.
    77  	id1, err := registry.InsertRequestInternal(ctx, "INSERT INTO test VALUES (_)")
    78  	require.NoError(t, err)
    79  	id2, err := registry.InsertRequestInternal(ctx, "SELECT x FROM test")
    80  	require.NoError(t, err)
    81  	id3, err := registry.InsertRequestInternal(ctx, "SELECT x FROM test WHERE x > _")
    82  	require.NoError(t, err)
    83  
    84  	// Run the queries in a different order.
    85  	_, err = db.Exec("SELECT x FROM test")
    86  	require.NoError(t, err)
    87  	checkCompleted(id2)
    88  
    89  	_, err = db.Exec("SELECT x FROM test WHERE x > 1")
    90  	require.NoError(t, err)
    91  	checkCompleted(id3)
    92  
    93  	_, err = db.Exec("INSERT INTO test VALUES (2)")
    94  	require.NoError(t, err)
    95  	checkCompleted(id1)
    96  }
    97  
    98  // Test that a different node can service a diagnostics request.
    99  func TestDiagnosticsRequestDifferentNode(t *testing.T) {
   100  	defer leaktest.AfterTest(t)()
   101  	tc := serverutils.StartTestCluster(t, 2, base.TestClusterArgs{})
   102  	ctx := context.Background()
   103  	defer tc.Stopper().Stop(ctx)
   104  	db0 := tc.ServerConn(0)
   105  	db1 := tc.ServerConn(1)
   106  	_, err := db0.Exec("CREATE TABLE test (x int PRIMARY KEY)")
   107  	require.NoError(t, err)
   108  
   109  	// Ask to trace a particular query using node 0.
   110  	registry := tc.Server(0).ExecutorConfig().(sql.ExecutorConfig).StmtDiagnosticsRecorder
   111  	reqID, err := registry.InsertRequestInternal(ctx, "INSERT INTO test VALUES (_)")
   112  	require.NoError(t, err)
   113  	reqRow := db0.QueryRow(
   114  		`SELECT completed, statement_diagnostics_id FROM system.statement_diagnostics_requests
   115  		 WHERE ID = $1`, reqID)
   116  	var completed bool
   117  	var traceID gosql.NullInt64
   118  	require.NoError(t, reqRow.Scan(&completed, &traceID))
   119  	require.False(t, completed)
   120  	require.False(t, traceID.Valid) // traceID should be NULL
   121  
   122  	// Repeatedly run the query through node 1 until we get a trace.
   123  	runUntilTraced := func(query string, reqID int64) {
   124  		testutils.SucceedsSoon(t, func() error {
   125  			// Run the query using node 1.
   126  			_, err = db1.Exec(query)
   127  			require.NoError(t, err)
   128  
   129  			// Check that the row from statement_diagnostics_request was marked as completed.
   130  			traceRow := db0.QueryRow(
   131  				`SELECT completed, statement_diagnostics_id FROM system.statement_diagnostics_requests
   132  				 WHERE ID = $1`, reqID)
   133  			require.NoError(t, traceRow.Scan(&completed, &traceID))
   134  			if !completed {
   135  				_, err := db0.Exec("DELETE FROM test")
   136  				require.NoError(t, err)
   137  				return fmt.Errorf("not completed yet")
   138  			}
   139  			return nil
   140  		})
   141  		require.True(t, traceID.Valid)
   142  
   143  		// Check the trace.
   144  		row := db0.QueryRow(`SELECT jsonb_pretty(trace) FROM system.statement_diagnostics
   145                           WHERE ID = $1`, traceID.Int64)
   146  		var json string
   147  		require.NoError(t, row.Scan(&json))
   148  		require.Contains(t, json, "traced statement")
   149  	}
   150  	runUntilTraced("INSERT INTO test VALUES (1)", reqID)
   151  
   152  	// Verify that we can handle multiple requests at the same time.
   153  	id1, err := registry.InsertRequestInternal(ctx, "INSERT INTO test VALUES (_)")
   154  	require.NoError(t, err)
   155  	id2, err := registry.InsertRequestInternal(ctx, "SELECT x FROM test")
   156  	require.NoError(t, err)
   157  	id3, err := registry.InsertRequestInternal(ctx, "SELECT x FROM test WHERE x > _")
   158  	require.NoError(t, err)
   159  
   160  	// Run the queries in a different order.
   161  	runUntilTraced("SELECT x FROM test", id2)
   162  	runUntilTraced("SELECT x FROM test WHERE x > 1", id3)
   163  	runUntilTraced("INSERT INTO test VALUES (2)", id1)
   164  }
   165  
   166  // TestChangePollInterval ensures that changing the polling interval takes effect.
   167  func TestChangePollInterval(t *testing.T) {
   168  	defer leaktest.AfterTest(t)()
   169  
   170  	// We'll inject a request filter to detect scans due to the polling.
   171  	tableStart := keys.SystemSQLCodec.TablePrefix(keys.StatementDiagnosticsRequestsTableID)
   172  	tableSpan := roachpb.Span{
   173  		Key:    tableStart,
   174  		EndKey: tableStart.PrefixEnd(),
   175  	}
   176  	var scanState = struct {
   177  		syncutil.Mutex
   178  		m map[uuid.UUID]struct{}
   179  	}{
   180  		m: map[uuid.UUID]struct{}{},
   181  	}
   182  	recordScan := func(id uuid.UUID) {
   183  		scanState.Lock()
   184  		defer scanState.Unlock()
   185  		scanState.m[id] = struct{}{}
   186  	}
   187  	numScans := func() int {
   188  		scanState.Lock()
   189  		defer scanState.Unlock()
   190  		return len(scanState.m)
   191  	}
   192  	waitForScans := func(atLeast int) (seen int) {
   193  		testutils.SucceedsSoon(t, func() error {
   194  			if seen = numScans(); seen < atLeast {
   195  				return errors.Errorf("expected at least %d scans, saw %d", atLeast, seen)
   196  			}
   197  			return nil
   198  		})
   199  		return seen
   200  	}
   201  	args := base.TestServerArgs{
   202  		Knobs: base.TestingKnobs{
   203  			Store: &kvserver.StoreTestingKnobs{
   204  				TestingRequestFilter: func(ctx context.Context, request roachpb.BatchRequest) *roachpb.Error {
   205  					if request.Txn == nil {
   206  						return nil
   207  					}
   208  					for _, req := range request.Requests {
   209  						if scan := req.GetScan(); scan != nil && scan.Span().Overlaps(tableSpan) {
   210  							recordScan(request.Txn.ID)
   211  							return nil
   212  						}
   213  					}
   214  					return nil
   215  				},
   216  			},
   217  		},
   218  	}
   219  	s, db, _ := serverutils.StartServer(t, args)
   220  	ctx := context.Background()
   221  	defer s.Stopper().Stop(ctx)
   222  
   223  	require.Equal(t, 1, waitForScans(1))
   224  	time.Sleep(time.Millisecond) // ensure no unexpected scan occur
   225  	require.Equal(t, 1, waitForScans(1))
   226  	_, err := db.Exec("SET CLUSTER SETTING sql.stmt_diagnostics.poll_interval = '200us'")
   227  	require.NoError(t, err)
   228  	waitForScans(10) // ensure several scans occur
   229  }