storage: add MVCCCheckForAcquireLock and MVCCAcquireLock functions

[nvanbenschoten]

Fixes #109646.
Informs #100193.

This PR adds and implements two new MVCC functions: MVCCCheckForAcquireLock and MVCCAcquireLock. The former scans the replicated lock table to determine whether a lock acquisition is permitted. It will be used by unreplicated lock acquisition. The latter does the same, but then also writes the lock to the replicated lock table if permitted. It will be used by replicated lock acquisition.

MVCCStats handling is left as a TODO for after #109645 is addressed.

---

The two functions are built using a new abstraction, the lockTableKeyScanner.

The lockTableKeyScanner uses a LockTableIterator to scan a single key in the replicated lock table. It searches for locks on the key that conflict with a (transaction, strength) pair and for locks that the transaction has already acquired on the key.

The purpose of a lockTableKeyScanner is to determine whether a transaction can acquire a lock on a key or perform an MVCC mutation on a key, and if so, what lock table keys the transaction should write to perform the operation. It is used by this commit to implement the two new MVCC functions. In a future commit, it will also start to be used by mvccPutInternal, the kernel of all MVCC mutations.

Release note: None

[cockroach-teamcity]

This change is 

[sumeerbhola]

flushing some preliminary comments -- have only read lock_table_key_scanner.go.

Reviewed 1 of 7 files at r5.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @arulajmani and @nvanbenschoten)

---

pkg/storage/lock_table_key_scanner.go line 28 at r5 (raw file):

// be used to iterate supported lock strengths in strength order (strongest to
// weakest).
var replicatedLockStrengths = [...]lock.Strength{lock.Intent, lock.Exclusive, lock.Shared}

nit: should there be an assertion that there is no strength higher than lock.Intent so that we remember to update this if more strengths are added?

---

pkg/storage/lock_table_key_scanner.go line 79 at r5 (raw file):

type lockTableKeyScanner struct {
	iter *LockTableIterator
	// The transaction attempting to acquire a lock. The ID be zero if a

nit: will be zero

---

pkg/storage/lock_table_key_scanner.go line 169 at r5 (raw file):

// beforeScan resets the scanner's state before a scan.
func (s *lockTableKeyScanner) beforeScan() {

nit: it is better to call this resetScanState()

---

pkg/storage/lock_table_key_scanner.go line 225 at r5 (raw file):

		return false
	}
	return s.advance()

nit: this seems over-abstracted. Given there is only 1 caller of advance why not just plop the code here.

Similarly for beforeScan and afterScan -- its more work to jump around to read the actual code.

This is subjective of course, and I am not suggesting one giant function with a for loop, but at some point they become too tiny.
But I am also fine with whatever you decide.

---

pkg/storage/lock_table_key_scanner.go line 284 at r5 (raw file):

		ltValueCopy = new(enginepb.MVCCMetadata)
	}
	*ltValueCopy = *ltValue

This will not do a deep copy so they will both share the same *TxnMeta, yes?
The unmarshal code looks like:

if m.Txn == nil {  
 m.Txn = &TxnMeta{}  
}  
if err := m.Txn.Unmarshal(dAtA\[iNdEx:postIndex\]); err != nil {  
 return err  
}

So won't this incorrectly overwrite the shared *TxnMeta?

[arulajmani]

This looks pretty good, the layering here came out quite nice. Should be for me after this round of comments.

Reviewed 9 of 9 files at r1, 1 of 1 files at r2, 5 of 5 files at r3, 2 of 2 files at r4, 7 of 7 files at r5, 18 of 18 files at r6, 18 of 18 files at r7, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @nvanbenschoten and @sumeerbhola)

---

-- commits line 102 at r5:
"are built"

---

-- commits line 113 at r5:

the kernel

🍿

---

pkg/storage/lock_table_key_scanner.go line 35 at r5 (raw file):

