sql,kv: reduce sets of writes included into LeafTxnInputState

This commit reduces the InFlightWrites and BufferedWrites sets that are
included in LeafTxnInputState to only those that overlap with the key
spans that will be read by the DistSQL processors. We now construct an
interval tree that contains all spans read by TableReaders and
TableID / IndexID specific key spans read by JoinReaders, InvertedJoiners,
and ZigzagJoiners, and all non-overlapping writes are not included into
the proto.

We skip constructing the reads tree if the txn hasn't performed any
writes (avoiding redundant work). We also introduce a session variable
that disables this optimization as an escape hatch.

All processors have been audited to fall into one of three categories:
- if a core might need to have access to the full write sets of the leaf
txn, it goes into the "unsafe" category;
- if a core is not on the "hot" query path (i.e. part of bulk operations,
etc), it goes into the "unoptimized" category;
- all other processors are examined more thoroughly in terms of what
operations the concrete spec might perform, and it's then decided between
"unsafe" and "safe".

Let's focus on the last category a bit more. There are two main ways
these processors worth optimizing can use the txn:
- explicitly, when fetching the data from disk (4 processor cores
mentioned above). (Note that vector search processors also use the txn
explicitly, but we currently put them into the first category since we
don't distribute plans with them.)
- implicitly, when evaluating some expression. All processor cores and
post-process specs that store `execinfrapb.Expression` are now checked
for being "unsafe".

Defining what expression is "unsafe" is a bit hand-wavy: if an
expression might use the txn somehow (e.g. it's a builtin function call),
then we deem it "unsafe".

LocalPlanNode gets special treatment too since an arbitrary planNode
might use the txn. This commit doesn't perform audit of those, so we
mark almost all cases as "unsafe".

Release note: None

Author: yuzefovich

pkg/kv/BUILD.bazel

@@ -32,6 +32,7 @@ go_library(
         "//pkg/util/duration",
         "//pkg/util/errorutil/unimplemented",
         "//pkg/util/hlc",
+        "//pkg/util/interval",
         "//pkg/util/log",
         "//pkg/util/protoutil",
         "//pkg/util/retry",

pkg/kv/kvclient/kvcoord/BUILD.bazel

@@ -78,6 +78,7 @@ go_library(
         "//pkg/util/grpcutil",
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
+        "//pkg/util/interval",
         "//pkg/util/iterutil",
         "//pkg/util/limit",
         "//pkg/util/log",

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/debugutil"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -191,7 +192,12 @@ type txnInterceptor interface {
 
 	// populateLeafInputState populates the given input payload
 	// for a LeafTxn.
-	populateLeafInputState(*roachpb.LeafTxnInputState)
+	//
+	// readsTree, when non-nil, specifies an interval tree of key spans that
+	// will be read by the caller. As such, any non-overlapping writes could be
+	// ignored when populating the LeafTxnInputState. If readsTree is nil, then
+	// all writes should be included.
+	populateLeafInputState(*roachpb.LeafTxnInputState, interval.Tree)
 
 	// initializeLeaf updates any internal state held inside the interceptor
 	// from the given LeafTxn input state.
@@ -1407,7 +1413,7 @@ func (tc *TxnCoordSender) Active() bool {
 
 // GetLeafTxnInputState is part of the kv.TxnSender interface.
 func (tc *TxnCoordSender) GetLeafTxnInputState(
-	ctx context.Context,
+	ctx context.Context, readsTree interval.Tree,
 ) (*roachpb.LeafTxnInputState, error) {
 	tis := new(roachpb.LeafTxnInputState)
 	tc.mu.Lock()
@@ -1421,7 +1427,7 @@ func (tc *TxnCoordSender) GetLeafTxnInputState(
 	// Copy mutable state so access is safe for the caller.
 	tis.Txn = tc.mu.txn
 	for _, reqInt := range tc.interceptorStack {
-		reqInt.populateLeafInputState(tis)
+		reqInt.populateLeafInputState(tis, readsTree)
 	}
 
 	// Also mark the TxnCoordSender as "active".  This prevents changing
@@ -1693,6 +1699,7 @@ func (tc *TxnCoordSender) TestingShouldRetry() bool {
 	return false
 }
 
+// TODO(148760): this doesn't work under Read Committed isolation.
 func (tc *TxnCoordSender) hasPerformedReadsLocked() bool {
 	return !tc.interceptorAlloc.txnSpanRefresher.refreshFootprint.empty()
 }

pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go

@@ -1004,7 +1004,7 @@ func TestTxnMultipleCoord(t *testing.T) {
 	}
 
 	// New create a second, leaf coordinator.
-	leafInputState, err := txn.GetLeafTxnInputState(ctx)
+	leafInputState, err := txn.GetLeafTxnInputState(ctx, nil /* readsTree */)
 	require.NoError(t, err)
 	txn2 := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState, nil /* header */)
 
@@ -2739,7 +2739,7 @@ func TestLeafTxnClientRejectError(t *testing.T) {
 
 	ctx := context.Background()
 	rootTxn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
-	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
+	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx, nil /* readsTree */)
 	require.NoError(t, err)
 
 	// New create a second, leaf coordinator.
@@ -2771,7 +2771,7 @@ func TestUpdateRootWithLeafFinalStateInAbortedTxn(t *testing.T) {
 	ctx := context.Background()
 
 	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
-	leafInputState, err := txn.GetLeafTxnInputState(ctx)
+	leafInputState, err := txn.GetLeafTxnInputState(ctx, nil /* readsTree */)
 	require.NoError(t, err)
 	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0, leafInputState, nil /* header */)
 
@@ -2785,7 +2785,7 @@ func TestUpdateRootWithLeafFinalStateInAbortedTxn(t *testing.T) {
 	}
 
 	// Check that the transaction was not updated.
-	leafInputState2, err := txn.GetLeafTxnInputState(ctx)
+	leafInputState2, err := txn.GetLeafTxnInputState(ctx, nil /* readsTree */)
 	require.NoError(t, err)
 	if leafInputState2.Txn.Status != roachpb.PENDING {
 		t.Fatalf("expected PENDING txn, got: %s", leafInputState2.Txn.Status)
@@ -2977,7 +2977,7 @@ func TestTxnTypeCompatibleWithBatchRequest(t *testing.T) {
 	defer s.Stop()
 
 	rootTxn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
-	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
+	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx, nil /* readsTree */)
 	require.NoError(t, err)
 	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState, nil /* header */)
 

pkg/kv/kvclient/kvcoord/txn_interceptor_committer.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/errors"
@@ -611,7 +612,7 @@ func (tc *txnCommitter) maybeDisable1PC(ba *kvpb.BatchRequest) {
 func (tc *txnCommitter) setWrapped(wrapped lockedSender) { tc.wrapped = wrapped }
 
 // populateLeafInputState is part of the txnInterceptor interface.
-func (*txnCommitter) populateLeafInputState(*roachpb.LeafTxnInputState) {}
+func (*txnCommitter) populateLeafInputState(*roachpb.LeafTxnInputState, interval.Tree) {}
 
 // initializeLeaf is part of the txnInterceptor interface.
 func (*txnCommitter) initializeLeaf(tis *roachpb.LeafTxnInputState) {}

pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/metamorphic"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
@@ -283,7 +284,7 @@ func (h *txnHeartbeater) setWrapped(wrapped lockedSender) {
 }
 
 // populateLeafInputState is part of the txnInterceptor interface.
-func (*txnHeartbeater) populateLeafInputState(*roachpb.LeafTxnInputState) {}
+func (*txnHeartbeater) populateLeafInputState(*roachpb.LeafTxnInputState, interval.Tree) {}
 
 // initializeLeaf is part of the txnInterceptor interface.
 func (*txnHeartbeater) initializeLeaf(tis *roachpb.LeafTxnInputState) {}

pkg/kv/kvclient/kvcoord/txn_interceptor_metric_recorder.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 )
 
@@ -59,7 +60,7 @@ func (m *txnMetricRecorder) SendLocked(
 func (m *txnMetricRecorder) setWrapped(wrapped lockedSender) { m.wrapped = wrapped }
 
 // populateLeafInputState is part of the txnInterceptor interface.
-func (*txnMetricRecorder) populateLeafInputState(*roachpb.LeafTxnInputState) {}
+func (*txnMetricRecorder) populateLeafInputState(*roachpb.LeafTxnInputState, interval.Tree) {}
 
 // initializeLeaf is part of the txnInterceptor interface.
 func (*txnMetricRecorder) initializeLeaf(tis *roachpb.LeafTxnInputState) {}

pkg/kv/kvclient/kvcoord/txn_interceptor_pipeliner.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
@@ -947,8 +948,26 @@ func (tp *txnPipeliner) setWrapped(wrapped lockedSender) {
 }
 
 // populateLeafInputState is part of the txnInterceptor interface.
-func (tp *txnPipeliner) populateLeafInputState(tis *roachpb.LeafTxnInputState) {
-	tis.InFlightWrites = tp.ifWrites.asSlice()
+func (tp *txnPipeliner) populateLeafInputState(
+	tis *roachpb.LeafTxnInputState, readsTree interval.Tree,
+) {
+	if tp.ifWrites.len() == 0 {
+		return
+	}
+	if readsTree == nil {
+		tis.InFlightWrites = tp.ifWrites.asSlice()
+		return
+	}
+	var sp roachpb.Span
+	tp.ifWrites.ascend(func(w *inFlightWrite) {
+		sp.Key = w.Key
+		sp.EndKey = w.Key.Next()
+		if overlaps := readsTree.DoMatching(
+			func(interval.Interface) (done bool) { return true }, sp.AsRange(),
+		); overlaps {
+			tis.InFlightWrites = append(tis.InFlightWrites, w.SequencedWrite)
+		}
+	})
 }
 
 // initializeLeaf is part of the txnInterceptor interface.

pkg/kv/kvclient/kvcoord/txn_interceptor_seq_num_allocator.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/errors"
 )
 
@@ -116,7 +117,9 @@ func (s *txnSeqNumAllocator) SendLocked(
 func (s *txnSeqNumAllocator) setWrapped(wrapped lockedSender) { s.wrapped = wrapped }
 
 // populateLeafInputState is part of the txnInterceptor interface.
-func (s *txnSeqNumAllocator) populateLeafInputState(tis *roachpb.LeafTxnInputState) {
+func (s *txnSeqNumAllocator) populateLeafInputState(
+	tis *roachpb.LeafTxnInputState, tree interval.Tree,
+) {
 	tis.Txn.Sequence = s.writeSeq
 	tis.SteppingModeEnabled = bool(s.steppingMode)
 	tis.ReadSeqNum = s.readSeq

pkg/kv/kvclient/kvcoord/txn_interceptor_span_refresher.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
@@ -712,7 +713,9 @@ func (sr *txnSpanRefresher) maxRefreshAttempts() int {
 func (sr *txnSpanRefresher) setWrapped(wrapped lockedSender) { sr.wrapped = wrapped }
 
 // populateLeafInputState is part of the txnInterceptor interface.
-func (sr *txnSpanRefresher) populateLeafInputState(tis *roachpb.LeafTxnInputState) {
+func (sr *txnSpanRefresher) populateLeafInputState(
+	tis *roachpb.LeafTxnInputState, tree interval.Tree,
+) {
 	tis.RefreshInvalid = sr.refreshInvalid
 }