concurrency,kvserver: limited scans do optimistic latching

Latches for the full spans get inserted up front in the
spanlatch.Manager, and the conflict checking happens after
evaluation, only over the spans that were actually read.
If there is a conflict, the existing inserted latches are waited
on, and execution switches to pessimistic latching and locking.
The existing cluster setting for optimistic locking is used to
gate this behavior.

Numbers for the new OptimisticEval benchmark show improvement
when real-contention=false, compared to master which only had
optimistic locking. There is a modest slowdown in the
real-contention=true case since every optimistic read has to
try twice. The benchmark has concurrent writes of two different
kinds: latches represents 1PC writers, so no intents, while
latches-and-locks represent writers that will leave intents.

For the latter, when the optimistic read finds an intent during
evaluation it cannot necessarily add it as a discovered lock,
if there is also a conflicting latch (since it could be racing
with intent resolution). This case is rare in these benchmarks
(latches-and-locks/real-contention=true): 13% found an intent
and had a conflicting latch when num-writers=1 and < 1% for
the same case when num-writers=4. The remainder (the common
case) was to find a conflict when looking for conflicting
latches/locks using CheckOptimisticNoConflicts.

name                                                                     old time/op    new time/op    delta
OptimisticEvalForLocks/real-contention=false-16                            28.2µs ± 4%    28.5µs ± 5%     ~     (p=0.393 n=10+10)
OptimisticEvalForLocks/real-contention=true-16                             5.52ms ± 1%    5.52ms ± 1%     ~     (p=0.796 n=10+10)
OptimisticEval/latches-and-locks/real-contention=false/num-writers=1-16     220µs ±62%      89µs ±38%  -59.37%  (p=0.000 n=10+8)
OptimisticEval/latches-and-locks/real-contention=false/num-writers=4-16     213µs ±65%     155µs ±82%  -27.33%  (p=0.015 n=10+10)
OptimisticEval/latches-and-locks/real-contention=true/num-writers=1-16     1.33ms ± 3%    1.27ms ±16%     ~     (p=0.829 n=8+10)
OptimisticEval/latches-and-locks/real-contention=true/num-writers=4-16     2.02ms ±10%    2.25ms ± 9%  +11.31%  (p=0.000 n=10+10)
OptimisticEval/latches/real-contention=false/num-writers=1-16               175µs ± 2%      45µs ± 5%  -74.05%  (p=0.000 n=10+10)
OptimisticEval/latches/real-contention=false/num-writers=4-16               613µs ± 2%      44µs ± 3%  -92.74%  (p=0.000 n=10+9)
OptimisticEval/latches/real-contention=true/num-writers=1-16                181µs ± 4%     179µs ± 5%     ~     (p=0.315 n=10+10)
OptimisticEval/latches/real-contention=true/num-writers=4-16                624µs ± 3%     620µs ± 3%     ~     (p=0.247 n=10+10)

name                                                                     old alloc/op   new alloc/op   delta
OptimisticEvalForLocks/real-contention=false-16                            8.40kB ± 7%    8.33kB ± 3%     ~     (p=1.000 n=10+8)
OptimisticEvalForLocks/real-contention=true-16                             31.8kB ± 7%    32.6kB ± 9%     ~     (p=0.382 n=8+8)
OptimisticEval/latches-and-locks/real-contention=false/num-writers=1-16    23.9kB ±21%    17.8kB ±25%  -25.55%  (p=0.003 n=10+10)
OptimisticEval/latches-and-locks/real-contention=false/num-writers=4-16    24.1kB ±20%    19.4kB ±22%  -19.56%  (p=0.015 n=10+10)
OptimisticEval/latches-and-locks/real-contention=true/num-writers=1-16      104kB ± 1%     101kB ± 1%   -2.89%  (p=0.000 n=8+9)
OptimisticEval/latches-and-locks/real-contention=true/num-writers=4-16      197kB ± 6%     217kB ±11%  +10.19%  (p=0.000 n=9+10)
OptimisticEval/latches/real-contention=false/num-writers=1-16              29.9kB ± 0%    13.4kB ± 1%  -55.07%  (p=0.000 n=9+10)
OptimisticEval/latches/real-contention=false/num-writers=4-16              94.9kB ± 0%    14.5kB ± 1%  -84.74%  (p=0.000 n=9+8)
OptimisticEval/latches/real-contention=true/num-writers=1-16               29.9kB ± 0%    31.3kB ± 0%   +4.59%  (p=0.000 n=10+10)
OptimisticEval/latches/real-contention=true/num-writers=4-16               94.8kB ± 0%    96.2kB ± 0%   +1.48%  (p=0.000 n=10+10)

