Merge #153040

153040: sql: unblock user create from long-running transactions r=bghal a=bghal

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


Co-authored-by: Brendan Gerrity <[email protected]>

Author: craig[bot]

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

@@ -557,16 +557,7 @@ func TestWaitForNewVersion(testingT *testing.T) {
 	defer leaktest.AfterTest(testingT)()
 	defer log.Scope(testingT).Close(testingT)
 
-	skip.WithIssue(testingT, 152051)
-
 	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()
 
@@ -589,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/conn_executor_jobs.go

@@ -84,9 +84,7 @@ func (ex *connExecutor) waitForNewVersionPropagation(
 			continue
 		}
 
-		// TODO(sql-foundations): Change this back to WaitForNewVersion once we
-		// address the flaky behavior around privilege and role membership caching.
-		if _, err := ex.planner.LeaseMgr().WaitForOneVersion(ex.Ctx(), idVersion.ID, cachedRegions,
+		if _, err := ex.planner.LeaseMgr().WaitForNewVersion(ex.Ctx(), idVersion.ID, cachedRegions,
 			retry.Options{
 				InitialBackoff: time.Millisecond,
 				MaxBackoff:     time.Second,

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`,