// Trying to use a lock strength that isn't supported with replicated locks to
// index into the lockTableScanner.ownLocks array will cause a runtime error.
var replicatedLockStrengthToIndexMap = func() [lock.MaxStrength + 1]int {

Third usage it's a pattern ;)

---

pkg/storage/lock_table_key_scanner.go line 84 at r5 (raw file):

	// Stop adding conflicting locks and abort scan once the maxConflicts limit
	// is reached. Ignored if zero.
	maxConflicts int64

nit: I thought you were in the maxConflicts is the ... camp for such comments, no?

---

pkg/storage/lock_table_key_scanner.go line 99 at r5 (raw file):

// to avoid heap allocations. It is extracted into a separate struct to avoid
// being cleared when lockTableKeyScanner is recycled.

Could you explain this bit, for my edification?

---

pkg/storage/lock_table_key_scanner.go line 229 at r5 (raw file):

// getLockTableKey decodes the current lock table key.
func (s *lockTableKeyScanner) getLockTableKey() (LockTableKey, bool) {

nit: consider adding comments about what the boolean value means to the caller (here, and elsewhere).

---

pkg/storage/lock_table_key_scanner.go line 273 at r5 (raw file):

// consumeOwnLock consumes a lock held by the scanning transaction.
func (s *lockTableKeyScanner) consumeOwnLock(

We're always returning true from this function -- could we avoid having a return value here altogether?

---

pkg/storage/lock_table_key_scanner.go line 284 at r5 (raw file):

Previously, sumeerbhola wrote…

This will not do a deep copy so they will both share the same *TxnMeta, yes?
The unmarshal code looks like:

if m.Txn == nil {  
 m.Txn = &TxnMeta{}  
}  
if err := m.Txn.Unmarshal(dAtA\[iNdEx:postIndex\]); err != nil {  
 return err  
}

So won't this incorrectly overwrite the shared *TxnMeta?

👀

---

pkg/storage/mvcc.go line 5315 at r5 (raw file):

// is returned. Otherwise, nil is returned. Unlike MVCCAcquireLock, this method
// does not actually acquire the lock (i.e. write to the lock table).
func MVCCCheckAcquireLock(

How do you feel about renaming this to either MVCCCheckForAcquireLock or MVCCCheckLockForAcquire, to make it clear that we're not actually acquiring the lock in the function name itself?

---

edit: MVCCCheckForAcquireLock it is!

---

pkg/storage/mvcc.go line 5316 at r5 (raw file):

// does not actually acquire the lock (i.e. write to the lock table).
func MVCCCheckAcquireLock(
	ctx context.Context,

nit: unused

---

pkg/storage/mvcc.go line 5395 at r5 (raw file):

			// This is both a performance optimization and a necessary check for
			// correctness. If we were to reacquire the lock at a newer sequence
			// number and clobber the existing locks with its older sequence

"locks" and "its" sounds off -- did you mean "existing lock with its older sequence number"?

---

pkg/storage/mvcc.go line 5414 at r5 (raw file):

	newMeta := &buf.newMeta
	newMeta.Txn = &txn.TxnMeta
	keyBytes, valBytes, err := buf.putLockMeta(rw, MakeMVCCMetadataKey(key), str, newMeta, rolledBack)

One difference re: ignored sequence numbers jumped out to me here, compared to unreplicated locks.

For unreplicated locks, when we're re-acquiring a lock, we roll back sequence numbers across all lock strenghts -- not just the one that it's being re-acquired with. However, for replicated locks, we're only doing so for the lock strength the lock is being acquired with.

Thought about this a bit, and I think the behaviour here is correct. Even though we're not releasing locks across all strengths that have been rolled back, newer lock acquisitions at lower lock strengths know which ones should be -- and you're correctly handling this, by checking if a higher lock strength has been rolled back, and only returning early if it hasn't.

It might be worth capturing this difference between replicated and unreplicated locks in a comment, mostly as a proof for future readers who find themselves in these ends.

---

pkg/storage/testdata/mvcc_histories/replicated_locks line 209 at r7 (raw file):

error: (*kvpb.LockConflictError:) conflicting locks on "k1"

# Intents are treated similarly to Exclusive locks.

Does the following case work with intents:

• Put @ key A, seq 1
• Put @ key A, seq 2
• Shared @ Key A, seq 3, ignored sequence numbers [2]

I'd expect us to not acquire a shared lock at sequence number 3, as we've already got an intent at sequence number 1. I may be missing something, but I don't think the rolledBack logic takes the intent history into account though. Should it?

[nvanbenschoten]

TFTRs!

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @arulajmani and @sumeerbhola)

---

-- commits line 102 at r5:

Previously, arulajmani (Arul Ajmani) wrote…

"are built"

Done.

---

pkg/storage/lock_table_key_scanner.go line 28 at r5 (raw file):

Previously, sumeerbhola wrote…

nit: should there be an assertion that there is no strength higher than lock.Intent so that we remember to update this if more strengths are added?

Done.

---

pkg/storage/lock_table_key_scanner.go line 35 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

Third usage it's a pattern ;)

There's probably some room to clean these up and unify them in the lock package.

---

pkg/storage/lock_table_key_scanner.go line 79 at r5 (raw file):

Previously, sumeerbhola wrote…

nit: will be zero

Done.

---

pkg/storage/lock_table_key_scanner.go line 84 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

nit: I thought you were in the maxConflicts is the ... camp for such comments, no?

For exported fields, I am. For explanations on internal fields, it feels a bit verbose. Though I'm sure I've contradicted that plenty of times.

---

pkg/storage/lock_table_key_scanner.go line 99 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

// to avoid heap allocations. It is extracted into a separate struct to avoid
// being cleared when lockTableKeyScanner is recycled.

Could you explain this bit, for my edification?

I was imagining that there would be more than one field here that we wanted to recycle, but it doesn't look like there is. In fact, two of the three we specifically don't want to recycle. So I've removed this.

---

pkg/storage/lock_table_key_scanner.go line 169 at r5 (raw file):

Previously, sumeerbhola wrote…

nit: it is better to call this resetScanState()

Done.

---

pkg/storage/lock_table_key_scanner.go line 225 at r5 (raw file):

Previously, sumeerbhola wrote…

nit: this seems over-abstracted. Given there is only 1 caller of advance why not just plop the code here.

Similarly for beforeScan and afterScan -- its more work to jump around to read the actual code.

This is subjective of course, and I am not suggesting one giant function with a for loop, but at some point they become too tiny.
But I am also fine with whatever you decide.

I think that's fair. There was more in these internal functions back before LockTableIterator extracted some of the complexity, so it made sense to split them out. I also tried to model this off of pebbleMVCCScanner.

---

pkg/storage/lock_table_key_scanner.go line 229 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

nit: consider adding comments about what the boolean value means to the caller (here, and elsewhere).

Done.

---

pkg/storage/lock_table_key_scanner.go line 273 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

We're always returning true from this function -- could we avoid having a return value here altogether?

We could, but the style here is to let each function independently determine whether the iteration can continue. I think it's appropriate to let consumeOwnLock decide whether it reached an invalid state or not.

---

pkg/storage/lock_table_key_scanner.go line 284 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

👀

They will both share the same *TxnMeta immediately after the copy, but this aliasing won't lead to incorrectly overwriting the shared *TxnMeta, because protoutil.Unmarshal resets the proto before unmarshalling into it.

This is subtle, so I added a comment to avoid confusion.

---

pkg/storage/mvcc.go line 5315 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

How do you feel about renaming this to either MVCCCheckForAcquireLock or MVCCCheckLockForAcquire, to make it clear that we're not actually acquiring the lock in the function name itself?

---

edit: MVCCCheckForAcquireLock it is!

Done.

---

pkg/storage/mvcc.go line 5316 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

nit: unused

Yeah, I'll keep it plumbed in to avoid churn if we want to log or trace in here.

---

pkg/storage/mvcc.go line 5395 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

"locks" and "its" sounds off -- did you mean "existing lock with its older sequence number"?

Done.

---

pkg/storage/mvcc.go line 5414 at r5 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

One difference re: ignored sequence numbers jumped out to me here, compared to unreplicated locks.

For unreplicated locks, when we're re-acquiring a lock, we roll back sequence numbers across all lock strenghts -- not just the one that it's being re-acquired with. However, for replicated locks, we're only doing so for the lock strength the lock is being acquired with.

Thought about this a bit, and I think the behaviour here is correct. Even though we're not releasing locks across all strengths that have been rolled back, newer lock acquisitions at lower lock strengths know which ones should be -- and you're correctly handling this, by checking if a higher lock strength has been rolled back, and only returning early if it hasn't.

It might be worth capturing this difference between replicated and unreplicated locks in a comment, mostly as a proof for future readers who find themselves in these ends.

Great point! Left a comment.

---

pkg/storage/testdata/mvcc_histories/replicated_locks line 209 at r7 (raw file):

Previously, arulajmani (Arul Ajmani) wrote…

Does the following case work with intents:

• Put @ key A, seq 1
• Put @ key A, seq 2
• Shared @ Key A, seq 3, ignored sequence numbers [2]

I'd expect us to not acquire a shared lock at sequence number 3, as we've already got an intent at sequence number 1. I may be missing something, but I don't think the rolledBack logic takes the intent history into account though. Should it?

That's a really interesting case. I added support and a test for it.

[sumeerbhola]

Reviewed 2 of 7 files at r8, 7 of 7 files at r9, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @arulajmani and @nvanbenschoten)

---

pkg/storage/mvcc.go line 5395 at r9 (raw file):

				// We don't currently do this for replicated locks to avoid
				// performing addition writes during the lock acquisition
				// process, but we could consider doing this in the future.

Can you add this comment to the lockTableImpl too, since that is where this logic is encoded.
We are acquiring a replicated lock in this method, so we will be doing a write already. Clearing the locks that have been rolled back seems a worthy thing to do.

Is the real issue that our current code structure does not make this easy in that we need to do a clear here and hook things up such that lockTableImpl is informed when the state machine is updated. And more importantly we must maintain the property that the replicated locks known to the lockTableImpl are a subset of the actual replicated locks, so we need to prevent a situation where a conflicting reader/writer concurrently finds the replicated lock in the persistent lock table and adds it to lockTableImpl even though it has been cleared. Presumably we are holding latches at str while doing MVCCAcquireLock, and the strength s we found here may be stronger than str, so there could be concurrent work that conflicts with s, but does not with str, so this race can happen.

---

pkg/storage/mvcc.go line 5415 at r9 (raw file):

		}

		if foundLock != nil && !rolledBack {

the if-block structuring seems odd. We are not updating foundLock so this conjunction could simply be if !rolledBack by placing it in the previous if-block. Also, we don't do anything for the case foundLock == nil so we could do something like the following and reduce the indentation.

for ... {
   ...
   if foundLock == nil {
     continue
   }
   // foundLock != nil
   ...
}

---

pkg/storage/mvcc.go line 5423 at r9 (raw file):

			// we would forget that the lock held at the older sequence number
			// had been and still should be held.
			log.VInfof(ctx, 2, "skipping lock acquisition for txn %s on key %s "+

would this logging be ever helpful in a production setting with high qps?
wouldn't a trace statement be more useful?

[arulajmani]

Reviewed 7 of 7 files at r8, 7 of 7 files at r9, all commit messages.
Reviewable status: complete! 2 of 0 LGTMs obtained (waiting on @nvanbenschoten and @sumeerbhola)

---

pkg/storage/lock_table_key_scanner.go line 273 at r5 (raw file):

Previously, nvanbenschoten (Nathan VanBenschoten) wrote…

We could, but the style here is to let each function independently determine whether the iteration can continue. I think it's appropriate to let consumeOwnLock decide whether it reached an invalid state or not.

Ack.

---

pkg/storage/testdata/mvcc_histories/replicated_locks line 209 at r7 (raw file):

Previously, nvanbenschoten (Nathan VanBenschoten) wrote…

That's a really interesting case. I added support and a test for it.

Nice!

[nvanbenschoten]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 2 stale) (waiting on @arulajmani, @nvanbenschoten, and @sumeerbhola)

---

pkg/storage/mvcc.go line 5395 at r9 (raw file):

Can you add this comment to the lockTableImpl too, since that is where this logic is encoded.

Done.

Presumably we are holding latches at str while doing MVCCAcquireLock, and the strength s we found here may be stronger than str, so there could be concurrent work that conflicts with s, but does not with str, so this race can happen.

I was not considering this case, but it's a great point. I'll add it to the comment.

---

pkg/storage/mvcc.go line 5415 at r9 (raw file):

Previously, sumeerbhola wrote…

the if-block structuring seems odd. We are not updating foundLock so this conjunction could simply be if !rolledBack by placing it in the previous if-block. Also, we don't do anything for the case foundLock == nil so we could do something like the following and reduce the indentation.

for ... {
   ...
   if foundLock == nil {
     continue
   }
   // foundLock != nil
   ...
}

Done.

---

pkg/storage/mvcc.go line 5423 at r9 (raw file):

Previously, sumeerbhola wrote…

would this logging be ever helpful in a production setting with high qps?
wouldn't a trace statement be more useful?

Switched to a log.VEventf(3) so that it will always be in traces and also the logs if a high enough verbosity is enabled.

[nvanbenschoten]

TFTRs!

bors r=sumeerbhola,arulajmani

[nvanbenschoten]

bors r-

[craig]

Canceled.

[nvanbenschoten]

bors r+

[craig]

Build succeeded:

• Bazel Essential CI (Cockroach)