[wip] improve contention event logic test

docs/generated/sql/functions.md

@@ -2640,7 +2640,7 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.num_inverted_index_entries"></a><code>crdb_internal.num_inverted_index_entries(val: jsonb, version: <a href="int.html">int</a>) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>This function is used only by CockroachDB’s developers for testing purposes.</p>
 </span></td></tr>
-<tr><td><a name="crdb_internal.payloads_for_span"></a><code>crdb_internal.payloads_for_span(span ID: <a href="int.html">int</a>) &rarr; jsonb</code></td><td><span class="funcdesc"><p>Returns the payload(s) of the span whose ID is passed in the argument.</p>
+<tr><td><a name="crdb_internal.payloads_for_span"></a><code>crdb_internal.payloads_for_span(span_id: <a href="int.html">int</a>) &rarr; tuple{string AS payload_type, jsonb AS payload_jsonb}</code></td><td><span class="funcdesc"><p>Returns the payload(s) of the requested span.</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.pretty_key"></a><code>crdb_internal.pretty_key(raw_key: <a href="bytes.html">bytes</a>, skip_fields: <a href="int.html">int</a>) &rarr; <a href="string.html">string</a></code></td><td><span class="funcdesc"><p>This function is used only by CockroachDB’s developers for testing purposes.</p>
 </span></td></tr>

pkg/sql/logictest/testdata/logic_test/builtin_function

@@ -2913,3 +2913,34 @@ query T
 SELECT regexp_split_to_array('3aaa0AAa1', 'a+', 'i')
 ----
 {3,0,1}
+
+subtest crdb_internal.trace_id
+
+# switch users -- this one has no permissions so expect errors
+user testuser
+
+query error insufficient privilege
+SELECT * FROM crdb_internal.trace_id()
+
+user root
+
+query B
+SELECT count(*) = 1 FROM crdb_internal.trace_id()
+----
+true
+
+subtest crdb_internal.payloads_for_span
+
+# switch users -- this one has no permissions so expect errors
+user testuser
+
+query error pq: only users with the admin role are allowed to use crdb_internal.payloads_for_span
+SELECT * FROM crdb_internal.payloads_for_span(0)
+
+user root
+
+query TT colnames
+SELECT * FROM crdb_internal.payloads_for_span(0)
+WHERE false
+----
+payload_type payload_jsonb

pkg/sql/logictest/testdata/logic_test/contention_event

@@ -9,7 +9,7 @@ GRANT ADMIN TO testuser
 
 statement ok
 CREATE TABLE kv (k VARCHAR PRIMARY KEY, v VARCHAR);
-ALTER TABLE kv SPLIT AT VALUES ('b'), ('d'), ('q'), ('z');
+ALTER TABLE kv SPLIT AT VALUES ('b'), ('d'), ('q'), ('z')
 
 query TT
 SELECT * FROM kv
@@ -18,7 +18,7 @@ SELECT * FROM kv
 user testuser
 
 statement ok
-BEGIN;
+BEGIN
 
 statement ok
 INSERT INTO kv VALUES('k', 'v')
@@ -32,7 +32,7 @@ user root
 statement ok
 BEGIN;
 SET TRANSACTION PRIORITY HIGH;
-SELECT * FROM kv;
+SELECT * FROM kv
 
 user testuser
 
@@ -49,16 +49,20 @@ user root
 #
 # NB: this needs the 5node-pretend59315 config because otherwise the span is not
 # tracked.
-#
-query B
+query T
 WITH spans AS (
- SELECT span_id FROM crdb_internal.node_inflight_trace_spans
+ SELECT span_id
+ FROM crdb_internal.node_inflight_trace_spans
  WHERE trace_id = crdb_internal.trace_id()
-), payload_types AS (
- SELECT jsonb_array_elements(crdb_internal.payloads_for_span(span_id))->>'@type' AS payload_type
- FROM spans
-) SELECT count(*) > 0
- FROM payload_types
- WHERE payload_type = 'type.googleapis.com/cockroach.roachpb.ContentionEvent';
+), payloads AS (
+ SELECT *
+ FROM spans, LATERAL crdb_internal.payloads_for_span(spans.span_id)
+) SELECT
+ -- Strip the first prefix because we don't want to have this test
+ -- become sensitive to changes in descriptor ID allocation.
+ crdb_internal.pretty_key(decode(payload_jsonb->>'key', 'base64'), 1)
+ FROM payloads
+ WHERE payload_type = 'roachpb.ContentionEvent'
+ LIMIT 1
 ----
-true
+/1/"k"/0

pkg/sql/sem/builtins/builtins.go

