Skip to content

Commit 911a71a

Browse files
ystaticybufferflies
ystaticy
authored and
bufferflies
committed
tso:support remove archive & tombstone keyspace from keyspace group (#420)
* support remove archive & tombstone keyspace from keyspace group Signed-off-by: ystaticy <y_static_y@sina.com> * support keyspace list Signed-off-by: ystaticy <y_static_y@sina.com> * no keyspace to remove ,return succ Signed-off-by: ystaticy <y_static_y@sina.com> * remove debug log Signed-off-by: ystaticy <y_static_y@sina.com> --------- Signed-off-by: ystaticy <y_static_y@sina.com> (cherry picked from commit 407ad06)
1 parent 048f0d8 commit 911a71a

File tree

4 files changed

+296
-4
lines changed

4 files changed

+296
-4
lines changed

‎pkg/keyspace/tso_keyspace_group.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -424,6 +424,68 @@ func (m *GroupManager) GetGroupByKeyspaceID(id uint32) (uint32, error) {
424424
return 0, errs.ErrKeyspaceNotInAnyKeyspaceGroup
425425
}
426426

427+
// RemoveKeyspacesFromGroup removes the specified keyspaces from the given keyspace group.
428+
// If a keyspace is not in the group, it will be skipped (no error).
429+
// It returns the updated keyspace group and any error encountered.
430+
func (m *GroupManager) RemoveKeyspacesFromGroup(groupID uint32, keyspaceIDs []uint32) (*endpoint.KeyspaceGroup, error) {
431+
m.Lock()
432+
defer m.Unlock()
433+
434+
var (
435+
kg *endpoint.KeyspaceGroup
436+
err error
437+
)
438+
439+
if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error {
440+
// Load the keyspace group
441+
kg, err = m.store.LoadKeyspaceGroup(txn, groupID)
442+
if err != nil {
443+
return err
444+
}
445+
if kg == nil {
446+
return errs.ErrKeyspaceGroupNotExists.FastGenByArgs(groupID)
447+
}
448+
449+
// Build a set of keyspaces to remove (excluding default keyspace)
450+
toRemove := make(map[uint32]struct{})
451+
for _, ksID := range keyspaceIDs {
452+
// Skip default keyspace
453+
if ksID == constant.DefaultKeyspaceID {
454+
continue
455+
}
456+
// Only add if it exists in the group (skip if not present)
457+
if slice.Contains(kg.Keyspaces, ksID) {
458+
toRemove[ksID] = struct{}{}
459+
}
460+
}
461+
462+
// If nothing to remove, return nil to skip update
463+
if len(toRemove) == 0 {
464+
return nil
465+
}
466+
467+
// Filter out keyspaces to remove
468+
newKeyspaces := make([]uint32, 0, len(kg.Keyspaces)-len(toRemove))
469+
for _, ks := range kg.Keyspaces {
470+
if _, shouldRemove := toRemove[ks]; !shouldRemove {
471+
newKeyspaces = append(newKeyspaces, ks)
472+
}
473+
}
474+
kg.Keyspaces = newKeyspaces
475+
476+
// Save the updated keyspace group
477+
return m.store.SaveKeyspaceGroup(txn, kg)
478+
}); err != nil {
479+
return nil, err
480+
}
481+
482+
// Update the cache
483+
userKind := endpoint.StringUserKind(kg.UserKind)
484+
m.groups[userKind].Put(kg)
485+
486+
return kg, nil
487+
}
488+
427489
var failpointOnce sync.Once
428490

429491
// UpdateKeyspaceForGroup updates the keyspace field for the keyspace group.

‎server/apiv2/handlers/tso_keyspace_group.go‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"github.com/gin-gonic/gin"
2323

2424
"github.com/pingcap/errors"
25+
"github.com/pingcap/kvproto/pkg/keyspacepb"
2526

2627
"github.com/tikv/pd/pkg/errs"
2728
"github.com/tikv/pd/pkg/keyspace/constant"
@@ -51,6 +52,7 @@ func RegisterTSOKeyspaceGroup(r *gin.RouterGroup) {
5152
router.DELETE("/:id/split", FinishSplitKeyspaceByID)
5253
router.POST("/:id/merge", MergeKeyspaceGroups)
5354
router.DELETE("/:id/merge", FinishMergeKeyspaceByID)
55+
router.DELETE("/:id/keyspaces", RemoveKeyspacesFromGroup)
5456
}
5557

