sql: unblock user create from long-running transactions

The changes in #150747 to unblock user create were partially reverted
in #152297 due to it introducing a race condition that caused tests to
flake.
This change restores that behavior and accounts for the race by updating
the lease manager before acquiring a lease on tables used for user
access management.
Epic: CRDB-49398
Fixes: #138691

Release note: None

Author: bghal

pkg/sql/catalog/lease/lease.go

@@ -1002,7 +1002,7 @@ func (m *Manager) AcquireFreshestFromStore(ctx context.Context, id descpb.ID) er
 	attemptsMade := 0
 	for {
 		// Acquire a fresh lease.
-		didAcquire, err := acquireNodeLease(ctx, m, id, AcquireFreshestBlock)
+		didAcquire, err := m.acquireNodeLease(ctx, id, AcquireFreshestBlock)
 		if err != nil {
 			return err
 		}
@@ -1024,8 +1024,8 @@ func (m *Manager) AcquireFreshestFromStore(ctx context.Context, id descpb.ID) er
 // being dropped or offline, the error will be of type inactiveTableError.
 // The boolean returned is true if this call was actually responsible for the
 // lease acquisition.
-func acquireNodeLease(
-	ctx context.Context, m *Manager, id descpb.ID, typ AcquireType,
+func (m *Manager) acquireNodeLease(
+	ctx context.Context, id descpb.ID, typ AcquireType,
 ) (bool, error) {
 	start := timeutil.Now()
 	log.VEventf(ctx, 2, "acquiring lease for descriptor %d...", id)
@@ -1037,7 +1037,7 @@ func acquireNodeLease(
 		},
 		func(ctx context.Context) (interface{}, error) {
 			if m.IsDraining() {
-				return nil, errLeaseManagerIsDraining
+				return false, errLeaseManagerIsDraining
 			}
 			newest := m.findNewest(id)
 			var currentVersion descpb.DescriptorVersion
@@ -1051,26 +1051,75 @@ func acquireNodeLease(
 			if err != nil {
 				return false, errors.Wrapf(err, "lease acquisition was unable to resolve liveness session")
 			}
-			desc, regionPrefix, err := m.storage.acquire(ctx, session, id, currentVersion, currentSessionID)
-			if err != nil {
-				return nil, err
-			}
-			// If a nil descriptor is returned, then the latest version has already
-			// been leased. So, nothing needs to be done here.
-			if desc == nil {
-				return true, nil
+
+			doAcquisition := func() (catalog.Descriptor, error) {
+				desc, _, err := m.storage.acquire(ctx, session, id, currentVersion, currentSessionID)
+				if err != nil {
+					return nil, err
+				}
+
+				return desc, nil
 			}
-			t := m.findDescriptorState(id, false /* create */)
-			if t == nil {
-				return nil, errors.AssertionFailedf("could not find descriptor state for id %d", id)
+
+			doUpsertion := func(desc catalog.Descriptor) error {
+				t := m.findDescriptorState(id, false /* create */)
+				if t == nil {
+					return errors.AssertionFailedf("could not find descriptor state for id %d", id)
+				}
+				t.mu.Lock()
+				t.mu.takenOffline = false
+				defer t.mu.Unlock()
+				err = t.upsertLeaseLocked(ctx, desc, session, m.storage.getRegionPrefix())
+				if err != nil {
+					return err
+				}
+
+				return nil
 			}
-			t.mu.Lock()
-			t.mu.takenOffline = false
-			defer t.mu.Unlock()
-			err = t.upsertLeaseLocked(ctx, desc, session, regionPrefix)
-			if err != nil {
-				return nil, err
+
+			// These tables are special and can have their versions bumped
+			// without blocking on other nodes converging to that version.
+			if newest != nil && (id == keys.UsersTableID ||
+				id == keys.RoleMembersTableID ||
+				id == keys.RoleOptionsTableID ||
+				m.isMaybeSystemPrivilegesTable(ctx, id)) {
+
+				// The two-version invariant allows an update in lease manager
+				// without (immediately) acquiring a new lease. This prevents
+				// a race on where the lease is acquired but the manager isn't
+				// yet updated.
+				desc, err := m.maybeGetDescriptorWithoutValidation(ctx, id, true /* existenceRequired */)
+				if err != nil {
+					return false, err
+				}
+				err = doUpsertion(desc)
+				if err != nil {
+					return false, err
+				}
+
+				desc, err = doAcquisition()
+				if err != nil {
+					return false, err
+				}
+				if desc == nil {
+					return true, nil
+				}
+			} else {
+				desc, err := doAcquisition()
+				if err != nil {
+					return false, err
+				}
+				// If a nil descriptor is returned, then the latest version has already
+				// been leased. So, nothing needs to be done here.
+				if desc == nil {
+					return true, nil
+				}
+				err = doUpsertion(desc)
+				if err != nil {
+					return false, err
+				}
 			}
+
 			return true, nil
 		})
 	if m.testingKnobs.LeaseStoreTestingKnobs.LeaseAcquireResultBlockEvent != nil {
@@ -1084,6 +1133,22 @@ func acquireNodeLease(
 	return didAcquire, nil
 }
 
+var systemPrivilegesTableDescID atomic.Uint32
+
+// isMaybeSystemPrivilegesTable tries to determine if the given descriptor is
+// for the system privileges table. It depends on the namecache to resolve the
+// table descriptor after which it's memoized.
+func (m *Manager) isMaybeSystemPrivilegesTable(ctx context.Context, id descpb.ID) bool {
+	if systemPrivilegesTableDescID.Load() == uint32(descpb.InvalidID) {
+		if privilegesDesc, _ := m.names.get(ctx, keys.SystemDatabaseID, keys.SystemPublicSchemaID, "privileges", hlc.Timestamp{}); privilegesDesc != nil {
+			systemPrivilegesTableDescID.CompareAndSwap(uint32(descpb.InvalidID), uint32(privilegesDesc.GetID()))
+			privilegesDesc.Release(ctx)
+		}
+	}
+
+	return uint32(id) == systemPrivilegesTableDescID.Load()
+}
+
 // releaseLease deletes an entry from system.lease.
 func releaseLease(ctx context.Context, lease *storedLease, m *Manager) (released bool) {
 	// Force the release to happen synchronously, if we are draining or, when we
@@ -1729,7 +1794,7 @@ func (m *Manager) Acquire(
 				t.markAcquisitionStart(ctx)
 				defer t.markAcquisitionDone(ctx)
 				// Renew lease and retry. This will block until the lease is acquired.
-				_, errLease := acquireNodeLease(ctx, m, id, AcquireBlock)
+				_, errLease := m.acquireNodeLease(ctx, id, AcquireBlock)
 				return errLease
 			}(); err != nil {
 				return nil, err
@@ -2371,7 +2436,7 @@ func (m *Manager) refreshSomeLeases(ctx context.Context, refreshAndPurgeAllDescr
 						return
 					}
 				}
-				if _, err := acquireNodeLease(ctx, m, id, AcquireBackground); err != nil {
+				if _, err := m.acquireNodeLease(ctx, id, AcquireBackground); err != nil {
 					log.Dev.Errorf(ctx, "refreshing descriptor: %d lease failed: %s", id, err)
 
 					if errors.Is(err, catalog.ErrDescriptorNotFound) || errors.Is(err, catalog.ErrDescriptorDropped) {

pkg/sql/catalog/lease/lease_internal_test.go

@@ -571,8 +571,8 @@ CREATE TABLE t.%s (k CHAR PRIMARY KEY, v CHAR);
 	}
 	expiration := lease.Expiration(context.Background())
 	// Acquire another lease.
-	if _, err := acquireNodeLease(
-		context.Background(), leaseManager, tableDesc.GetID(), AcquireBlock,
+	if _, err := leaseManager.acquireNodeLease(
+		context.Background(), tableDesc.GetID(), AcquireBlock,
 	); err != nil {
 		t.Fatal(err)
 	}

pkg/sql/catalog/lease/lease_test.go

@@ -558,13 +558,6 @@ func TestWaitForNewVersion(testingT *testing.T) {
 	defer log.Scope(testingT).Close(testingT)
 
 	var params base.TestClusterArgs
-	params.ServerArgs.Knobs = base.TestingKnobs{
-		SQLLeaseManager: &lease.ManagerTestingKnobs{
-			LeaseStoreTestingKnobs: lease.StorageTestingKnobs{
-				LeaseAcquiredEvent: nil, // TODO
-			},
-		},
-	}
 	t := newLeaseTest(testingT, params)
 	defer t.cleanup()
 
@@ -587,7 +580,10 @@ func TestWaitForNewVersion(testingT *testing.T) {
 		defer cancel()
 
 		_, err := leaseMgr.WaitForNewVersion(timeoutCtx, descID, nil, retry.Options{})
-		require.ErrorIs(t, err, context.DeadlineExceeded)
+		if !(sqltestutils.IsClientSideQueryCanceledErr(err) ||
+			errors.Is(err, context.DeadlineExceeded)) {
+			t.Fatalf("The client or the context should have timed out. Unexpected error: %v", err)
+		}
 	}
 
 	t.mustAcquire(3, descID)

pkg/sql/grant_revoke_system.go

@@ -128,10 +128,11 @@ VALUES ($1, $2, $3, $4, (
 					return err
 				}
 				userPrivs, found := syntheticPrivDesc.FindUser(user)
+				emptyPrivs := !found || userPrivs.Privileges == 0
 
 				// For Public role and virtual tables, leave an empty
 				// row to indicate that SELECT has been revoked.
-				if !found && (n.grantOn == privilege.VirtualTable && user == username.PublicRoleName()) {
+				if emptyPrivs && (n.grantOn == privilege.VirtualTable && user == username.PublicRoleName()) {
 					_, err := params.p.InternalSQLTxn().ExecEx(
 						params.ctx,
 						`insert-system-privilege`,
@@ -151,7 +152,7 @@ VALUES ($1, $2, $3, $4, (
 
 				// If there are no entries remaining on the PrivilegeDescriptor for the user
 				// we can remove the entire row for the user.
-				if !found {
+				if emptyPrivs {
 					_, err := params.p.InternalSQLTxn().ExecEx(
 						params.ctx,
 						`delete-system-privilege`,

pkg/sql/tests/allow_role_memberships_to_change_during_transaction_test.go

@@ -12,18 +12,15 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
-	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
 )
 
-// TestAllowRoleMembershipsToChangeDuringTransaction ensures that when the
-// corresponding session variable is set for all sessions using cached
-// role memberships, performing new grants do not need to wait for open
-// transactions to complete.
+// TestAllowRoleMembershipsToChangeDuringTransaction ensures user operations
+// are not blocked by transactions with the corresponding session variable set.
 func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -40,14 +37,15 @@ func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
 		}
 	}
 
-	// Create three users: foo, bar, and baz.
+	// Create four users: foo, bar, biz, and baz.
 	// Use one of these users to hold open a transaction which uses a lease
 	// on the role_memberships table. Ensure that initially granting does
 	// wait on that transaction. Then set the session variable and ensure
 	// that it does not.
 	tdb := sqlutils.MakeSQLRunner(sqlDB)
 	tdb.Exec(t, "CREATE USER foo PASSWORD 'foo'")
 	tdb.Exec(t, "CREATE USER bar")
+	tdb.Exec(t, "CREATE USER biz")
 	tdb.Exec(t, "CREATE USER baz")
 	tdb.Exec(t, "GRANT admin TO foo")
 	tdb.Exec(t, "CREATE DATABASE db2")
@@ -56,20 +54,27 @@ func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
 	fooDB, cleanupFoo := openUser("foo", "db2")
 	defer cleanupFoo()
 
-	// In this first test, we want to make sure that the query cannot proceed
-	// until the outer transaction commits. We launch the grant while the
-	// transaction is open, wait a tad, make sure it hasn't made progress,
-	// then we commit the relevant transaction and ensure that it does finish.
+	// Ensure that the outer transaction blocks the user admin operations until
+	// it is committed and its lease on the initial versions of the tables are
+	// released.
 	t.Run("normal path waits", func(t *testing.T) {
-		fooTx, err := fooDB.BeginTx(ctx, nil)
+		outerTx, err := fooDB.BeginTx(ctx, nil)
 		require.NoError(t, err)
+		defer func() { _ = outerTx.Rollback() }()
+
 		var count int
 		require.NoError(t,
-			fooTx.QueryRow("SELECT count(*) FROM t").Scan(&count))
+			outerTx.QueryRow("SELECT count(*) FROM t").Scan(&count))
 		require.Equal(t, 0, count)
+
+		// the first user operation shouldn't block
+		_, err = sqlDB.Exec("GRANT biz TO bar;")
+		require.NoError(t, err)
+
+		// the outer tx still has a lease on the initial table version so this should block
 		errCh := make(chan error, 1)
 		go func() {
-			_, err := sqlDB.Exec("GRANT bar TO baz;")
+			_, err = sqlDB.Exec("GRANT baz TO bar;")
 			errCh <- err
 		}()
 		select {
@@ -79,74 +84,38 @@ func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
 		case err := <-errCh:
 			t.Fatalf("expected transaction to block, got err: %v", err)
 		}
-		require.NoError(t, fooTx.Commit())
+		require.NoError(t, outerTx.Commit())
 		require.NoError(t, <-errCh)
 	})
 
-	// In this test we ensure that we can perform role grant and revoke
-	// operations while the transaction which uses the relevant roles
-	// remains open. We ensure that the transaction still succeeds and
-	// that the operations occur in a timely manner.
-	t.Run("session variable prevents waiting during GRANT and REVOKE", func(t *testing.T) {
+	// When the session variable is set, the outer transaction should not block
+	// the user admin operations as it releases the leases after every
+	// statement.
+	t.Run("session variable prevents blocking", func(t *testing.T) {
 		fooConn, err := fooDB.Conn(ctx)
 		require.NoError(t, err)
 		defer func() { _ = fooConn.Close() }()
-		_, err = fooConn.ExecContext(ctx, "SET allow_role_memberships_to_change_during_transaction = true;")
+
+		outerTx, err := fooConn.BeginTx(ctx, nil)
 		require.NoError(t, err)
-		fooTx, err := fooConn.BeginTx(ctx, nil)
+		defer func() { _ = outerTx.Rollback() }()
+
+		_, err = outerTx.ExecContext(ctx, "SET allow_role_memberships_to_change_during_transaction = true;")
 		require.NoError(t, err)
+
 		var count int
 		require.NoError(t,
-			fooTx.QueryRow("SELECT count(*) FROM t").Scan(&count))
+			outerTx.QueryRow("SELECT count(*) FROM t").Scan(&count))
 		require.Equal(t, 0, count)
-		conn, err := sqlDB.Conn(ctx)
-		require.NoError(t, err)
-		defer func() { _ = conn.Close() }()
-		// Set a timeout on the SQL operations to ensure that they both
-		// happen in a timely manner.
-		grantRevokeTimeout, cancel := context.WithTimeout(
-			ctx, testutils.DefaultSucceedsSoonDuration,
-		)
-		defer cancel()
-
-		_, err = conn.ExecContext(grantRevokeTimeout, "GRANT foo TO baz;")
-		require.NoError(t, err)
-		_, err = conn.ExecContext(grantRevokeTimeout, "REVOKE bar FROM baz;")
-		require.NoError(t, err)
 
-		// Ensure the transaction we held open commits without issue.
-		require.NoError(t, fooTx.Commit())
-	})
-
-	t.Run("session variable prevents waiting during CREATE and DROP role", func(t *testing.T) {
-		fooConn, err := fooDB.Conn(ctx)
-		require.NoError(t, err)
-		defer func() { _ = fooConn.Close() }()
-		_, err = fooConn.ExecContext(ctx, "SET allow_role_memberships_to_change_during_transaction = true;")
-		require.NoError(t, err)
-		fooTx, err := fooConn.BeginTx(ctx, nil)
-		require.NoError(t, err)
-		// We need to use show roles because that access the system.users table.
-		_, err = fooTx.Exec("SHOW ROLES")
+		// the first user operation shouldn't block
+		_, err = sqlDB.Exec("GRANT biz TO bar;")
 		require.NoError(t, err)
 
-		conn, err := sqlDB.Conn(ctx)
-		require.NoError(t, err)
-		defer func() { _ = conn.Close() }()
-		// Set a timeout on the SQL operations to ensure that they both
-		// happen in a timely manner.
-		grantRevokeTimeout, cancel := context.WithTimeout(
-			ctx, testutils.DefaultSucceedsSoonDuration,
-		)
-		defer cancel()
-
-		_, err = conn.ExecContext(grantRevokeTimeout, "CREATE ROLE new_role;")
-		require.NoError(t, err)
-		_, err = conn.ExecContext(grantRevokeTimeout, "DROP ROLE new_role;")
+		// nor should the second
+		_, err = sqlDB.Exec("GRANT baz TO bar;")
 		require.NoError(t, err)
 
-		// Ensure the transaction we held open commits without issue.
-		require.NoError(t, fooTx.Commit())
+		require.NoError(t, outerTx.Commit())
 	})
-
 }