@@ -73,7 +73,6 @@ import (
  "github.com/cockroachdb/cockroach/pkg/util/unaccent"
  "github.com/cockroachdb/cockroach/pkg/util/uuid"
  "github.com/cockroachdb/errors"
- pbtypes "github.com/gogo/protobuf/types"
  "github.com/knz/strtime"
 )
 
@@ -3618,7 +3617,9 @@ may increase either contention or retry errors, or both.`,
  return nil, err
  }
  if !isAdmin {
- return nil, pgerror.Newf(pgcode.InsufficientPrivilege, "user needs the admin role to view trace ID")
+ if err := checkPrivilegedUser(ctx); err != nil {
+ return nil, err
+ }
  }
 
  sp := tracing.SpanFromContext(ctx.Context)
@@ -3633,50 +3634,6 @@ may increase either contention or retry errors, or both.`,
  },
  ),
 
- "crdb_internal.payloads_for_span": makeBuiltin(
- tree.FunctionProperties{Category: categorySystemInfo},
- tree.Overload{
- Types: tree.ArgTypes{{"span ID", types.Int}},
- ReturnType: tree.FixedReturnType(types.Jsonb),
- Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
- // The user must be an admin to use this builtin.
- isAdmin, err := ctx.SessionAccessor.HasAdminRole(ctx.Context)
- if err != nil {
- return nil, err
- }
- if !isAdmin {
- return nil, pgerror.Newf(pgcode.InsufficientPrivilege, "user needs the admin role to view payloads")
- }
-
- builder := json.NewArrayBuilder(len(args))
-
- spanID := uint64(*(args[0].(*tree.DInt)))
- span, found := ctx.Settings.Tracer.GetActiveSpanFromID(spanID)
- // A span may not be found if its ID was surfaced previously but its
- // corresponding trace has ended by the time this builtin was run.
- if !found {
- // Returns an empty JSON array.
- return tree.NewDJSON(builder.Build()), nil
- }
-
- for _, rec := range span.GetRecording() {
- rec.Structured(func(item *pbtypes.Any) {
- payload, err := protoreflect.MessageToJSON(item, true /* emitDefaults */)
- if err != nil {
- return
- }
- if payload != nil {
- builder.Add(payload)
- }
- })
- }
- return tree.NewDJSON(builder.Build()), nil
- },
- Info: "Returns the payload(s) of the span whose ID is passed in the argument.",
- Volatility: tree.VolatilityVolatile,
- },
- ),
-
  "crdb_internal.locality_value": makeBuiltin(
  tree.FunctionProperties{Category: categorySystemInfo},
  tree.Overload{

pkg/sql/sem/builtins/generator_builtins.go

@@ -13,6 +13,7 @@ package builtins
 import (
  "bytes"
  "context"
+ "strings"
  "time"
 
  "github.com/cockroachdb/cockroach/pkg/keys"
@@ -22,13 +23,16 @@ import (
  "github.com/cockroachdb/cockroach/pkg/sql/lex"
  "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
  "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+ "github.com/cockroachdb/cockroach/pkg/sql/protoreflect"
  "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
  "github.com/cockroachdb/cockroach/pkg/sql/types"
  "github.com/cockroachdb/cockroach/pkg/util/arith"
  "github.com/cockroachdb/cockroach/pkg/util/duration"
  "github.com/cockroachdb/cockroach/pkg/util/errorutil"
  "github.com/cockroachdb/cockroach/pkg/util/json"
+ "github.com/cockroachdb/cockroach/pkg/util/tracing"
  "github.com/cockroachdb/errors"
+ pbtypes "github.com/gogo/protobuf/types"
 )
 
 // See the comments at the start of generators.go for details about
@@ -326,6 +330,22 @@ var generators = map[string]builtinDefinition{
  tree.VolatilityVolatile,
  ),
  ),
+
+ "crdb_internal.payloads_for_span": makeBuiltin(
+ tree.FunctionProperties{
+ Class: tree.GeneratorClass,
+ Category: categorySystemInfo,
+ },
+ makeGeneratorOverload(
+ tree.ArgTypes{
+ {Name: "span_id", Typ: types.Int},
+ },
+ payloadsForSpanGeneratorType,
+ makePayloadsForSpanGenerator,
+ "Returns the payload(s) of the requested span.",
+ tree.VolatilityVolatile,
+ ),
+ ),
 }
 
 func makeGeneratorOverload(
@@ -1363,3 +1383,133 @@ func (rk *rangeKeyIterator) Values() (tree.Datums, error) {
 
 // Close implements the tree.ValueGenerator interface.
 func (rk *rangeKeyIterator) Close() {}
+
+var payloadsForSpanGeneratorLabels = []string{"payload_type", "payload_jsonb"}
+
+var payloadsForSpanGeneratorType = types.MakeLabeledTuple(
+ []*types.T{types.String, types.Jsonb},
+ payloadsForSpanGeneratorLabels,
+)
+
+// payloadsForSpanGenerator is a value generator that iterates over all payloads
+// over all recordings for a given Span.
+type payloadsForSpanGenerator struct {
+ // The span to iterate over.
+ span *tracing.Span
+
+ // recordingIndex maintains the current position of the index of the iterator
+ // in the list of recordings surfaced by a given span. The payloads of the
+ // recording that this iterator points to are buffered in `payloads`
+ recordingIndex int
+
+ // payloads represents all payloads for a given recording currently accessed
+ // by the iterator, and accesses more in a streaming fashion.
+ payloads []json.JSON
+
+ // payloadIndex maintains the current position of the index of the iterator
+ // in the list of `payloads` associated with a given recording.
+ payloadIndex int
+}
+
+func makePayloadsForSpanGenerator(
+ ctx *tree.EvalContext, args tree.Datums,
+) (tree.ValueGenerator, error) {
+ // The user must be an admin to use this builtin.
+ isAdmin, err := ctx.SessionAccessor.HasAdminRole(ctx.Context)
+ if err != nil {
+ return nil, err
+ }
+ if !isAdmin {
+ return nil, pgerror.Newf(
+ pgcode.InsufficientPrivilege,
+ "only users with the admin role are allowed to use crdb_internal.payloads_for_span",
+ )
+ }
+ spanID := uint64(*(args[0].(*tree.DInt)))
+ span, found := ctx.Settings.Tracer.GetActiveSpanFromID(spanID)
+ if !found {
+ return nil, nil
+ }
+
+ return &payloadsForSpanGenerator{span: span}, nil
+}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) ResolvedType() *types.T {
+ return payloadsForSpanGeneratorType
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Start(_ context.Context, _ *kv.Txn) error {
+ // The user of the generator first calls Next(), then Values(), so the index
+ // managing the iterator's position needs to start at -1 instead of 0.
+ p.recordingIndex = -1
+ p.payloadIndex = -1
+
+ return nil
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Next(_ context.Context) (bool, error) {
+ p.payloadIndex++
+
+ // If payloadIndex is within payloads and there are some payloads, then we
+ // have more buffered payloads to return.
+ if p.payloads != nil && p.payloadIndex < len(p.payloads) {
+ return true, nil
+ }
+
+ // Otherwise either there are no payloads or we have exhausted the payloads in
+ // our current recording, and we need to access another set of payloads from
+ // another recording.
+ p.payloads = nil
+
+ // Keep searching recordings for one with a valid (non-nil) payload.
+ for p.payloads == nil {
+ p.recordingIndex++
+ // If there are no more recordings, then we cannot continue.
+ if !(p.recordingIndex < p.span.GetRecording().Len()) {
+ return false, nil
+ }
+ currRecording := p.span.GetRecording()[p.recordingIndex]
+ currRecording.Structured(func(item *pbtypes.Any) {
+ payload, err := protoreflect.MessageToJSON(item, true /* emitDefaults */)
+ if err != nil {
+ return
+ }
+ if payload != nil {
+ p.payloads = append(p.payloads, payload)
+ }
+ })
+ }
+
+ p.payloadIndex = 0
+ return true, nil
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Values() (tree.Datums, error) {
+ payload := p.payloads[p.payloadIndex]
+ payloadTypeAsJSON, err := payload.FetchValKey("@type")
+ if err != nil {
+ return nil, err
+ }
+
+ // We trim the proto type prefix as well as the enclosing double quotes
+ // leftover from JSON value conversion.
+ payloadTypeAsString := strings.TrimSuffix(
+ strings.TrimPrefix(
+ payloadTypeAsJSON.String(),
+ "\"type.googleapis.com/cockroach.",
+ ),
+ "\"",
+ )
+
+ return tree.Datums{
+ tree.NewDString(payloadTypeAsString),
+ tree.NewDJSON(payload),
+ }, nil
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Close() {}

pkg/sql/sem/tree/eval.go

@@ -3940,7 +3940,7 @@ func EvalComparisonExprWithSubOperator(
  return evalDatumsCmp(ctx, expr.Operator, expr.SubOperator, expr.Fn, left, datums)
 }
 
-// EvalArgsAndGetGenerator evaluates the arguments and instanciates a
+// EvalArgsAndGetGenerator evaluates the arguments and instantiates a
 // ValueGenerator for use by set projections.
 func (expr *FuncExpr) EvalArgsAndGetGenerator(ctx *EvalContext) (ValueGenerator, error) {
  if expr.fn == nil || expr.fnProps.Class != GeneratorClass {