name                                                                     old allocs/op  new allocs/op  delta
OptimisticEvalForLocks/real-contention=false-16                              68.6 ± 1%      69.6 ± 2%   +1.52%  (p=0.005 n=9+10)
OptimisticEvalForLocks/real-contention=true-16                                271 ± 2%       272 ± 2%     ~     (p=0.336 n=8+8)
OptimisticEval/latches-and-locks/real-contention=false/num-writers=1-16       149 ±17%       117 ±18%  -21.00%  (p=0.002 n=10+10)
OptimisticEval/latches-and-locks/real-contention=false/num-writers=4-16       151 ±16%       126 ±18%  -16.31%  (p=0.013 n=10+10)
OptimisticEval/latches-and-locks/real-contention=true/num-writers=1-16        610 ± 0%       559 ± 1%   -8.32%  (p=0.000 n=8+9)
OptimisticEval/latches-and-locks/real-contention=true/num-writers=4-16      1.12k ± 5%     1.19k ±12%     ~     (p=0.053 n=9+10)
OptimisticEval/latches/real-contention=false/num-writers=1-16                 162 ± 0%        91 ± 0%  -43.83%  (p=0.001 n=8+9)
OptimisticEval/latches/real-contention=false/num-writers=4-16                 445 ± 0%        96 ± 0%  -78.44%  (p=0.000 n=9+9)
OptimisticEval/latches/real-contention=true/num-writers=1-16                  163 ± 0%       184 ± 0%  +13.07%  (p=0.000 n=10+10)
OptimisticEval/latches/real-contention=true/num-writers=4-16                  445 ± 0%       467 ± 0%   +4.94%  (p=0.000 n=10+10)

Fixes #9521

Release note (performance improvement): A limited scan now checks
for conflicting latches in an optimistic manner, which means it will
not conflict with latches that were held in the scan's full spans,
but were not in the spans that were scanned until the limit was
reached. This behavior can be turned off (along with optimistic
locking) by changing the value of the cluster setting
kv.concurrency.optimistic_eval_limited_scans.enabled to false.

pkg/kv/kvserver/client_replica_test.go

@@ -3904,7 +3904,71 @@ func TestOptimisticEvalNoContention(t *testing.T) {
 	require.NoError(t, txn1.Commit(ctx))
 }
 
-func BenchmarkOptimisticEval(b *testing.B) {
+// TestOptimisticEvalWithConcurrentWriters tests concurrently running writes
+// and optimistic reads where the latter always conflict. This is just a
+// sanity check to confirm that nothing fails.
+func TestOptimisticEvalWithConcurrentWriters(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s, db := setupDBAndWriteAAndB(t)
+	defer s.Stopper().Stop(ctx)
+
+	finish := make(chan struct{})
+	var workers sync.WaitGroup
+	for i := 0; i < 4; i++ {
+		workers.Add(1)
+		go func() {
+			for {
+				require.NoError(t, db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) (err error) {
+					if err := txn.Put(ctx, "a", "a"); err != nil {
+						return err
+					}
+					return txn.Commit(ctx)
+				}))
+				select {
+				case _, recv := <-finish:
+					if !recv {
+						workers.Done()
+						return
+					}
+				default:
+				}
+			}
+		}()
+	}
+	for i := 0; i < 4; i++ {
+		workers.Add(1)
+		go func() {
+			for {
+				require.NoError(t, db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) (err error) {
+					_, err = txn.Scan(ctx, "a", "c", 1)
+					if err != nil {
+						return err
+					}
+					err = txn.Commit(ctx)
+					return err
+				}))
+				select {
+				case _, recv := <-finish:
+					if !recv {
+						workers.Done()
+						return
+					}
+				default:
+				}
+			}
+		}()
+	}
+	time.Sleep(10 * time.Second)
+	close(finish)
+	workers.Wait()
+}
+
+// BenchmarkOptimisticEvalForLocks benchmarks optimistic evaluation when the
+// potentially conflicting lock is explicitly held for a duration of time.
+func BenchmarkOptimisticEvalForLocks(b *testing.B) {
 	defer log.Scope(b).Close(b)
 	ctx := context.Background()
 	args := base.TestServerArgs{}
@@ -3980,3 +4044,102 @@ func BenchmarkOptimisticEval(b *testing.B) {
 			})
 	}
 }
