Merge #140673 #141724 #141928

140673: sqlstats: add new MaybeFlushWithDrainer method to persistedSqlStats r=kyle-a-wong a=kyle-a-wong

Adds a new SSDrainer interface that PersistedSQLStats uses to flush sql stats.

Instead of calling ConsumeStats on sslocal.SQLStats, persistedsqlstats will use a provided SSDrainer to either DrainStats or Reset.

The ConsumeStats method on sslocal.SQLStats has been removed and its logic has been redistributed to persistedsqlstats, as that was the only consumer of the method. Since all the logic now lives in persistedsqlstats, the flush related functions can be called directly instead of relying on nested callbacks.

Resolves: #139143
Epic: CRDB-45771
Release note: None


Note: This is a stack committed, only the last commit in the chain needs to be  reviewed

141724: kvclient: ensure locking Get requests are always sent to KV  r=yuzefovich a=arulajmani

First commit from #141672

----

A Get request to a key that's been previously buffered on the client
must always be served locally. However, if the Get request is locking
in nature, we must also acquire a lock at the leaseholder. This patch
addresses a TODO to do exactly this, while ensuring the response from
KV is discarded and served locally instead.

Informs #139054

Release note: None


141928: catalog/lease: populate SessionID for historical leases r=fqazi a=fqazi

Previously, we were seeing intermittent panics in scenarios where
historical descriptors were accessed followed by reads of newer
descriptor versions. This would happen because the SessionID on
historical versions would be blank, causing our checks for detecting
changes to fails. To address this, this patch always populates the
session ID inside the descriptor versions.

Fixes: #141877
Fixes: #141876
Fixes: #141855

Release note: None

Co-authored-by: Kyle Wong <[email protected]>
Co-authored-by: Arul Ajmani <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 4 people

pkg/ccl/serverccl/BUILD.bazel

