keyspace: add preconditions for keyspace config update

hbisheng · hbisheng · commit 6897830815f1 · 2026-02-02T14:04:11.000+08:00
Changes:
- Extend /pd/api/v2/keyspaces/{name}/config PATCH with "preconditions" to support CAS-like config updates (equal / absent).
- Add pd-ctl support: keyspace update-config --expect and --expect-absent.

Rationale:
We need PD to provide a simple CAS-style guard for keyspace config updates, so concurrent key-rotation workflows can coordinate safely and only the expected state transition is allowed.
diff --git a/pkg/errs/errno.go b/pkg/errs/errno.go
@@ -247,6 +247,8 @@ var (
 	ErrIllegalOperation = errors.Normalize("unknown operation", errors.RFCCodeText("PD:keyspace:ErrIllegalOperation"))
 	// ErrUnsupportedOperationInKeyspace is used to indicate this is an unsupported operation.
 	ErrUnsupportedOperationInKeyspace = errors.Normalize("it's a unsupported operation", errors.RFCCodeText("PD:keyspace:ErrUnsupportedOperationInKeyspace"))
+	// ErrKeyspaceConfigPreconditionFailed is used to indicate the precondition for updating keyspace config is not met.
+	ErrKeyspaceConfigPreconditionFailed = errors.Normalize("keyspace config precondition failed, %s", errors.RFCCodeText("PD:keyspace:ErrKeyspaceConfigPreconditionFailed"))
 	// ErrKeyspaceGroupPrimaryNotFound is used to indicate primary of target keyspace group does not exist.
 	ErrKeyspaceGroupPrimaryNotFound = errors.Normalize("primary of keyspace group does not exist", errors.RFCCodeText("PD:keyspace:ErrKeyspaceGroupPrimaryNotFound"))
 	// ErrKeyspaceGroupNotExists is used to indicate target keyspace group does not exist.
diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go
@@ -645,6 +645,68 @@ const (
 // UpdateKeyspaceConfig changes target keyspace's config in the order specified in mutations.
 // It returns error if saving failed, operation not allowed, or if keyspace not exists.
 func (manager *Manager) UpdateKeyspaceConfig(name string, mutations []*Mutation) (*keyspacepb.KeyspaceMeta, error) {
+	return manager.updateKeyspaceConfig(name, func(meta *keyspacepb.KeyspaceMeta) error {
+		for _, mutation := range mutations {
+			switch mutation.Op {
+			case OpPut:
+				meta.Config[mutation.Key] = mutation.Value
+			case OpDel:
+				delete(meta.Config, mutation.Key)
+			default:
+				return errs.ErrIllegalOperation
+			}
+		}
+		return nil
+	})
+}
+
+// UpdateKeyspaceConfigWithPreconditions changes target keyspace's config in the order specified in mutations if the
+// given preconditions are satisfied.
+// Preconditions use a JSON-merge-patch-like encoding:
+// - key -> null means the key must be absent.
+// - key -> "value" means the key must exist and equal "value".
+func (manager *Manager) UpdateKeyspaceConfigWithPreconditions(name string, mutations []*Mutation, preconditions map[string]*string) (*keyspacepb.KeyspaceMeta, error) {
+	if len(preconditions) == 0 {
+		return manager.UpdateKeyspaceConfig(name, mutations)
+	}
+	return manager.updateKeyspaceConfig(name, func(meta *keyspacepb.KeyspaceMeta) error {
+		if err := checkKeyspaceConfigPreconditions(meta.GetConfig(), preconditions); err != nil {
+			return err
+		}
+		for _, mutation := range mutations {
+			switch mutation.Op {
+			case OpPut:
+				meta.Config[mutation.Key] = mutation.Value
+			case OpDel:
+				delete(meta.Config, mutation.Key)
+			default:
+				return errs.ErrIllegalOperation
+			}
+		}
+		return nil
+	})
+}
+
+func checkKeyspaceConfigPreconditions(config map[string]string, preconditions map[string]*string) error {
+	for k, expected := range preconditions {
+		actual, exists := config[k]
+		if expected == nil {
+			if exists {
+				return errs.ErrKeyspaceConfigPreconditionFailed.FastGenByArgs(k + " must be absent")
+			}
+			continue
+		}
+		if !exists {
+			return errs.ErrKeyspaceConfigPreconditionFailed.FastGenByArgs(k + " does not exist")
+		}
+		if actual != *expected {
+			return errs.ErrKeyspaceConfigPreconditionFailed.FastGenByArgs("key=" + k + " expected=" + *expected + " actual=" + actual)
+		}
+	}
+	return nil
+}
+
+func (manager *Manager) updateKeyspaceConfig(name string, update func(meta *keyspacepb.KeyspaceMeta) error) (*keyspacepb.KeyspaceMeta, error) {
 	var meta *keyspacepb.KeyspaceMeta
 	oldConfig := make(map[string]string)
 	err := manager.store.RunInTxn(manager.ctx, func(txn kv.Txn) error {
@@ -677,16 +739,9 @@ func (manager *Manager) UpdateKeyspaceConfig(name string, mutations []*Mutation)
 		for k, v := range meta.GetConfig() {
 			oldConfig[k] = v
 		}
-		// Update keyspace config according to mutations.
-		for _, mutation := range mutations {
-			switch mutation.Op {
-			case OpPut:
-				meta.Config[mutation.Key] = mutation.Value
-			case OpDel:
-				delete(meta.Config, mutation.Key)
-			default:
-				return errs.ErrIllegalOperation
-			}
+		// Update keyspace config.
+		if err := update(meta); err != nil {
+			return err
 		}
 		newConfig := meta.GetConfig()
 		oldUserKind := endpoint.StringUserKind(oldConfig[UserKindKey])
diff --git a/pkg/keyspace/keyspace_test.go b/pkg/keyspace/keyspace_test.go
@@ -16,6 +16,7 @@ package keyspace
 
 import (
 	"context"
+	goerrors "errors"
 	"fmt"
 	"math"
 	"strconv"
@@ -390,6 +391,68 @@ func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfig() {
 	checkMutations(re, nil, updated.Config, mutations)
 }
 
+func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfigWithPreconditions() {
+	re := suite.Require()
+	manager := suite.manager
+
+	ksName := "precond_test"
+	_, err := manager.CreateKeyspace(&CreateKeyspaceRequest{
+		Name:       ksName,
+		Config:     map[string]string{},
+		CreateTime: time.Now().Unix(),
+	})
+	re.NoError(err)
+
+	currentKey := "current_file_id"
+	nextKey := "next_file_id"
+
+	// 1) only set next_file_id if it doesn't exist already.
+	nextV1 := "1000"
+	meta, err := manager.UpdateKeyspaceConfigWithPreconditions(ksName, []*Mutation{
+		{Op: OpPut, Key: nextKey, Value: nextV1},
+	}, map[string]*string{
+		nextKey: nil,
+	})
+	re.NoError(err)
+	re.Equal(nextV1, meta.Config[nextKey])
+
+	nextV2 := "2000"
+	_, err = manager.UpdateKeyspaceConfigWithPreconditions(ksName, []*Mutation{
+		{Op: OpPut, Key: nextKey, Value: nextV2},
+	}, map[string]*string{
+		nextKey: nil,
+	})
+	re.Error(err)
+	re.True(goerrors.Is(err, errs.ErrKeyspaceConfigPreconditionFailed))
+
+	// 2) Update current_file_id to the next_file_id (guarded by next_file_id == expected).
+	meta, err = manager.UpdateKeyspaceConfigWithPreconditions(ksName, []*Mutation{
+		{Op: OpPut, Key: currentKey, Value: nextV1},
+	}, map[string]*string{
+		nextKey: &nextV1,
+	})
+	re.NoError(err)
+	re.Equal(nextV1, meta.Config[currentKey])
+
+	// 3) Delete next_file_id if it matches my expected value.
+	wrongExpected := "999"
+	_, err = manager.UpdateKeyspaceConfigWithPreconditions(ksName, []*Mutation{
+		{Op: OpDel, Key: nextKey},
+	}, map[string]*string{
+		nextKey: &wrongExpected,
+	})
+	re.Error(err)
+	re.True(goerrors.Is(err, errs.ErrKeyspaceConfigPreconditionFailed))
+
+	meta, err = manager.UpdateKeyspaceConfigWithPreconditions(ksName, []*Mutation{
+		{Op: OpDel, Key: nextKey},
+	}, map[string]*string{
+		nextKey: &nextV1,
+	})
+	re.NoError(err)
+	re.NotContains(meta.GetConfig(), nextKey)
+}
+
 func (suite *keyspaceTestSuite) TestUpdateKeyspaceState() {
 	re := suite.Require()
 	manager := suite.manager
diff --git a/server/apiv2/handlers/keyspace.go b/server/apiv2/handlers/keyspace.go
@@ -16,6 +16,7 @@ package handlers
 
 import (
 	"encoding/json"
+	goerrors "errors"
 	"net/http"
 	"strconv"
 	"strings"
@@ -308,6 +309,10 @@ func LoadAllKeyspaces(c *gin.Context) {
 // which will both be set to "" if value type is string during binding.
 type UpdateConfigParams struct {
 	Config map[string]*string `json:"config"`
+	// Preconditions specifies prerequisites for updating config, using a JSON-merge-patch-like encoding:
+	// - key -> null means the key must be absent.
+	// - key -> "value" means the key must exist and equal "value".
+	Preconditions map[string]*string `json:"preconditions,omitempty"`
 }
 
 // UpdateKeyspaceConfig updates target keyspace's config.
@@ -350,8 +355,16 @@ func UpdateKeyspaceConfig(c *gin.Context) {
 		}
 	}
 
-	meta, err := manager.UpdateKeyspaceConfig(name, mutations)
+	meta, err := manager.UpdateKeyspaceConfigWithPreconditions(name, mutations, configParams.Preconditions)
 	if err != nil {
+		if goerrors.Is(err, errs.ErrKeyspaceConfigPreconditionFailed) {
+			c.AbortWithStatusJSON(http.StatusConflict, err.Error())
+			return
+		}
+		if goerrors.Is(err, errs.ErrEtcdTxnConflict) {
+			c.AbortWithStatusJSON(http.StatusConflict, err.Error())
+			return
+		}
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
 		return
 	}
diff --git a/tests/server/apiv2/handlers/keyspace_test.go b/tests/server/apiv2/handlers/keyspace_test.go
@@ -113,6 +113,132 @@ func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfig() {
 	}
 }
 
+func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfigPreconditions() {
+	re := suite.Require()
+	created := MustCreateKeyspace(re, suite.server, &handlers.CreateKeyspaceParams{
+		Name:   "test_keyspace_cas",
+		Config: map[string]string{},
+	})
+	re.NotNil(created)
+
+	currentKey := "current_file_id"
+	nextKey := "next_file_id"
+
+	next := "1000"
+	status, body, meta := tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+		Config: map[string]*string{
+			nextKey: &next,
+		},
+		Preconditions: map[string]*string{
+			nextKey: nil,
+		},
+	})
+	re.Equal(http.StatusOK, status, body)
+	re.Equal(next, meta.Config[nextKey])
+
+	otherNext := "2000"
+	status, body, meta = tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+		Config: map[string]*string{
+			nextKey: &otherNext,
+		},
+		Preconditions: map[string]*string{
+			nextKey: nil,
+		},
+	})
+	re.Equal(http.StatusConflict, status)
+	re.Nil(meta)
+	re.Contains(body, "precondition failed")
+
+	status, _, meta = tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+		Config: map[string]*string{
+			currentKey: &next,
+		},
+		Preconditions: map[string]*string{
+			nextKey: &next,
+		},
+	})
+	re.Equal(http.StatusOK, status)
+	re.Equal(next, meta.Config[currentKey])
+	re.Equal(next, meta.Config[nextKey])
+
+	wrongExpected := "999"
+	status, body, meta = tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+		Config: map[string]*string{
+			nextKey: nil,
+		},
+		Preconditions: map[string]*string{
+			nextKey: &wrongExpected,
+		},
+	})
+	re.Equal(http.StatusConflict, status)
+	re.Nil(meta)
+	re.Contains(body, "expected")
+
+	status, _, meta = tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+		Config: map[string]*string{
+			nextKey: nil,
+		},
+		Preconditions: map[string]*string{
+			nextKey: &next,
+		},
+	})
+	re.Equal(http.StatusOK, status)
+	re.Equal(next, meta.Config[currentKey])
+	re.NotContains(meta.GetConfig(), nextKey)
+}
+
+func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfigPreconditionsConcurrentSetNextIfAbsent() {
+	re := suite.Require()
+	created := MustCreateKeyspace(re, suite.server, &handlers.CreateKeyspaceParams{
+		Name:   "test_keyspace_cas2",
+		Config: map[string]string{},
+	})
+	re.NotNil(created)
+
+	nextKey := "next_file_id"
+	start := make(chan struct{})
+	results := make(chan int, 2)
+
+	go func() {
+		<-start
+		next := "1000"
+		status, body, _ := tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+			Config: map[string]*string{
+				nextKey: &next,
+			},
+			Preconditions: map[string]*string{
+				nextKey: nil,
+			},
+		})
+		if status != http.StatusOK && status != http.StatusConflict {
+			re.FailNow("unexpected status", "status=%d body=%s", status, body)
+		}
+		results <- status
+	}()
+	go func() {
+		<-start
+		next := "2000"
+		status, body, _ := tryUpdateKeyspaceConfig(re, suite.server, created.Name, &handlers.UpdateConfigParams{
+			Config: map[string]*string{
+				nextKey: &next,
+			},
+			Preconditions: map[string]*string{
+				nextKey: nil,
+			},
+		})
+		if status != http.StatusOK && status != http.StatusConflict {
+			re.FailNow("unexpected status", "status=%d body=%s", status, body)
+		}
+		results <- status
+	}()
+
+	close(start)
+
+	s1 := <-results
+	s2 := <-results
+	re.ElementsMatch([]int{http.StatusOK, http.StatusConflict}, []int{s1, s2})
+}
+
 func (suite *keyspaceTestSuite) TestUpdateKeyspaceState() {
 	re := suite.Require()
 	keyspaces := mustMakeTestKeyspaces(re, suite.server, 10)
diff --git a/tests/server/apiv2/handlers/testutil.go b/tests/server/apiv2/handlers/testutil.go
@@ -152,6 +152,24 @@ func mustUpdateKeyspaceConfig(re *require.Assertions, server *tests.TestServer,
 	return meta.KeyspaceMeta
 }
 
+func tryUpdateKeyspaceConfig(re *require.Assertions, server *tests.TestServer, name string, request *handlers.UpdateConfigParams) (int, string, *keyspacepb.KeyspaceMeta) {
+	data, err := json.Marshal(request)
+	re.NoError(err)
+	httpReq, err := http.NewRequest(http.MethodPatch, server.GetAddr()+keyspacesPrefix+"/"+name+"/config", bytes.NewBuffer(data))
+	re.NoError(err)
+	resp, err := tests.TestDialClient.Do(httpReq)
+	re.NoError(err)
+	defer resp.Body.Close()
+	data, err = io.ReadAll(resp.Body)
+	re.NoError(err)
+	if resp.StatusCode != http.StatusOK {
+		return resp.StatusCode, string(data), nil
+	}
+	meta := &handlers.KeyspaceMeta{}
+	re.NoError(json.Unmarshal(data, meta))
+	return resp.StatusCode, string(data), meta.KeyspaceMeta
+}
+
 func mustLoadKeyspaces(re *require.Assertions, server *tests.TestServer, name string) *keyspacepb.KeyspaceMeta {
 	resp, err := tests.TestDialClient.Get(server.GetAddr() + keyspacesPrefix + "/" + name)
 	re.NoError(err)
diff --git a/tools/pd-ctl/pdctl/command/keyspace_command.go b/tools/pd-ctl/pdctl/command/keyspace_command.go
diff --git a/tools/pd-ctl/tests/keyspace/keyspace_test.go b/tools/pd-ctl/tests/keyspace/keyspace_test.go
