diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index 4b7642adcf5..a8042e36699 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -424,6 +424,74 @@ func (m *GroupManager) GetGroupByKeyspaceID(id uint32) (uint32, error) { return 0, errs.ErrKeyspaceNotInAnyKeyspaceGroup } +// RemoveKeyspacesFromGroup removes the specified keyspaces from the given keyspace group. +// If a keyspace is not in the group, it will be skipped (no error). +// It returns the updated keyspace group and any error encountered. +func (m *GroupManager) RemoveKeyspacesFromGroup(groupID uint32, keyspaceIDs []uint32) (*endpoint.KeyspaceGroup, error) { + m.Lock() + defer m.Unlock() + + var ( + kg *endpoint.KeyspaceGroup + err error + ) + + if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { + // Load the keyspace group + kg, err = m.store.LoadKeyspaceGroup(txn, groupID) + if err != nil { + return err + } + if kg == nil { + return errs.ErrKeyspaceGroupNotExists.FastGenByArgs(groupID) + } + if kg.IsSplitting() { + return errs.ErrKeyspaceGroupInSplit.FastGenByArgs(groupID) + } + if kg.IsMerging() { + return errs.ErrKeyspaceGroupInMerging.FastGenByArgs(groupID) + } + + // Build a set of keyspaces to remove (excluding protected bootstrap/system keyspace) + toRemove := make(map[uint32]struct{}) + for _, ksID := range keyspaceIDs { + // Keep the protected bootstrap/system keyspace in the default group. + if isProtectedKeyspaceID(ksID) { + continue + } + // Only add if it exists in the group (skip if not present) + if slice.Contains(kg.Keyspaces, ksID) { + toRemove[ksID] = struct{}{} + } + } + + // If nothing to remove, return nil to skip update + if len(toRemove) == 0 { + return nil + } + + // Filter out keyspaces to remove + newKeyspaces := make([]uint32, 0, len(kg.Keyspaces)-len(toRemove)) + for _, ks := range kg.Keyspaces { + if _, shouldRemove := toRemove[ks]; !shouldRemove { + newKeyspaces = append(newKeyspaces, ks) + } + } + kg.Keyspaces = newKeyspaces + + // Save the updated keyspace group + return m.store.SaveKeyspaceGroup(txn, kg) + }); err != nil { + return nil, err + } + + // Update the cache + userKind := endpoint.StringUserKind(kg.UserKind) + m.groups[userKind].Put(kg) + + return kg, nil +} + var failpointOnce sync.Once // UpdateKeyspaceForGroup updates the keyspace field for the keyspace group. diff --git a/server/apiv2/handlers/tso_keyspace_group.go b/server/apiv2/handlers/tso_keyspace_group.go index f976625ca9e..ed6200ab7bb 100644 --- a/server/apiv2/handlers/tso_keyspace_group.go +++ b/server/apiv2/handlers/tso_keyspace_group.go @@ -22,6 +22,7 @@ import ( "github.com/gin-gonic/gin" "github.com/pingcap/errors" + "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/tikv/pd/pkg/errs" "github.com/tikv/pd/pkg/keyspace/constant" @@ -51,6 +52,7 @@ func RegisterTSOKeyspaceGroup(r *gin.RouterGroup) { router.DELETE("/:id/split", FinishSplitKeyspaceByID) router.POST("/:id/merge", MergeKeyspaceGroups) router.DELETE("/:id/merge", FinishMergeKeyspaceByID) + router.DELETE("/:id/keyspaces", RemoveKeyspacesFromGroup) } // CreateKeyspaceGroupParams defines the params for creating keyspace groups. @@ -558,3 +560,89 @@ func parseNodeAddress(c *gin.Context) (string, error) { func isValid(id uint32) bool { return id >= constant.DefaultKeyspaceGroupID && id <= mcs.MaxKeyspaceGroupCountInUse } + +// RemoveKeyspacesFromGroupParams defines the params for removing keyspaces from a keyspace group. +type RemoveKeyspacesFromGroupParams struct { + Keyspaces []uint32 `json:"keyspaces"` +} + +// RemoveKeyspacesFromGroup removes the specified keyspaces from the given keyspace group. +// Keyspaces in archived or tombstone state will be removed. Keyspaces not in the group will be skipped. +func RemoveKeyspacesFromGroup(c *gin.Context) { + // Parse and validate group ID + groupID, err := validateKeyspaceGroupID(c) + if err != nil { + c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id") + return + } + + // Parse request body + var params RemoveKeyspacesFromGroupParams + if err := c.BindJSON(¶ms); err != nil { + c.AbortWithStatusJSON(http.StatusBadRequest, errs.ErrBindJSON.Wrap(err).GenWithStackByCause().Error()) + return + } + + if len(params.Keyspaces) == 0 { + c.AbortWithStatusJSON(http.StatusBadRequest, "keyspaces list cannot be empty") + return + } + + svr := c.MustGet(middlewares.ServerContextKey).(*server.Server) + keyspaceManager := svr.GetKeyspaceManager() + if keyspaceManager == nil { + c.AbortWithStatusJSON(http.StatusInternalServerError, managerUninitializedErr) + return + } + + groupManager := svr.GetKeyspaceGroupManager() + if groupManager == nil { + c.AbortWithStatusJSON(http.StatusInternalServerError, GroupManagerUninitializedErr) + return + } + + // Filter keyspaces: only keep those in ARCHIVED or TOMBSTONE state + var validKeyspaces []uint32 + for _, keyspaceID := range params.Keyspaces { + // Load the keyspace meta to check its state + keyspaceMeta, err := keyspaceManager.LoadKeyspaceByID(keyspaceID) + if err != nil { + // Skip if keyspace doesn't exist + if errors.ErrorEqual(err, errs.ErrKeyspaceNotFound) { + continue + } + c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) + return + } + + // Check if the keyspace is in archived or tombstone state + state := keyspaceMeta.GetState() + if state == keyspacepb.KeyspaceState_ARCHIVED || state == keyspacepb.KeyspaceState_TOMBSTONE { + validKeyspaces = append(validKeyspaces, keyspaceID) + } + } + + // If no valid keyspaces to remove, load and return the group without modification + if len(validKeyspaces) == 0 { + kg, err := groupManager.GetKeyspaceGroupByID(groupID) + if err != nil { + c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) + return + } + if kg == nil { + c.AbortWithStatusJSON(http.StatusInternalServerError, errs.ErrKeyspaceGroupNotExists.FastGenByArgs(groupID).Error()) + return + } + c.IndentedJSON(http.StatusOK, kg) + return + } + + // Remove the keyspaces from the keyspace group + kg, err := groupManager.RemoveKeyspacesFromGroup(groupID, validKeyspaces) + if err != nil { + c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) + return + } + + c.IndentedJSON(http.StatusOK, kg) +} diff --git a/tests/server/apiv2/handlers/testutil.go b/tests/server/apiv2/handlers/testutil.go index ce7cb565ad0..565cb695b66 100644 --- a/tests/server/apiv2/handlers/testutil.go +++ b/tests/server/apiv2/handlers/testutil.go @@ -326,3 +326,40 @@ func MustMergeKeyspaceGroup(re *require.Assertions, server *tests.TestServer, id re.NoError(err) re.Equal(http.StatusOK, resp.StatusCode, string(data)) } + +// MustRemoveKeyspacesFromGroup removes keyspaces from a keyspace group with HTTP API. +func MustRemoveKeyspacesFromGroup(re *require.Assertions, server *tests.TestServer, groupID uint32, keyspaceIDs []uint32) *endpoint.KeyspaceGroup { + params := &handlers.RemoveKeyspacesFromGroupParams{ + Keyspaces: keyspaceIDs, + } + data, err := json.Marshal(params) + re.NoError(err) + httpReq, err := http.NewRequest(http.MethodDelete, server.GetAddr()+keyspaceGroupsPrefix+fmt.Sprintf("/%d/keyspaces", groupID), bytes.NewBuffer(data)) + re.NoError(err) + resp, err := tests.TestDialClient.Do(httpReq) + re.NoError(err) + defer resp.Body.Close() + respData, err := io.ReadAll(resp.Body) + re.NoError(err) + re.Equal(http.StatusOK, resp.StatusCode, string(respData)) + var kg endpoint.KeyspaceGroup + re.NoError(json.Unmarshal(respData, &kg)) + return &kg +} + +// FailRemoveKeyspacesFromGroupWithCode fails to remove keyspaces from a keyspace group with HTTP API. +func FailRemoveKeyspacesFromGroupWithCode(re *require.Assertions, server *tests.TestServer, groupID uint32, keyspaceIDs []uint32, expectCode int) { + params := &handlers.RemoveKeyspacesFromGroupParams{ + Keyspaces: keyspaceIDs, + } + data, err := json.Marshal(params) + re.NoError(err) + httpReq, err := http.NewRequest(http.MethodDelete, server.GetAddr()+keyspaceGroupsPrefix+fmt.Sprintf("/%d/keyspaces", groupID), bytes.NewBuffer(data)) + re.NoError(err) + resp, err := tests.TestDialClient.Do(httpReq) + re.NoError(err) + defer resp.Body.Close() + respData, err := io.ReadAll(resp.Body) + re.NoError(err) + re.Equal(expectCode, resp.StatusCode, string(respData)) +} diff --git a/tests/server/apiv2/handlers/tso_keyspace_group_test.go b/tests/server/apiv2/handlers/tso_keyspace_group_test.go index 5e48f4ec6ee..d5a11591468 100644 --- a/tests/server/apiv2/handlers/tso_keyspace_group_test.go +++ b/tests/server/apiv2/handlers/tso_keyspace_group_test.go @@ -23,7 +23,12 @@ import ( "github.com/stretchr/testify/suite" - "github.com/tikv/pd/pkg/mcs/utils/constant" + "github.com/pingcap/failpoint" + "github.com/pingcap/kvproto/pkg/keyspacepb" + + "github.com/tikv/pd/pkg/keyspace" + kgconstant "github.com/tikv/pd/pkg/keyspace/constant" + mcsconstant "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/storage/endpoint" "github.com/tikv/pd/server/apiv2/handlers" "github.com/tikv/pd/tests" @@ -100,7 +105,7 @@ func (suite *keyspaceGroupTestSuite) TestCreateKeyspaceGroups() { // invalid ID. kgs = &handlers.CreateKeyspaceGroupParams{KeyspaceGroups: []*endpoint.KeyspaceGroup{ { - ID: constant.MaxKeyspaceGroupCount + 1, + ID: mcsconstant.MaxKeyspaceGroupCount + 1, UserKind: endpoint.Standard.String(), }, }} @@ -141,7 +146,7 @@ func (suite *keyspaceGroupTestSuite) TestSplitKeyspaceGroup() { ID: uint32(1), UserKind: endpoint.Standard.String(), Keyspaces: []uint32{111, 222, 333}, - Members: make([]endpoint.KeyspaceGroupMember, constant.DefaultKeyspaceGroupReplicaCount), + Members: make([]endpoint.KeyspaceGroupMember, mcsconstant.DefaultKeyspaceGroupReplicaCount), }, }} @@ -248,7 +253,6 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupErrorMessage() { re.NoError(json.NewDecoder(resp.Body).Decode(&errorMsg)) re.NotEmpty(errorMsg, "Error message should not be empty") re.Contains(errorMsg, "invalid", "Error message should indicate invalid input") - // Test SetPriorityForKeyspaceGroup with invalid JSON httpReq, err = http.NewRequest( http.MethodPatch, @@ -266,3 +270,114 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupErrorMessage() { re.NotEmpty(errorMsg, "Error message should not be empty") re.Contains(errorMsg, "invalid", "Error message should indicate invalid input") } + +func (suite *keyspaceGroupTestSuite) TestRemoveKeyspacesFromGroup() { + re := suite.Require() + re.NoError(failpoint.Enable("github.com/tikv/pd/server/delayStartServerLoop", `return(true)`)) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/server/delayStartServerLoop")) + }() + re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)")) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion")) + }() + + keyspaceManager := suite.server.GetKeyspaceManager() + re.NotNil(keyspaceManager) + + // Create test keyspaces (automatically added to default keyspace group 0) + keyspaceMeta1, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{ + Name: "test_keyspace_1", + CreateTime: 0, + }) + re.NoError(err) + keyspaceID1 := keyspaceMeta1.GetId() + + keyspaceMeta2, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{ + Name: "test_keyspace_2", + CreateTime: 0, + }) + re.NoError(err) + keyspaceID2 := keyspaceMeta2.GetId() + + keyspaceMeta3, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{ + Name: "test_keyspace_3", + CreateTime: 0, + }) + re.NoError(err) + keyspaceID3 := keyspaceMeta3.GetId() + + // Verify all keyspaces are in the default group + kg := MustLoadKeyspaceGroupByID(re, suite.server, kgconstant.DefaultKeyspaceGroupID) + re.Contains(kg.Keyspaces, keyspaceID1) + re.Contains(kg.Keyspaces, keyspaceID2) + re.Contains(kg.Keyspaces, keyspaceID3) + + // Test 1: Try to remove ENABLED keyspaces (should succeed but nothing removed) + kg = MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID, + []uint32{keyspaceID1}) + // Verify nothing is removed (keyspace is still there because it's ENABLED) + re.Contains(kg.Keyspaces, keyspaceID1) + + // Test 2: Update keyspaces to ARCHIVED/TOMBSTONE state and batch remove + // Set keyspace1 to ARCHIVED + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID1, keyspacepb.KeyspaceState_DISABLED, 0) + re.NoError(err) + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID1, keyspacepb.KeyspaceState_ARCHIVED, 0) + re.NoError(err) + + // Set keyspace2 to TOMBSTONE + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_DISABLED, 0) + re.NoError(err) + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_ARCHIVED, 0) + re.NoError(err) + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_TOMBSTONE, 0) + re.NoError(err) + + // Batch remove keyspace1 and keyspace2 + MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID, + []uint32{keyspaceID1, keyspaceID2}) + + // Verify both keyspaces are removed + kg = MustLoadKeyspaceGroupByID(re, suite.server, kgconstant.DefaultKeyspaceGroupID) + re.NotContains(kg.Keyspaces, keyspaceID1) + re.NotContains(kg.Keyspaces, keyspaceID2) + re.Contains(kg.Keyspaces, keyspaceID3) // keyspace3 should still be there + + // Test 3: Mix valid and invalid keyspaces + // Set keyspace3 to ARCHIVED + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID3, keyspacepb.KeyspaceState_DISABLED, 0) + re.NoError(err) + _, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID3, keyspacepb.KeyspaceState_ARCHIVED, 0) + re.NoError(err) + + // Include: valid (keyspace3), already removed (keyspace1), non-existent (99999) + // Should only remove keyspace3, others are skipped + MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID, + []uint32{keyspaceID3, keyspaceID1, 99999}) + + // Verify only keyspace3 is removed + kg = MustLoadKeyspaceGroupByID(re, suite.server, kgconstant.DefaultKeyspaceGroupID) + re.NotContains(kg.Keyspaces, keyspaceID3) + + // Test 4: Try to remove from non-existent group + FailRemoveKeyspacesFromGroupWithCode(re, suite.server, 999, + []uint32{keyspaceID1}, http.StatusInternalServerError) + + // Test 5: Try to remove with empty keyspace list (should fail - empty list) + FailRemoveKeyspacesFromGroupWithCode(re, suite.server, kgconstant.DefaultKeyspaceGroupID, + []uint32{}, http.StatusBadRequest) + + // Test 6: All keyspaces in wrong state (should succeed but nothing removed) + keyspaceMeta4, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{ + Name: "test_keyspace_4", + CreateTime: 0, + }) + re.NoError(err) + keyspaceID4 := keyspaceMeta4.GetId() + + kg = MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID, + []uint32{keyspaceID4}) // ENABLED state, will be skipped + // Verify keyspace4 is still there + re.Contains(kg.Keyspaces, keyspaceID4) +}