@@ -16,7 +16,6 @@ go_library(
         "//pkg/server/serverpb",
         "//pkg/sql",
         "//pkg/sql/contention",
-        "//pkg/sql/sqlstats/persistedsqlstats",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",
         "//pkg/util/httputil",

pkg/ccl/serverccl/statusccl/tenant_status_test.go

@@ -591,8 +591,16 @@ func testResetSQLStatsRPCForTenant(
 			}
 
 			if flushed {
-				testTenant.TenantSQLStats().MaybeFlush(ctx, testTenant.GetTenant().AppStopper())
-				controlCluster.TenantSQLStats(serverccl.RandomServer).MaybeFlush(ctx, controlCluster.Tenant(0).GetTenant().AppStopper())
+				testTenantServer := testTenant.TenantSQLServer()
+				testTenantServer.GetSQLStatsProvider().MaybeFlush(
+					ctx,
+					testTenant.GetTenant().AppStopper(),
+				)
+				randomTenantServer := controlCluster.TenantSQLServer(serverccl.RandomServer)
+				randomTenantServer.GetSQLStatsProvider().MaybeFlush(
+					ctx,
+					controlCluster.Tenant(0).GetTenant().AppStopper(),
+				)
 			}
 
 			status := testTenant.TenantStatusSrv()

pkg/ccl/serverccl/tenant_test_utils.go

@@ -19,7 +19,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/contention"
-	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/persistedsqlstats"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/httputil"
@@ -41,7 +40,7 @@ type testTenant struct {
 	tenantConn               *gosql.DB
 	tenantDB                 *sqlutils.SQLRunner
 	tenantStatus             serverpb.SQLStatusServer
-	tenantSQLStats           *persistedsqlstats.PersistedSQLStats
+	tenantSQLServer          *sql.Server
 	tenantContentionRegistry *contention.Registry
 }
 
@@ -53,8 +52,8 @@ func (h *testTenant) GetTenantConn() *sqlutils.SQLRunner {
 	return h.tenantDB
 }
 
-func (h *testTenant) TenantSQLStats() *persistedsqlstats.PersistedSQLStats {
-	return h.tenantSQLStats
+func (h *testTenant) TenantSQLServer() *sql.Server {
+	return h.tenantSQLServer
 }
 
 func (h *testTenant) TenantStatusSrv() serverpb.SQLStatusServer {
@@ -78,7 +77,7 @@ type TestTenant interface {
 	GetTenant() serverutils.ApplicationLayerInterface
 	GetTenantDB() *gosql.DB
 	GetTenantConn() *sqlutils.SQLRunner
-	TenantSQLStats() *persistedsqlstats.PersistedSQLStats
+	TenantSQLServer() *sql.Server
 	TenantStatusSrv() serverpb.SQLStatusServer
 	TenantContentionRegistry() *contention.Registry
 	GetRPCContext() *rpc.Context
@@ -95,15 +94,15 @@ func newTestTenant(
 	tenant, tenantConn := serverutils.StartTenant(t, server, args)
 	sqlDB := sqlutils.MakeSQLRunner(tenantConn)
 	status := tenant.StatusServer().(serverpb.SQLStatusServer)
-	sqlStats := tenant.SQLServer().(*sql.Server).GetSQLStatsProvider()
+	sqlServer := tenant.SQLServer().(*sql.Server)
 	contentionRegistry := tenant.ExecutorConfig().(sql.ExecutorConfig).ContentionRegistry
 
 	return &testTenant{
 		tenant:                   tenant,
 		tenantConn:               tenantConn,
 		tenantDB:                 sqlDB,
 		tenantStatus:             status,
-		tenantSQLStats:           sqlStats,
+		tenantSQLServer:          sqlServer,
 		tenantContentionRegistry: contentionRegistry,
 	}
 }
@@ -197,7 +196,7 @@ type TenantClusterHelper interface {
 	TenantDB(idx serverIdx) *gosql.DB
 	TenantHTTPClient(t *testing.T, idx serverIdx, isAdmin bool) *httpClient
 	TenantAdminHTTPClient(t *testing.T, idx serverIdx) *httpClient
-	TenantSQLStats(idx serverIdx) *persistedsqlstats.PersistedSQLStats
+	TenantSQLServer(idx serverIdx) *sql.Server
 	TenantStatusSrv(idx serverIdx) serverpb.SQLStatusServer
 	TenantContentionRegistry(idx serverIdx) *contention.Registry
 	Cleanup(t *testing.T)
@@ -249,8 +248,8 @@ func (c tenantCluster) TenantAdminHTTPClient(t *testing.T, idx serverIdx) *httpC
 	return c.TenantHTTPClient(t, idx, true /* isAdmin */)
 }
 
-func (c tenantCluster) TenantSQLStats(idx serverIdx) *persistedsqlstats.PersistedSQLStats {
-	return c.Tenant(idx).TenantSQLStats()
+func (c tenantCluster) TenantSQLServer(idx serverIdx) *sql.Server {
+	return c.Tenant(idx).TenantSQLServer()
 }
 
 func (c tenantCluster) TenantStatusSrv(idx serverIdx) serverpb.SQLStatusServer {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -9,6 +9,7 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
@@ -289,9 +290,10 @@ func (twb *txnWriteBuffer) applyTransformations(
 			// TODO(yuzefovich): ensure that we elide the lock acquisition
 			// whenever possible (e.g. blind UPSERT in an implicit txn).
 			ts = append(ts, transformation{
-				stripped: true,
-				index:    i,
-				resp:     ru,
+				stripped:    true,
+				index:       i,
+				origRequest: req,
+				resp:        ru,
 			})
 			twb.addToBuffer(t.Key, t.Value, t.Sequence)
 
@@ -304,9 +306,10 @@ func (twb *txnWriteBuffer) applyTransformations(
 				FoundKey: false,
 			})
 			ts = append(ts, transformation{
-				stripped: true,
-				index:    i,
-				resp:     ru,
+				stripped:    true,
+				index:       i,
+				origRequest: req,
+				resp:        ru,
 			})
 			twb.addToBuffer(t.Key, roachpb.Value{}, t.Sequence)
 
@@ -315,25 +318,36 @@ func (twb *txnWriteBuffer) applyTransformations(
 			val, served := twb.maybeServeRead(t.Key, t.Sequence)
 			if served {
 				log.VEventf(ctx, 2, "serving %s on key %s from the buffer", t.Method(), t.Key)
-				var ru kvpb.ResponseUnion
+				var resp kvpb.ResponseUnion
 				getResp := &kvpb.GetResponse{}
 				if val.IsPresent() {
 					getResp.Value = val
 				}
-				ru.MustSetInner(getResp)
+				resp.MustSetInner(getResp)
+
+				stripped := true
+				if t.KeyLockingStrength != lock.None {
+					// Even though the Get request must be served from the buffer, as the
+					// transaction performed a previous write to the key, we still need to
+					// acquire a lock at the leaseholder. As a result, we can't strip the
+					// request from the batch.
+					//
+					// TODO(arul): we could eschew sending this request if we knew there
+					// was a sufficiently strong lock already present on the key.
+					stripped = false
+					baRemote.Requests = append(baRemote.Requests, ru)
+				}
+
 				ts = append(ts, transformation{
-					stripped: true,
-					index:    i,
-					resp:     ru,
+					stripped:    stripped,
+					index:       i,
+					origRequest: req,
+					resp:        resp,
 				})
 				// We've constructed a response that we'll stitch together with the
 				// result on the response path; eschew sending the request to the KV
 				// layer.
 				//
-				// TODO(arul): if this is a locking Get request, we can't omit the KV
-				// request. At least, not unless we know a lock is already present on
-				// the key.
-				//
 				// TODO(arul): if the ReturnRawMVCCValues flag is set, we'll need to
 				// flush the buffer.
 				continue
@@ -579,14 +593,14 @@ type transformation struct {
 	// index of the request in the original batch to which the transformation
 	// applies.
 	index int
+	// origRequest is the original request that was transformed.
+	origRequest kvpb.Request
 	// resp is locally produced response that needs to be merged with any
 	// responses returned by the KV layer. This is set for requests that can be
 	// evaluated locally (e.g. blind writes, reads that can be served entirely
-	// from the buffer). If non-empty, stripped must also be true.
+	// from the buffer). Must be set if stripped is true, but the converse doesn't
+	// hold.
 	resp kvpb.ResponseUnion
-	// origRequest is the original request that was transformed. Set iff stripped
-	// is false.
-	origRequest kvpb.Request
 }
 
 // toResp returns the response that should be added to the batch response as
@@ -599,8 +613,21 @@ func (t transformation) toResp(
 	}
 
 	var ru kvpb.ResponseUnion
-	switch br.GetInner().(type) {
-	case *kvpb.ScanResponse:
+	switch t.origRequest.(type) {
+	case *kvpb.PutRequest:
+		ru = t.resp
+	case *kvpb.DeleteRequest:
+		ru = t.resp
+	case *kvpb.GetRequest:
+		// Get requests must be served from the local buffer if a transaction
+		// performed a previous write to the key being read. However, Get requests
+		// must be sent to the KV layer (i.e. not be stripped) iff they are locking
+		// in nature.
+		req := t.origRequest.(*kvpb.GetRequest)
+		assertTrue(t.stripped == (req.KeyLockingStrength == lock.None),
+			"Get requests should either be stripped or be locking")
+		ru = t.resp
+	case *kvpb.ScanRequest:
 		scanResp, err := twb.mergeWithScanResp(
 			t.origRequest.(*kvpb.ScanRequest), br.GetInner().(*kvpb.ScanResponse),
 		)

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -11,8 +11,10 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
@@ -825,3 +827,108 @@ func TestTxnWriteBufferServesOverlappingReadsCorrectly(t *testing.T) {
 	require.Len(t, br.Responses, 1)
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
+
+// TestTxnWriteBufferLockingGetRequests ensures that locking get requests are
+// handled appropriately -- they're sent to KV, to acquire a lock, but the
+// read is served from the buffer (upholding read-your-own-write semantics).
+func TestTxnWriteBufferLockingGetRequests(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer()
+
+	txn := makeTxnProto()
+	txn.Sequence = 10
+	keyA := roachpb.Key("a")
+	valA := "val"
+
+	// Blindly write to keys A.
+	ba := &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	putA := putArgs(keyA, valA, txn.Sequence)
+	ba.Add(putA)
+
+	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.
+	require.Equal(t, numCalled, mockSender.NumCalled())
+	// Even though the txnWriteBuffer did not send any Put requests to the KV
+	// layer above, the responses should still be populated.
+	require.Len(t, br.Responses, 1)
+	require.Equal(t, br.Responses[0].GetInner(), &kvpb.PutResponse{})
+
+	// Verify the writes were buffered correctly.
+	expBufferedWrites := []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue(valA, 10)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Perform a locking read on keyA. Ensure a request is sent to the KV layer,
+	// but the response is served from the buffer.
+	testutils.RunValues(t, "str", []lock.Strength{lock.None, lock.Shared, lock.Exclusive, lock.Update}, func(t *testing.T, strength lock.Strength) {
+		txn.Sequence = 11
+		ba = &kvpb.BatchRequest{}
+		ba.Header = kvpb.Header{Txn: &txn}
+		getA := &kvpb.GetRequest{
+			RequestHeader:      kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+			KeyLockingStrength: strength,
+		}
+		ba.Add(getA)
+		numCalled = mockSender.NumCalled()
+
+		mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+			require.Len(t, ba.Requests, 1)
+			require.IsType(t, &kvpb.GetRequest{}, ba.Requests[0].GetInner())
+			require.Equal(t, strength, ba.Requests[0].GetInner().(*kvpb.GetRequest).KeyLockingStrength)
+			require.True(t, strength != lock.None) // non-locking Gets aren't sent to KV
+
+			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.Equal(t, roachpb.MakeValueFromString(valA), *br.Responses[0].GetInner().(*kvpb.GetResponse).Value)
+
+		var expNumCalled int
+		if strength == lock.None {
+			expNumCalled = numCalled // nothing should be sent to KV
+		} else {
+			expNumCalled = numCalled + 1 // a locking request should still be sent to KV
+		}
+		require.Equal(t, expNumCalled, mockSender.NumCalled())
+	})
+
+	// Lastly, for completeness, commit the transaction and ensure that the buffer
+	// is correctly flushed.
+	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)
+
+		// We now expect the buffer to be flushed along with the commit.
+		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)
+
+	// Even though we flushed the buffer, responses from the blind writes should
+	// not be returned.
+	require.Len(t, br.Responses, 1)
+	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
+}

pkg/sql/catalog/lease/descriptor_state.go

@@ -181,13 +181,12 @@ func newDescriptorVersionState(
 			version: int(desc.GetVersion()),
 		}
 		descState.mu.lease.sessionID = session.ID().UnsafeBytes()
-		descState.mu.session = session
-
 		if buildutil.CrdbTestBuild && !expiration.IsEmpty() {
 			panic(errors.AssertionFailedf("expiration should always be empty for "+
 				"session based leases (got: %s on Desc: %s(%d))", expiration.String(), desc.GetName(), desc.GetID()))
 		}
 	}
+	descState.mu.session = session
 	descState.mu.expiration = expiration
 
 	return descState

pkg/sql/catalog/lease/descriptor_version_state.go

@@ -60,9 +60,7 @@ type descriptorVersionState struct {
 		// when the version isn't associated with a lease.
 		expiration hlc.Timestamp
 
-		// The session that was used to acquire this descriptor version, which is
-		// only populated when the session based leasing mode is *at least* dual
-		// write.
+		// The session that was used to acquire this descriptor version.
 		session sqlliveness.Session
 
 		refcount int