Merge #143991 #144073 #144093

143991: changefeedccl: add roachtest and parallel test for enriched envelope r=andyyang890 a=asg0451

Add a roachtest and a parallelized test for enriched envelope feeds.

Part of: #139660

Release note: None

144073: catalog/lease: fix race condition releasing old versions r=fqazi a=fqazi

Previously, a race condition existed between the name cache and the logic used to clean up old versions of leased descriptors. This issue was introduced when atomic operations were used to manipulate the descriptor version states. Specifically, after ensuring the reference count is zero, it's necessary to acquire a lock before releasing the stored lease. This patch adds logic to acquire and release a mutex when cleaning up old versions.

Fixes: #143815

Release note: None

144093: kvclient: enable savepoint rollbacks with buffered writes r=arulajmani a=arulajmani

First commit from #144058

----

This patch correctly handles savepoint rollbacks when buffered writes
are enabled. Whenever a savepoint is rolled back, any writes that belong
to the savepoint are removed from the buffer.

Closes #139059

Release note: None

Co-authored-by: Miles Frankel <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>
Co-authored-by: Arul Ajmani <[email protected]>

Author: 4 people

pkg/ccl/changefeedccl/cdctest/validator.go

@@ -1017,21 +1017,26 @@ func ParseJSONValueTimestamps(v []byte) (updated, resolved hlc.Timestamp, err er
 	var valueRaw struct {
 		Resolved string `json:"resolved"`
 		Updated  string `json:"updated"`
+		Source   struct {
+			TsHLC string `json:"ts_hlc"`
+		} `json:"source"`
 	}
 	if err := gojson.Unmarshal(v, &valueRaw); err != nil {
 		return hlc.Timestamp{}, hlc.Timestamp{}, errors.Wrapf(err, "parsing [%s] as json", v)
 	}
+
 	if valueRaw.Updated != `` {
-		var err error
-		updated, err = hlc.ParseHLC(valueRaw.Updated)
-		if err != nil {
+		if updated, err = hlc.ParseHLC(valueRaw.Updated); err != nil {
 			return hlc.Timestamp{}, hlc.Timestamp{}, err
 		}
 	}
 	if valueRaw.Resolved != `` {
-		var err error
-		resolved, err = hlc.ParseHLC(valueRaw.Resolved)
-		if err != nil {
+		if resolved, err = hlc.ParseHLC(valueRaw.Resolved); err != nil {
+			return hlc.Timestamp{}, hlc.Timestamp{}, err
+		}
+	}
+	if valueRaw.Source.TsHLC != `` {
+		if updated, err = hlc.ParseHLC(valueRaw.Source.TsHLC); err != nil {
 			return hlc.Timestamp{}, hlc.Timestamp{}, err
 		}
 	}

pkg/ccl/changefeedccl/changefeed_test.go

@@ -4193,6 +4193,99 @@ func TestChangefeedEnriched(t *testing.T) {
 	}
 }
 
+// TestChangefeedsParallelEnriched tests that multiple changefeeds can run in
+// parallel with the enriched envelope. It is most useful under race, to ensure
+// that there is no accidental data races in the encoders and source providers.
+func TestChangefeedsParallelEnriched(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	const numFeeds = 10
+	const maxIterations = 1_000_000_000
+	const maxRows = 100
+
+	testFn := func(t *testing.T, s TestServer, f cdctest.TestFeedFactory) {
+		db := sqlutils.MakeSQLRunner(s.DB)
+		db.Exec(t, `CREATE TABLE foo (a INT PRIMARY KEY, b STRING)`)
+
+		ctx, cancel := context.WithCancel(ctx)
+
+		var wg sync.WaitGroup
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			db := sqlutils.MakeSQLRunner(s.Server.SQLConn(t))
+			var i int
+			for i = 0; i < maxIterations && ctx.Err() == nil; i++ {
+				db.Exec(t, `UPSERT INTO d.foo VALUES ($1, $2)`, i%maxRows, fmt.Sprintf("hello %d", i))
+			}
+		}()
+
+		opts := `envelope='enriched'`
+
+		_, isKafka := f.(*kafkaFeedFactory)
+		useAvro := isKafka && rand.Intn(2) == 0
+		if useAvro {
+			t.Logf("using avro")
+			opts += `, format='avro'`
+		}
+		var feeds []cdctest.TestFeed
+		for range numFeeds {
+			feed := feed(t, f, fmt.Sprintf(`CREATE CHANGEFEED FOR foo WITH %s`, opts))
+			feeds = append(feeds, feed)
+		}
+
+		// consume from the feeds
+		for _, feed := range feeds {
+			feed := feed
+			msgCount := 0
+
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				for ctx.Err() == nil {
+					_, err := feed.Next()
+					if err != nil {
+						if errors.Is(err, context.Canceled) {
+							t.Errorf("error reading from feed: %v", err)
+						}
+						break
+					}
+					msgCount++
+				}
+				assert.GreaterOrEqual(t, msgCount, 0)
+			}()
+		}
+
+		// let the feeds run for a few seconds
+		select {
+		case <-time.After(5 * time.Second):
+		case <-ctx.Done():
+			t.Fatalf("%v", ctx.Err())
+		}
+
+		cancel()
+
+		for _, feed := range feeds {
+			closeFeed(t, feed)
+		}
+
+		doneWaiting := make(chan struct{})
+		go func() {
+			defer close(doneWaiting)
+			wg.Wait()
+		}()
+		select {
+		case <-doneWaiting:
+		case <-time.After(5 * time.Second):
+			t.Fatalf("timed out waiting for goroutines to finish")
+		}
+	}
+	// Sinkless testfeeds have some weird shutdown behaviours, so exclude them for now.
+	cdcTest(t, testFn, feedTestRestrictSinks("kafka", "pubsub", "webhook"))
+}
+
 func TestChangefeedEnrichedAvro(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/cmd/roachtest/tests/cdc.go

@@ -517,6 +517,7 @@ func makeDefaultFeatureFlags() cdcFeatureFlags {
 type feedArgs struct {
 	sinkType        sinkType
 	targets         []string
+	envelope        string
 	opts            map[string]string
 	assumeRole      string
 	tolerateErrors  bool
@@ -547,15 +548,15 @@ func (ct *cdcTester) newChangefeed(args feedArgs) changefeedJob {
 
 	targetsStr := strings.Join(args.targets, ", ")
 
+	if args.envelope == "" {
+		args.envelope = "wrapped"
+	}
+
 	feedOptions := make(map[string]string)
 	feedOptions["min_checkpoint_frequency"] = "'10s'"
+	feedOptions["envelope"] = args.envelope
 	if args.sinkType == cloudStorageSink || args.sinkType == webhookSink {
-		// Webhook and cloudstorage don't have a concept of keys and therefore
-		// require envelope=wrapped
-		feedOptions["envelope"] = "wrapped"
-
 		feedOptions["resolved"] = "'10s'"
-
 	} else {
 		feedOptions["resolved"] = ""
 	}
@@ -2311,6 +2312,40 @@ func registerCDC(r registry.Registry) {
 		Timeout:          1 * time.Hour,
 		Run:              runCDCMultipleSchemaChanges,
 	})
+	r.Add(registry.TestSpec{
+		Name:             "cdc/tpcc-100/10min/sink=kafka/envelope=enriched",
+		Owner:            registry.OwnerCDC,
+		Benchmark:        true,
+		Cluster:          r.MakeClusterSpec(4, spec.WorkloadNode(), spec.CPU(16)),
+		Leases:           registry.MetamorphicLeases,
+		CompatibleClouds: registry.AllClouds,
+		Suites:           registry.Suites(registry.Nightly),
+		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
+			ct := newCDCTester(ctx, t, c)
+			defer ct.Close()
+
+			ct.runTPCCWorkload(tpccArgs{warehouses: 100, duration: "10m"})
+
+			feed := ct.newChangefeed(feedArgs{
+				sinkType: kafkaSink,
+				envelope: "enriched",
+				targets:  allTpccTargets,
+				kafkaArgs: kafkaFeedArgs{
+					validateOrder: true,
+				},
+				opts: map[string]string{
+					"initial_scan":        "'no'",
+					"updated":             "",
+					"enriched_properties": "source",
+				},
+			})
+			ct.runFeedLatencyVerifier(feed, latencyTargets{
+				initialScanLatency: 3 * time.Minute,
+				steadyLatency:      10 * time.Minute,
+			})
+			ct.waitForWorkload()
+		},
+	})
 }
 
 const (

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/binary"
 	"slices"
+	"sort"
 	"unsafe"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
@@ -423,7 +424,36 @@ func (twb *txnWriteBuffer) epochBumpedLocked() {}
 func (twb *txnWriteBuffer) createSavepointLocked(context.Context, *savepoint) {}
 
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
-func (twb *txnWriteBuffer) rollbackToSavepointLocked(ctx context.Context, s savepoint) {}
+func (twb *txnWriteBuffer) rollbackToSavepointLocked(ctx context.Context, s savepoint) {
+	toDelete := make([]*bufferedWrite, 0)
+	it := twb.buffer.MakeIter()
+	for it.First(); it.Valid(); it.Next() {
+		bufferedVals := it.Cur().vals
+		// NB: the savepoint is being rolled back to s.seqNum (inclusive). So,
+		// idx is the index of the first value that is considered rolled back.
+		idx := sort.Search(len(bufferedVals), func(i int) bool {
+			return bufferedVals[i].seq >= s.seqNum
+		})
+		if idx == len(bufferedVals) {
+			// No writes are being rolled back.
+			continue
+		}
+		// Update size bookkeeping for the values we're rolling back.
+		for i := idx; i < len(bufferedVals); i++ {
+			twb.bufferSize -= bufferedVals[i].size()
+		}
+		// Rollback writes by truncating the buffered values.
+		it.Cur().vals = bufferedVals[:idx]
+		if len(it.Cur().vals) == 0 {
+			// All writes have been rolled back; we should remove this key from
+			// the buffer entirely.
+			toDelete = append(toDelete, it.Cur())
+		}
+	}
+	for _, bw := range toDelete {
+		twb.removeFromBuffer(bw)
+	}
+}
 
 // closeLocked implements the txnInterceptor interface.
 func (twb *txnWriteBuffer) closeLocked() {}
@@ -1057,6 +1087,12 @@ func (twb *txnWriteBuffer) addToBuffer(key roachpb.Key, val roachpb.Value, seq e
 	}
 }
 
+// removeFromBuffer removes all buffered writes on a given key from the buffer.
+func (twb *txnWriteBuffer) removeFromBuffer(bw *bufferedWrite) {
+	twb.buffer.Delete(bw)
+	twb.bufferSize -= bw.size()
+}
+
 // flushBufferAndSendBatch flushes all buffered writes when sending the supplied
 // batch request to the KV layer. This is done by pre-pending the buffered
 // writes to the requests in the batch.
@@ -1092,11 +1128,11 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 
 	reqs := make([]kvpb.RequestUnion, 0, numBuffered+len(ba.Requests))
 
-	// Next, remove the buffered writes from the buffer and collect them into requests.
+	// Next, remove the buffered writes from the buffer and collect them into
+	// requests.
 	for _, bw := range toFlushBufferedWrites {
 		reqs = append(reqs, bw.toRequest())
-		twb.buffer.Delete(&bw)
-		twb.bufferSize -= bw.size()
+		twb.removeFromBuffer(&bw)
 	}
 
 	// Layers below us expect that writes inside a batch are in sequence number

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1614,3 +1614,101 @@ func TestTxnWriteBufferDeleteRange(t *testing.T) {
 	require.Len(t, br.Responses, 1)
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
+
+// TestTxnWriteBufferRollbackToSavepoint tests the savepoint rollback logic.
+func TestTxnWriteBufferRollbackToSavepoint(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+
+	txn := makeTxnProto()
+	txn.Sequence = 10
+	keyA, keyB, keyC := roachpb.Key("a"), roachpb.Key("b"), roachpb.Key("c")
+	valA, valA2, valB := "valA", "valA2", "valB"
+
+	ba := &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	putA := putArgs(keyA, valA, txn.Sequence)
+	txn.Sequence++
+	delC := delArgs(keyC, txn.Sequence)
+	ba.Add(putA)
+	ba.Add(delC)
+
+	numCalled := mockSender.NumCalled()
+	br, pErr := twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	// All the requests should be buffered and not make it past the
+	// txnWriteBuffer. The response returned should be indistinguishable.
+	require.Equal(t, numCalled, mockSender.NumCalled())
+	require.Len(t, br.Responses, 2)
+	require.IsType(t, &kvpb.PutResponse{}, br.Responses[0].GetInner())
+	require.IsType(t, &kvpb.DeleteResponse{}, br.Responses[1].GetInner())
+	// Verify the
+	expBufferedWrites := []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue("valA", 10)),
+		makeBufferedWrite(keyC, makeBufferedValue("", 11)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Create a savepoint. This is inclusive of the buffered Delete on keyC.
+	savepoint := &savepoint{seqNum: txn.Sequence}
+	twb.createSavepointLocked(ctx, savepoint)
+
+	// Add some new writes. A second write to keyA and a new one to keyB.
+	ba = &kvpb.BatchRequest{}
+	txn.Sequence++
+	putA2 := putArgs(keyA, valA2, txn.Sequence)
+	ba.Add(putA2)
+	txn.Sequence++
+	putB := putArgs(keyB, valB, txn.Sequence)
+	ba.Add(putB)
+
+	// All these writes should be buffered as well, with no KV requests sent yet.
+	numCalled = mockSender.NumCalled()
+	br, pErr = twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+
+	require.Equal(t, numCalled, mockSender.NumCalled())
+	require.Len(t, br.Responses, 2)
+	require.IsType(t, &kvpb.PutResponse{}, br.Responses[0].GetInner())
+	require.IsType(t, &kvpb.PutResponse{}, br.Responses[1].GetInner())
+	// Verify the state of the write buffer.
+	expBufferedWrites = []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue("valA", 10), makeBufferedValue("valA2", 12)),
+		makeBufferedWrite(keyB, makeBufferedValue("valB", 13)),
+		makeBufferedWrite(keyC, makeBufferedValue("", 11)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Now, Rollback to the savepoint. This should leave just one write in the
+	// buffer, that on keyA at seqnum 10.
+	twb.rollbackToSavepointLocked(ctx, *savepoint)
+	expBufferedWrites = []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue("valA", 10)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Commit the transaction.
+	ba = &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	ba.Add(&kvpb.EndTxnRequest{Commit: true})
+
+	mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+		require.Len(t, ba.Requests, 2)
+		require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+		require.IsType(t, &kvpb.EndTxnRequest{}, ba.Requests[1].GetInner())
+
+		br = ba.CreateReply()
+		br.Txn = ba.Txn
+		return br, nil
+	})
+
+	br, pErr = twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	require.Len(t, br.Responses, 1)
+	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
+}

pkg/sql/catalog/lease/descriptor_state.go

@@ -203,10 +203,15 @@ func (t *descriptorState) removeInactiveVersions() []*storedLease {
 	for _, desc := range append([]*descriptorVersionState(nil), t.mu.active.data...) {
 		if desc.refcount.Load() == 0 {
 			t.mu.active.remove(desc)
-			if l := desc.mu.lease; l != nil {
-				desc.mu.lease = nil
-				leases = append(leases, l)
-			}
+			func() {
+				desc.mu.Lock()
+				defer desc.mu.Unlock()
+				if l := desc.mu.lease; l != nil {
+					desc.mu.lease = nil
+					leases = append(leases, l)
+
+				}
+			}()
 		}
 	}
 	return leases