reduce mut contention with a pessimistic lock

Author: aep

Makefile

@@ -10,7 +10,7 @@ apogy: api .PHONY
 
 
 test:
-	go test github.com/aep/apogy/server/...
+	go test -count=1  -v github.com/aep/apogy/server/...
 
 
 .PHONY:

kv/kv.go

@@ -13,6 +13,7 @@ type KeyAndValue struct {
 type KV interface {
 	Close()
 	Write() Write
+	ExclusiveWrite(ctx context.Context, keys ...[]byte) (Write, error)
 	Read() Read
 	Ping() error
 }

kv/kv_test.go

@@ -2,8 +2,11 @@ package kv
 
 import (
 	"context"
-	"github.com/stretchr/testify/require"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 )
 
 func TestKVDoesPreventOutdatedWrites(t *testing.T) {
@@ -47,3 +50,52 @@ func TestKVDoesNotPreventDup(t *testing.T) {
 	err = w1.Commit(context.Background())
 	require.NoError(t, err, "kv must allow setting the same key twice within a tx")
 }
+
+func TestKVLockReordering(t *testing.T) {
+	k, err := NewTikv()
+	require.NoError(t, err)
+	defer k.Close()
+
+	var ctx = t.Context()
+	var key = []byte("reordertest")
+
+	w1, err := k.ExclusiveWrite(ctx, key)
+	require.NoError(t, err)
+	defer w1.Rollback()
+
+	var w1commited = atomic.NewBool(false)
+	go func() {
+		time.Sleep(time.Millisecond * 101)
+
+		err = w1.Put(key, []byte("1"))
+		require.NoError(t, err)
+
+		err = w1.Commit(ctx)
+		w1commited.Store(true)
+		require.NoError(t, err, "first commit must work")
+	}()
+
+	w2, err := k.ExclusiveWrite(ctx, key)
+	require.NoError(t, err)
+	defer w2.Rollback()
+
+	if !w1commited.Load() {
+		panic("not reordered")
+	}
+
+	current, err := w2.Get(ctx, key)
+	require.NoError(t, err)
+	require.Equal(t, "1", string(current))
+
+	require.NoError(t, err)
+	err = w2.Put(key, []byte("2"))
+	require.NoError(t, err)
+
+	err = w2.Commit(ctx)
+	require.NoError(t, err, "second commit must work")
+
+	actual, err := k.Read().Get(ctx, key)
+	require.NoError(t, err)
+	require.Equal(t, "2", string(actual))
+
+}

kv/tikv.go

@@ -2,14 +2,18 @@ package kv
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"iter"
 	"log/slog"
 	"os"
+	"time"
 
 	pingcaplog "github.com/pingcap/log"
 
 	"github.com/lmittmann/tint"
+	tikverr "github.com/tikv/client-go/v2/error"
+	"github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/txnkv"
 	"github.com/tikv/client-go/v2/txnkv/txnsnapshot"
 
@@ -45,6 +49,12 @@ type TikvWrite struct {
 	txn      *txnkv.KVTxn
 	err      error
 	commited bool
+
+	statLockRetries int
+}
+
+func (w *TikvWrite) Stat() (statLockRetries int) {
+	return w.statLockRetries
 }
 
 func (w *TikvWrite) Commit(ctx context.Context) error {
@@ -249,8 +259,66 @@ func (t *Tikv) Close() {
 
 func (t *Tikv) Write() Write {
 	txn, err := t.k.Begin()
-	txn.SetEnable1PC(true)
-	return &TikvWrite{txn, err, false}
+	//txn.SetEnable1PC(true)
+	return &TikvWrite{txn: txn, err: err}
+}
+
+func (t *Tikv) ExclusiveWrite(ctx context.Context, keys ...[]byte) (Write, error) {
+
+	var waitMs = int64(100)
+
+	// DO NOT use aggressive locking.
+	// it's a deadlock trap
+	// r.txn.StartAggressiveLocking()
+	// instead do the retry loop in the client (us)
+
+	txn, err := t.k.Begin()
+	if err != nil {
+		return nil, err
+	}
+	txn.SetPessimistic(true)
+
+	retries := 0
+	for {
+
+		waitMs += 1
+		retries += 1
+
+		lkctx := kv.NewLockCtx(txn.StartTS(), waitMs, time.Now())
+
+		err = txn.LockKeys(ctx, lkctx, keys...)
+		if err == nil {
+			break
+		}
+
+		// we got the lock but someone changed the key we're locking
+		// get a new txn with the current start time
+		if tikverr.IsErrWriteConflict(err) {
+			txn.Rollback()
+			txn, err = t.k.Begin()
+			if err != nil {
+				return nil, err
+			}
+			txn.SetPessimistic(true)
+			continue
+		}
+
+		if !errors.Is(err, tikverr.ErrLockWaitTimeout) {
+			return nil, err
+		}
+
+		select {
+		case <-ctx.Done():
+			if ctx.Err() == nil {
+				return nil, err
+			}
+			return nil, ctx.Err()
+		default:
+			continue
+		}
+	}
+
+	return &TikvWrite{txn: txn, statLockRetries: retries}, nil
 }
 
 func (t *Tikv) Read() Read {

server/crud.go

@@ -3,12 +3,12 @@ package server
 import (
 	"context"
 	"fmt"
-	"log/slog"
 	"net/http"
 	"reflect"
 	"time"
 
 	openapi "github.com/aep/apogy/api/go"
+	"github.com/aep/apogy/kv"
 	"github.com/labstack/echo/v4"
 	tikerr "github.com/tikv/client-go/v2/error"
 	"go.opentelemetry.io/otel/attribute"
@@ -39,13 +39,13 @@ func (s *server) PutDocument(c echo.Context) error {
 
 	switch doc.Model {
 	case "Model":
-		err := s.validateSchemaSchema(ctx, doc)
+		err = s.validateSchemaSchema(ctx, doc)
 		if err != nil {
 			span.RecordError(err)
 			return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("validation error: %s", err))
 		}
 	case "Reactor":
-		err := s.validateReactorSchema(ctx, doc)
+		err = s.validateReactorSchema(ctx, doc)
 		if err != nil {
 			span.RecordError(err)
 			return echo.NewHTTPError(http.StatusBadRequest, fmt.Sprintf("validation error: %s", err))
@@ -60,65 +60,32 @@ func (s *server) PutDocument(c echo.Context) error {
 		return echo.NewHTTPError(http.StatusBadRequest, err.Error())
 	}
 
-	// Try to save document with retry logic
-	attempts := 0
-	for i := 0; ; i++ {
-		attempts++
-		retrySpan := trace.SpanFromContext(ctx)
-		retrySpan.SetAttributes(attribute.Int("attempt", i+1))
-
-		err := s.putDocument1(c, ctx, doc, model)
-		if err == nil {
-			retrySpan.SetAttributes(attribute.Bool("success", true))
-			break
-		}
-
-		retrySpan.RecordError(err)
-
-		if echoErr, ok := err.(*echo.HTTPError); ok {
-			// Record failed commit
-			if echoErr.Code == http.StatusInternalServerError {
-				kvCommitFailures.WithLabelValues("put_document", "internal_error").Inc()
-			} else if echoErr.Code == http.StatusConflict {
-				kvCommitFailures.WithLabelValues("put_document", "write_conflict").Inc()
-			} else {
-				kvCommitFailures.WithLabelValues("put_document", "other").Inc()
-			}
+	path, err := safeDBPath(doc.Model, doc.Id)
+	if err != nil {
+		return echo.NewHTTPError(http.StatusBadRequest, err.Error())
+	}
 
-			// Record retries for failed operation
-			kvCommitRetries.WithLabelValues("put_document", "failed").Observe(float64(attempts - 1))
-			return echoErr
-		}
+	// for contested keys, use pessimistic locking
+	var hotKeys = [][]byte{}
+	if doc.Mut != nil {
+		hotKeys = append(hotKeys, path)
+	}
 
-		if i > 10 {
-			time.Sleep(100 * time.Millisecond)
-		} else {
-			time.Sleep(10 * time.Millisecond)
+	var w2 kv.Write
+	if len(hotKeys) == 0 {
+		w2 = s.kv.Write()
+	} else {
+		w2, err = s.kv.ExclusiveWrite(ctx, hotKeys...)
+		if err != nil {
+			return err
 		}
 
-		slog.Warn("putDocument1", "err", err)
+		statLockRetries := w2.(*kv.TikvWrite).Stat()
+		span.SetAttributes(attribute.Int("lockRetries", statLockRetries))
+		kvLockRetries.WithLabelValues("hot", "true").Observe(float64(statLockRetries))
 	}
-
-	// Record successful commit metrics
-	span.SetAttributes(attribute.Int("attempts", attempts))
-	kvCommitRetries.WithLabelValues("put_document", "success").Observe(float64(attempts - 1))
-
-	return c.JSON(http.StatusOK, doc)
-}
-
-func (s *server) putDocument1(c echo.Context, ctx context.Context, doc_ *openapi.Document, model *Model) error {
-
-	doccpy := *doc_
-	doc := &doccpy
-
-	w2 := s.kv.Write()
 	defer w2.Close()
 
-	path, err := safeDBPath(doc.Model, doc.Id)
-	if err != nil {
-		return echo.NewHTTPError(http.StatusBadRequest, err.Error())
-	}
-
 	now := time.Now()
 	doc.History = &openapi.History{
 		Created: &now,
@@ -149,7 +116,7 @@ func (s *server) putDocument1(c echo.Context, ctx context.Context, doc_ *openapi
 			}
 
 			if reflect.DeepEqual(old.Val, doc.Val) {
-				return nil
+				return c.JSON(http.StatusOK, doc)
 			}
 
 			if old.Version != nil && doc.Version != nil {
@@ -255,7 +222,7 @@ func (s *server) putDocument1(c echo.Context, ctx context.Context, doc_ *openapi
 		return echo.NewHTTPError(http.StatusUnprocessableEntity, err.Error())
 	}
 
-	return nil
+	return c.JSON(http.StatusOK, doc)
 }
 
 func (s *server) GetDocument(c echo.Context, model string, id string) error {

server/health.go

@@ -42,9 +42,9 @@ var (
 		[]string{"operation"},
 	)
 
-	kvCommitRetries = prometheus.NewHistogramVec(
+	kvLockRetries = prometheus.NewHistogramVec(
 		prometheus.HistogramOpts{
-			Name:    "kv_commit_retries",
+			Name:    "kv_lock_retries",
 			Help:    "Number of commit retries",
 			Buckets: []float64{0, 1, 2, 3, 4, 5, 10, 20, 40, 80, 160, 320},
 		},
@@ -75,7 +75,7 @@ func init() {
 	promRegistry.MustRegister(httpRequestsTotal)
 	promRegistry.MustRegister(httpRequestDuration)
 	promRegistry.MustRegister(kvCommitDuration)
-	promRegistry.MustRegister(kvCommitRetries)
+	promRegistry.MustRegister(kvLockRetries)
 	promRegistry.MustRegister(kvCommitFailures)
 }