+
+// BenchmarkOptimisticEval benchmarks optimistic evaluation with
+// - potentially conflicting latches held by 1PC transactions doing writes.
+// - potentially conflicting latches or locks held by transactions doing
+//   writes.
+func BenchmarkOptimisticEval(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ctx := context.Background()
+	args := base.TestServerArgs{}
+
+	for _, latches := range []bool{false, true} {
+		conflictWith := "latches-and-locks"
+		if latches {
+			conflictWith = "latches"
+		}
+		b.Run(conflictWith, func(b *testing.B) {
+			for _, realContention := range []bool{false, true} {
+				b.Run(fmt.Sprintf("real-contention=%t", realContention), func(b *testing.B) {
+					for _, numWriters := range []int{1, 4} {
+						b.Run(fmt.Sprintf("num-writers=%d", numWriters), func(b *testing.B) {
+							// Since we are doing writes in the benchmark, start with a
+							// fresh server each time so that we start with a fresh engine
+							// without many versions for a key.
+							s, _, db := serverutils.StartServer(b, args)
+							defer s.Stopper().Stop(ctx)
+
+							require.NoError(b, db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) (err error) {
+								if err := txn.Put(ctx, "a", "a"); err != nil {
+									return err
+								}
+								if err := txn.Put(ctx, "b", "b"); err != nil {
+									return err
+								}
+								return txn.Commit(ctx)
+							}))
+							tup, err := db.Get(ctx, "a")
+							require.NoError(b, err)
+							require.NotNil(b, tup.Value)
+							tup, err = db.Get(ctx, "b")
+							require.NoError(b, err)
+							require.NotNil(b, tup.Value)
+
+							writeKey := "b"
+							if realContention {
+								writeKey = "a"
+							}
+							finishWrites := make(chan struct{})
+							var writers sync.WaitGroup
+							for i := 0; i < numWriters; i++ {
+								writers.Add(1)
+								go func() {
+									for {
+										if latches {
+											require.NoError(b, db.Put(ctx, writeKey, "foo"))
+
+										} else {
+											require.NoError(b, db.Txn(ctx,
+												func(ctx context.Context, txn *kv.Txn) (err error) {
+													if err := txn.Put(ctx, writeKey, "foo"); err != nil {
+														return err
+													}
+													return txn.Commit(ctx)
+												}))
+										}
+										select {
+										case _, recv := <-finishWrites:
+											if !recv {
+												writers.Done()
+												return
+											}
+										default:
+										}
+									}
+								}()
+							}
+							b.ResetTimer()
+							for i := 0; i < b.N; i++ {
+								_ = db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) (err error) {
+									_, err = txn.Scan(ctx, "a", "c", 1)
+									if err != nil {
+										panic(err)
+									}
+									err = txn.Commit(ctx)
+									if err != nil {
+										panic(err)
+									}
+									return err
+								})
+							}
+							b.StopTimer()
+							close(finishWrites)
+							writers.Wait()
+						})
+					}
+				})
+			}
+		})
+	}
+}

pkg/kv/kvserver/concurrency/concurrency_control.go

