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 }