server: Add verification of whether lock was called within out outside of apply

scripts/test.sh

@@ -41,6 +41,7 @@ set -o pipefail
 # e.g. add/update missing dependencies. Such divergences should be 
 # detected and trigger a failure that needs explicit developer's action.
 export GOFLAGS=-mod=readonly
+export ETCD_VERIFY=all
 
 source ./scripts/test_lib.sh
 source ./scripts/build.sh

server/storage/backend/batch_tx.go

@@ -53,6 +53,8 @@ type BatchTx interface {
 	Commit()
 	// CommitAndStop commits the previous tx and does not create a new one.
 	CommitAndStop()
+	LockInsideApply()
+	LockOutsideApply()
 }
 
 type batchTx struct {
@@ -67,6 +69,16 @@ func (t *batchTx) Lock() {
 	t.Mutex.Lock()
 }
 
+func (t *batchTx) LockInsideApply() {
+	ValidateCalledInsideApply(t.backend.lg)
+	t.Lock()
+}
+
+func (t *batchTx) LockOutsideApply() {
+	ValidateCalledOutSideApply(t.backend.lg)
+	t.Lock()
+}
+
 func (t *batchTx) Unlock() {
 	if t.pending >= t.backend.batchLimit {
 		t.commit(false)

server/storage/backend/verify.go

@@ -0,0 +1,56 @@
+// Copyright 2022 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package backend
+
+import (
+	"os"
+	"runtime/debug"
+	"strings"
+
+	"go.uber.org/zap"
+)
+
+const (
+	ENV_VERIFY           = "ETCD_VERIFY"
+	ENV_VERIFY_ALL_VALUE = "all"
+	ENV_VERIFY_LOCK      = "lock"
+)
+
+func ValidateCalledInsideApply(lg *zap.Logger) {
+	if !verifyLockEnabled() {
+		return
+	}
+	if !insideApply() {
+		lg.Panic("Called outside of APPLY!", zap.Stack("stacktrace"))
+	}
+}
+
+func ValidateCalledOutSideApply(lg *zap.Logger) {
+	if !verifyLockEnabled() {
+		return
+	}
+	if insideApply() {
+		lg.Panic("Called inside of APPLY!", zap.Stack("stacktrace"))
+	}
+}
+
+func verifyLockEnabled() bool {
+	return os.Getenv(ENV_VERIFY) == ENV_VERIFY_ALL_VALUE || os.Getenv(ENV_VERIFY) == ENV_VERIFY_LOCK
+}
+
+func insideApply() bool {
+	stackTraceStr := string(debug.Stack())
+	return strings.Contains(stackTraceStr, ".applyEntries")
+}

server/storage/backend/verify_test.go

@@ -0,0 +1,91 @@
+// Copyright 2022 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package backend_test
+
+import (
+	"fmt"
+	"os"
+	"testing"
+	"time"
+
+	"go.etcd.io/etcd/server/v3/storage/backend"
+	betesting "go.etcd.io/etcd/server/v3/storage/backend/testing"
+)
+
+func TestLockVerify(t *testing.T) {
+	tcs := []struct {
+		insideApply bool
+		lock        func(tx backend.BatchTx)
+		expectPanic bool
+	}{
+		{
+			insideApply: true,
+			lock:        lockInsideApply,
+			expectPanic: false,
+		},
+		{
+			insideApply: false,
+			lock:        lockInsideApply,
+			expectPanic: true,
+		},
+		{
+			insideApply: false,
+			lock:        lockOutsideApply,
+			expectPanic: false,
+		},
+		{
+			insideApply: true,
+			lock:        lockOutsideApply,
+			expectPanic: true,
+		},
+	}
+	env := os.Getenv("ETCD_VERIFY")
+	os.Setenv("ETCD_VERIFY", "lock")
+	defer func() {
+		os.Setenv("ETCD_VERIFY", env)
+	}()
+	for i, tc := range tcs {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+
+			be, _ := betesting.NewTmpBackend(t, time.Hour, 10000)
+
+			hasPaniced := handlePanic(func() {
+				if tc.insideApply {
+					applyEntries(be, tc.lock)
+				} else {
+					tc.lock(be.BatchTx())
+				}
+			}) != nil
+			if hasPaniced != tc.expectPanic {
+				t.Errorf("%v != %v", hasPaniced, tc.expectPanic)
+			}
+		})
+	}
+}
+
+func handlePanic(f func()) (result interface{}) {
+	defer func() {
+		result = recover()
+	}()
+	f()
+	return result
+}
+
+func applyEntries(be backend.Backend, f func(tx backend.BatchTx)) {
+	f(be.BatchTx())
+}
+
+func lockInsideApply(tx backend.BatchTx)  { tx.LockInsideApply() }
+func lockOutsideApply(tx backend.BatchTx) { tx.LockOutsideApply() }

server/storage/mvcc/kvstore_test.go

@@ -904,8 +904,10 @@ func (b *fakeBatchTx) UnsafeDelete(bucket backend.Bucket, key []byte) {
 func (b *fakeBatchTx) UnsafeForEach(bucket backend.Bucket, visitor func(k, v []byte) error) error {
 	return nil
 }
-func (b *fakeBatchTx) Commit()        {}
-func (b *fakeBatchTx) CommitAndStop() {}
+func (b *fakeBatchTx) Commit()           {}
+func (b *fakeBatchTx) CommitAndStop()    {}
+func (b *fakeBatchTx) LockInsideApply()  {}
+func (b *fakeBatchTx) LockOutsideApply() {}
 
 type fakeBackend struct {
 	tx *fakeBatchTx