@@ -394,6 +394,7 @@ type Request struct {
 type Guard struct {
 	Req Request
 	lg  latchGuard
+	lm  latchManager
 	ltg lockTableGuard
 	// The latest RequestEvalKind passed to SequenceReq.
 	EvalKind RequestEvalKind
@@ -411,13 +412,30 @@ type Error = roachpb.Error
 // Internal Structure Interfaces //
 ///////////////////////////////////
 
-// latchManager serializes access to keys and key ranges.
+// latchManager serializes access to keys and key ranges. The
+// {AcquireOptimistic,CheckOptimisticNoConflicts,WaitUntilAcquired} methods
+// are only for use in optimistic latching.
 //
 // See additional documentation in pkg/storage/spanlatch.
 type latchManager interface {
 	// Acquires latches, providing mutual exclusion for conflicting requests.
 	Acquire(context.Context, Request) (latchGuard, *Error)
 
+	// AcquireOptimistic is like Acquire in that it inserts latches, but it does
+	// not wait for conflicting latches on overlapping spans to be released
+	// before returning. This should be followed by CheckOptimisticNoConflicts
+	// to validate that not waiting did not violate correctness.
+	AcquireOptimistic(req Request) latchGuard
+
+	// CheckOptimisticNoConflicts returns true iff the spans in the provided
+	// spanset do not conflict with existing latches.
+	CheckOptimisticNoConflicts(lg latchGuard, spans *spanset.SpanSet) bool
+
+	// WaitUntilAcquired is meant to be called when CheckOptimisticNoConflicts
+	// returned false, or some other occurrence (like conflicting locks) is
+	// causing this request to switch to pessimistic latching.
+	WaitUntilAcquired(ctx context.Context, lg latchGuard) (latchGuard, *Error)
+
 	// Releases latches, relinquish its protection from conflicting requests.
 	Release(latchGuard)
 

pkg/kv/kvserver/concurrency/concurrency_manager.go

@@ -189,25 +189,43 @@ func (m *managerImpl) sequenceReqWithGuard(
 
 	// Only the first iteration can sometimes already be holding latches -- we
 	// use this to assert below.
-	first := true
+	firstIteration := true
 	for {
-		if !first {
-			g.AssertNoLatches()
-		}
-		first = false
 		if !g.HoldingLatches() {
-			// TODO(sumeer): optimistic requests could register their need for
-			// latches, but not actually wait until acquisition.
-			// https://github.com/cockroachdb/cockroach/issues/9521
-
-			// Acquire latches for the request. This synchronizes the request
-			// with all conflicting in-flight requests.
-			log.Event(ctx, "acquiring latches")
-			g.lg, err = m.lm.Acquire(ctx, req)
+			if g.EvalKind == OptimisticEval {
+				if !firstIteration {
+					// The only way we loop more than once is when conflicting locks are
+					// found -- see below where that happens and the comment there on
+					// why it will never happen with OptimisticEval.
+					panic("optimistic eval should not loop in sequenceReqWithGuard")
+				}
+				log.Event(ctx, "optimistically acquiring latches")
+				g.lg = m.lm.AcquireOptimistic(req)
+				g.lm = m.lm
+			} else {
+				// Acquire latches for the request. This synchronizes the request
+				// with all conflicting in-flight requests.
+				log.Event(ctx, "acquiring latches")
+				g.lg, err = m.lm.Acquire(ctx, req)
+				if err != nil {
+					return nil, err
+				}
+				g.lm = m.lm
+			}
+		} else {
+			if !firstIteration {
+				panic(errors.AssertionFailedf("second or later iteration cannot be holding latches"))
+			}
+			if g.EvalKind != PessimisticAfterFailedOptimisticEval {
+				panic("must not be holding latches")
+			}
+			log.Event(ctx, "optimistic failed, so waiting for latches")
+			g.lg, err = m.lm.WaitUntilAcquired(ctx, g.lg)
 			if err != nil {
 				return nil, err
 			}
 		}
+		firstIteration = false
 
 		// Some requests don't want the wait on locks.
 		if req.LockSpans.Empty() {
@@ -226,7 +244,9 @@ func (m *managerImpl) sequenceReqWithGuard(
 			g.ltg = m.lt.ScanAndEnqueue(g.Req, g.ltg)
 		}
 
-		// Wait on conflicting locks, if necessary.
+		// Wait on conflicting locks, if necessary. Note that this will never be
+		// true if ScanOptimistic was called above. Therefore it will also never
+		// be true if latchManager.AcquireOptimistic was called.
 		if g.ltg.ShouldWait() {
 			m.lm.Release(g.moveLatchGuard())
 
@@ -547,21 +567,44 @@ func (g *Guard) AssertNoLatches() {
 	}
 }
 
-// CheckOptimisticNoConflicts checks that the lockSpansRead do not have a
-// conflicting lock.
-func (g *Guard) CheckOptimisticNoConflicts(lockSpansRead *spanset.SpanSet) (ok bool) {
+// CheckOptimisticNoConflicts checks that the {latch,lock}SpansRead do not
+// have a conflicting latch, lock.
+func (g *Guard) CheckOptimisticNoConflicts(
+	latchSpansRead *spanset.SpanSet, lockSpansRead *spanset.SpanSet,
+) (ok bool) {
 	if g.EvalKind != OptimisticEval {
 		panic(errors.AssertionFailedf("unexpected EvalKind: %d", g.EvalKind))
 	}
-	if g.ltg == nil {
+	if g.lg == nil && g.ltg == nil {
+		return true
+	}
+	if g.lg == nil {
+		panic("expected non-nil latchGuard")
+	}
+	// First check the latches, since a conflict there could mean that racing
+	// requests in the lock table caused a conflicting lock to not be noticed.
+	if g.lm.CheckOptimisticNoConflicts(g.lg, latchSpansRead) {
+		return g.ltg.CheckOptimisticNoConflicts(lockSpansRead)
+	}
+	return false
+}
+
+// CheckOptimisticNoLatchConflicts checks that the declared latch spans for
+// the request do not have a conflicting latch.
+func (g *Guard) CheckOptimisticNoLatchConflicts() (ok bool) {
+	if g.EvalKind != OptimisticEval {
+		panic(errors.AssertionFailedf("unexpected EvalKind: %d", g.EvalKind))
+	}
+	if g.lg == nil {
 		return true
 	}
-	return g.ltg.CheckOptimisticNoConflicts(lockSpansRead)
+	return g.lm.CheckOptimisticNoConflicts(g.lg, g.Req.LatchSpans)
 }
 
 func (g *Guard) moveLatchGuard() latchGuard {
 	lg := g.lg
 	g.lg = nil
+	g.lm = nil
 	return lg
 }
 

pkg/kv/kvserver/concurrency/concurrency_manager_test.go

@@ -300,8 +300,8 @@ func TestConcurrencyManagerBasic(t *testing.T) {
 					d.Fatalf(t, "unknown request: %s", reqName)
 				}
 				reqs, _ := scanRequests(t, d, c)
-				_, lockSpans := c.collectSpans(t, g.Req.Txn, g.Req.Timestamp, reqs)
-				return fmt.Sprintf("no-conflicts: %t", g.CheckOptimisticNoConflicts(lockSpans))
+				latchSpans, lockSpans := c.collectSpans(t, g.Req.Txn, g.Req.Timestamp, reqs)
+				return fmt.Sprintf("no-conflicts: %t", g.CheckOptimisticNoConflicts(latchSpans, lockSpans))
 
 			case "on-lock-acquired":
 				var reqName string

pkg/kv/kvserver/concurrency/latch_manager.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanlatch"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 )
 
@@ -31,6 +32,25 @@ func (m *latchManagerImpl) Acquire(ctx context.Context, req Request) (latchGuard
 	return lg, nil
 }
 
+func (m *latchManagerImpl) AcquireOptimistic(req Request) latchGuard {
+	lg := m.m.AcquireOptimistic(req.LatchSpans)
+	return lg
+}
+
+func (m *latchManagerImpl) CheckOptimisticNoConflicts(lg latchGuard, spans *spanset.SpanSet) bool {
+	return m.m.CheckOptimisticNoConflicts(lg.(*spanlatch.Guard), spans)
+}
+
+func (m *latchManagerImpl) WaitUntilAcquired(
+	ctx context.Context, lg latchGuard,
+) (latchGuard, *Error) {
+	lg, err := m.m.WaitUntilAcquired(ctx, lg.(*spanlatch.Guard))
+	if err != nil {
+		return nil, roachpb.NewError(err)
+	}
+	return lg, nil
+}
+
 func (m *latchManagerImpl) Release(lg latchGuard) {
 	m.m.Release(lg.(*spanlatch.Guard))
 }