5658
// CreateKeyspaceGroupParams defines the params for creating keyspace groups.
@@ -558,3 +560,85 @@ func parseNodeAddress(c *gin.Context) (string, error) {
558560
func isValid(id uint32) bool {
559561
return id >= constant.DefaultKeyspaceGroupID && id <= mcs.MaxKeyspaceGroupCountInUse
560562
}
563+
564+
// RemoveKeyspacesFromGroupParams defines the params for removing keyspaces from a keyspace group.
565+
type RemoveKeyspacesFromGroupParams struct {
566+
Keyspaces []uint32 `json:"keyspaces"`
567+
}
568+
569+
// RemoveKeyspacesFromGroup removes the specified keyspaces from the given keyspace group.
570+
// Keyspaces in archived or tombstone state will be removed. Keyspaces not in the group will be skipped.
571+
func RemoveKeyspacesFromGroup(c *gin.Context) {
572+
// Parse and validate group ID
573+
groupID, err := validateKeyspaceGroupID(c)
574+
if err != nil {
575+
c.AbortWithStatusJSON(http.StatusBadRequest, "invalid keyspace group id")
576+
return
577+
}
578+
579+
// Parse request body
580+
var params RemoveKeyspacesFromGroupParams
581+
if err := c.BindJSON(&params); err != nil {
582+
c.AbortWithStatusJSON(http.StatusBadRequest, errs.ErrBindJSON.Wrap(err).GenWithStackByCause())
583+
return
584+
}
585+
586+
if len(params.Keyspaces) == 0 {
587+
c.AbortWithStatusJSON(http.StatusBadRequest, "keyspaces list cannot be empty")
588+
return
589+
}
590+
591+
svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
592+
keyspaceManager := svr.GetKeyspaceManager()
593+
if keyspaceManager == nil {
594+
c.AbortWithStatusJSON(http.StatusInternalServerError, managerUninitializedErr)
595+
return
596+
}
597+
598+
groupManager := svr.GetKeyspaceGroupManager()
599+
if groupManager == nil {
600+
c.AbortWithStatusJSON(http.StatusInternalServerError, GroupManagerUninitializedErr)
601+
return
602+
}
603+
604+
// Filter keyspaces: only keep those in ARCHIVED or TOMBSTONE state
605+
var validKeyspaces []uint32
606+
for _, keyspaceID := range params.Keyspaces {
607+
// Load the keyspace meta to check its state
608+
keyspaceMeta, err := keyspaceManager.LoadKeyspaceByID(keyspaceID)
609+
if err != nil {
610+
// Skip if keyspace doesn't exist
611+
if errors.ErrorEqual(err, errs.ErrKeyspaceNotFound) {
612+
continue
613+
}
614+
c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
615+
return
616+
}
617+
618+
// Check if the keyspace is in archived or tombstone state
619+
state := keyspaceMeta.GetState()
620+
if state == keyspacepb.KeyspaceState_ARCHIVED || state == keyspacepb.KeyspaceState_TOMBSTONE {
621+
validKeyspaces = append(validKeyspaces, keyspaceID)
622+
}
623+
}
624+
625+
// If no valid keyspaces to remove, load and return the group without modification
626+
if len(validKeyspaces) == 0 {
627+
kg, err := groupManager.GetKeyspaceGroupByID(groupID)
628+
if err != nil {
629+
c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
630+
return
631+
}
632+
c.IndentedJSON(http.StatusOK, kg)
633+
return
634+
}
635+
636+
// Remove the keyspaces from the keyspace group
637+
kg, err := groupManager.RemoveKeyspacesFromGroup(groupID, validKeyspaces)
638+
if err != nil {
639+
c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
640+
return
641+
}
642+
643+
c.IndentedJSON(http.StatusOK, kg)
644+
}

‎tests/server/apiv2/handlers/testutil.go‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,3 +326,40 @@ func MustMergeKeyspaceGroup(re *require.Assertions, server *tests.TestServer, id
326326
re.NoError(err)
327327
re.Equal(http.StatusOK, resp.StatusCode, string(data))
328328
}
329+
330+
// MustRemoveKeyspacesFromGroup removes keyspaces from a keyspace group with HTTP API.
331+
func MustRemoveKeyspacesFromGroup(re *require.Assertions, server *tests.TestServer, groupID uint32, keyspaceIDs []uint32) *endpoint.KeyspaceGroup {
332+
params := &handlers.RemoveKeyspacesFromGroupParams{
333+
Keyspaces: keyspaceIDs,
334+
}
335+
data, err := json.Marshal(params)
336+
re.NoError(err)
337+
httpReq, err := http.NewRequest(http.MethodDelete, server.GetAddr()+keyspaceGroupsPrefix+fmt.Sprintf("/%d/keyspaces", groupID), bytes.NewBuffer(data))
338+
re.NoError(err)
339+
resp, err := tests.TestDialClient.Do(httpReq)
340+
re.NoError(err)
341+
defer resp.Body.Close()
342+
respData, err := io.ReadAll(resp.Body)
343+
re.NoError(err)
344+
re.Equal(http.StatusOK, resp.StatusCode, string(respData))
345+
var kg endpoint.KeyspaceGroup
346+
re.NoError(json.Unmarshal(respData, &kg))
347+
return &kg
348+
}
349+
350+
// FailRemoveKeyspacesFromGroupWithCode fails to remove keyspaces from a keyspace group with HTTP API.
351+
func FailRemoveKeyspacesFromGroupWithCode(re *require.Assertions, server *tests.TestServer, groupID uint32, keyspaceIDs []uint32, expectCode int) {
352+
params := &handlers.RemoveKeyspacesFromGroupParams{
353+
Keyspaces: keyspaceIDs,
354+
}
355+
data, err := json.Marshal(params)
356+
re.NoError(err)
357+
httpReq, err := http.NewRequest(http.MethodDelete, server.GetAddr()+keyspaceGroupsPrefix+fmt.Sprintf("/%d/keyspaces", groupID), bytes.NewBuffer(data))
358+
re.NoError(err)
359+
resp, err := tests.TestDialClient.Do(httpReq)
360+
re.NoError(err)
361+
defer resp.Body.Close()
362+
respData, err := io.ReadAll(resp.Body)
363+
re.NoError(err)
364+
re.Equal(expectCode, resp.StatusCode, string(respData))
365+
}

‎tests/server/apiv2/handlers/tso_keyspace_group_test.go‎

Lines changed: 113 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,12 @@ import (
2323

2424
"github.com/stretchr/testify/suite"
2525

26-
"github.com/tikv/pd/pkg/mcs/utils/constant"
26+
"github.com/pingcap/failpoint"
27+
"github.com/pingcap/kvproto/pkg/keyspacepb"
28+
29+
"github.com/tikv/pd/pkg/keyspace"
30+
kgconstant "github.com/tikv/pd/pkg/keyspace/constant"
31+
mcsconstant "github.com/tikv/pd/pkg/mcs/utils/constant"
2732
"github.com/tikv/pd/pkg/storage/endpoint"
2833
"github.com/tikv/pd/server/apiv2/handlers"
2934
"github.com/tikv/pd/tests"
@@ -100,7 +105,7 @@ func (suite *keyspaceGroupTestSuite) TestCreateKeyspaceGroups() {
100105
// invalid ID.
101106
kgs = &handlers.CreateKeyspaceGroupParams{KeyspaceGroups: []*endpoint.KeyspaceGroup{
102107
{
103-
ID: constant.MaxKeyspaceGroupCount + 1,
108+
ID: mcsconstant.MaxKeyspaceGroupCount + 1,
104109
UserKind: endpoint.Standard.String(),
105110
},
106111
}}
@@ -141,7 +146,7 @@ func (suite *keyspaceGroupTestSuite) TestSplitKeyspaceGroup() {
141146
ID: uint32(1),
142147
UserKind: endpoint.Standard.String(),
143148
Keyspaces: []uint32{111, 222, 333},
144-
Members: make([]endpoint.KeyspaceGroupMember, constant.DefaultKeyspaceGroupReplicaCount),
149+
Members: make([]endpoint.KeyspaceGroupMember, mcsconstant.DefaultKeyspaceGroupReplicaCount),
145150
},
146151
}}
147152

@@ -248,7 +253,6 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupErrorMessage() {
248253
re.NoError(json.NewDecoder(resp.Body).Decode(&errorMsg))
249254
re.NotEmpty(errorMsg, "Error message should not be empty")
250255
re.Contains(errorMsg, "invalid", "Error message should indicate invalid input")
251-
252256
// Test SetPriorityForKeyspaceGroup with invalid JSON
253257
httpReq, err = http.NewRequest(
254258
http.MethodPatch,
@@ -266,3 +270,108 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupErrorMessage() {
266270
re.NotEmpty(errorMsg, "Error message should not be empty")
267271
re.Contains(errorMsg, "invalid", "Error message should indicate invalid input")
268272
}
273+
274+
func (suite *keyspaceGroupTestSuite) TestRemoveKeyspacesFromGroup() {
275+
re := suite.Require()
276+
re.NoError(failpoint.Enable("github.com/tikv/pd/server/delayStartServerLoop", `return(true)`))
277+
re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)"))
278+
279+
keyspaceManager := suite.server.GetKeyspaceManager()
280+
re.NotNil(keyspaceManager)
281+
282+
// Create test keyspaces (automatically added to default keyspace group 0)
283+
keyspaceMeta1, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
284+
Name: "test_keyspace_1",
285+
CreateTime: 0,
286+
})
287+
re.NoError(err)
288+
keyspaceID1 := keyspaceMeta1.GetId()
289+
290+
keyspaceMeta2, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
291+
Name: "test_keyspace_2",
292+
CreateTime: 0,
293+
})
294+
re.NoError(err)
295+
keyspaceID2 := keyspaceMeta2.GetId()
296+
297+
keyspaceMeta3, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
298+
Name: "test_keyspace_3",
299+
CreateTime: 0,
300+
})
301+
re.NoError(err)
302+
keyspaceID3 := keyspaceMeta3.GetId()
303+
304+
// Verify all keyspaces are in the default group
305+
kg := MustLoadKeyspaceGroupByID(re, suite.server, kgconstant.DefaultKeyspaceGroupID)
306+
re.Contains(kg.Keyspaces, keyspaceID1)
307+
re.Contains(kg.Keyspaces, keyspaceID2)
308+
re.Contains(kg.Keyspaces, keyspaceID3)
309+
310+
// Test 1: Try to remove ENABLED keyspaces (should succeed but nothing removed)
311+
kg = MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID,
312+
[]uint32{keyspaceID1})
313+
// Verify nothing is removed (keyspace is still there because it's ENABLED)
314+
re.Contains(kg.Keyspaces, keyspaceID1)
315+
316+
// Test 2: Update keyspaces to ARCHIVED/TOMBSTONE state and batch remove
317+
// Set keyspace1 to ARCHIVED
318+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID1, keyspacepb.KeyspaceState_DISABLED, 0)
319+
re.NoError(err)
320+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID1, keyspacepb.KeyspaceState_ARCHIVED, 0)
321+
re.NoError(err)
322+
323+
// Set keyspace2 to TOMBSTONE
324+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_DISABLED, 0)
325+
re.NoError(err)
326+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_ARCHIVED, 0)
327+
re.NoError(err)
328+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_TOMBSTONE, 0)
329+
re.NoError(err)
330+
331+
// Batch remove keyspace1 and keyspace2
332+
MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID,
333+
[]uint32{keyspaceID1, keyspaceID2})
334+
335+
// Verify both keyspaces are removed
336+
kg = MustLoadKeyspaceGroupByID(re, suite.server, kgconstant.DefaultKeyspaceGroupID)
337+
re.NotContains(kg.Keyspaces, keyspaceID1)
338+
re.NotContains(kg.Keyspaces, keyspaceID2)
339+
re.Contains(kg.Keyspaces, keyspaceID3) // keyspace3 should still be there
340+
341+
// Test 3: Mix valid and invalid keyspaces
342+
// Set keyspace3 to ARCHIVED
343+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID3, keyspacepb.KeyspaceState_DISABLED, 0)
344+
re.NoError(err)
345+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID3, keyspacepb.KeyspaceState_ARCHIVED, 0)
346+
re.NoError(err)
347+
348+
// Include: valid (keyspace3), already removed (keyspace1), non-existent (99999)
349+
// Should only remove keyspace3, others are skipped
350+
MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID,
351+
[]uint32{keyspaceID3, keyspaceID1, 99999})
352+
353+
// Verify only keyspace3 is removed
354+
kg = MustLoadKeyspaceGroupByID(re, suite.server, kgconstant.DefaultKeyspaceGroupID)
355+
re.NotContains(kg.Keyspaces, keyspaceID3)
356+
357+
// Test 4: Try to remove from non-existent group
358+
FailRemoveKeyspacesFromGroupWithCode(re, suite.server, 999,
359+
[]uint32{keyspaceID1}, http.StatusInternalServerError)
360+
361+
// Test 5: Try to remove with empty keyspace list (should fail - empty list)
362+
FailRemoveKeyspacesFromGroupWithCode(re, suite.server, kgconstant.DefaultKeyspaceGroupID,
363+
[]uint32{}, http.StatusBadRequest)
364+
365+
// Test 6: All keyspaces in wrong state (should succeed but nothing removed)
366+
keyspaceMeta4, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
367+
Name: "test_keyspace_4",
368+
CreateTime: 0,
369+
})
370+
re.NoError(err)
371+
keyspaceID4 := keyspaceMeta4.GetId()
372+
373+
kg = MustRemoveKeyspacesFromGroup(re, suite.server, kgconstant.DefaultKeyspaceGroupID,
374+
[]uint32{keyspaceID4}) // ENABLED state, will be skipped
375+
// Verify keyspace4 is still there
376+
re.Contains(kg.Keyspaces, keyspaceID4)
377+
}

0 commit comments

Comments
 (0)
⚡