Skip to content

Commit 407ad06

Browse files
ystaticyystaticy
authored
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>
1 parent a9d35dc commit 407ad06

File tree

4 files changed

+292
-0
lines changed

4 files changed

+292
-0
lines changed

‎pkg/keyspace/tso_keyspace_group.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,68 @@ func (m *GroupManager) GetGroupByKeyspaceID(id uint32) (uint32, error) {
411411
return 0, ErrKeyspaceNotInAnyKeyspaceGroup
412412
}
413413

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

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

‎server/apiv2/handlers/tso_keyspace_group.go‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ import (
2121

2222
"github.com/gin-gonic/gin"
2323
"github.com/pingcap/errors"
24+
"github.com/pingcap/kvproto/pkg/keyspacepb"
2425
"github.com/tikv/pd/pkg/errs"
26+
"github.com/tikv/pd/pkg/keyspace"
2527
"github.com/tikv/pd/pkg/mcs/utils/constant"
2628
"github.com/tikv/pd/pkg/slice"
2729
"github.com/tikv/pd/pkg/storage/endpoint"
@@ -48,6 +50,7 @@ func RegisterTSOKeyspaceGroup(r *gin.RouterGroup) {
4850
router.DELETE("/:id/split", FinishSplitKeyspaceByID)
4951
router.POST("/:id/merge", MergeKeyspaceGroups)
5052
router.DELETE("/:id/merge", FinishMergeKeyspaceByID)
53+
router.DELETE("/:id/keyspaces", RemoveKeyspacesFromGroup)
5154
}
5255

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

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

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

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

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@ import (
1919
"net/http"
2020
"testing"
2121

22+
"github.com/pingcap/failpoint"
23+
"github.com/pingcap/kvproto/pkg/keyspacepb"
2224
"github.com/stretchr/testify/suite"
25+
"github.com/tikv/pd/pkg/keyspace"
2326
"github.com/tikv/pd/pkg/mcs/utils/constant"
2427
"github.com/tikv/pd/pkg/storage/endpoint"
2528
"github.com/tikv/pd/server/apiv2/handlers"
@@ -173,3 +176,108 @@ func (suite *keyspaceGroupTestSuite) TestSplitKeyspaceGroup() {
173176
kg2 = MustLoadKeyspaceGroupByID(re, suite.server, 2)
174177
re.False(kg2.IsSplitting())
175178
}
179+
180+
func (suite *keyspaceGroupTestSuite) TestRemoveKeyspacesFromGroup() {
181+
re := suite.Require()
182+
re.NoError(failpoint.Enable("github.com/tikv/pd/server/delayStartServerLoop", `return(true)`))
183+
re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)"))
184+
185+
keyspaceManager := suite.server.GetKeyspaceManager()
186+
re.NotNil(keyspaceManager)
187+
188+
// Create test keyspaces (automatically added to default keyspace group 0)
189+
keyspaceMeta1, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
190+
Name: "test_keyspace_1",
191+
CreateTime: 0,
192+
})
193+
re.NoError(err)
194+
keyspaceID1 := keyspaceMeta1.GetId()
195+
196+
keyspaceMeta2, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
197+
Name: "test_keyspace_2",
198+
CreateTime: 0,
199+
})
200+
re.NoError(err)
201+
keyspaceID2 := keyspaceMeta2.GetId()
202+
203+
keyspaceMeta3, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
204+
Name: "test_keyspace_3",
205+
CreateTime: 0,
206+
})
207+
re.NoError(err)
208+
keyspaceID3 := keyspaceMeta3.GetId()
209+
210+
// Verify all keyspaces are in the default group
211+
kg := MustLoadKeyspaceGroupByID(re, suite.server, constant.DefaultKeyspaceGroupID)
212+
re.Contains(kg.Keyspaces, keyspaceID1)
213+
re.Contains(kg.Keyspaces, keyspaceID2)
214+
re.Contains(kg.Keyspaces, keyspaceID3)
215+
216+
// Test 1: Try to remove ENABLED keyspaces (should succeed but nothing removed)
217+
kg = MustRemoveKeyspacesFromGroup(re, suite.server, constant.DefaultKeyspaceGroupID,
218+
[]uint32{keyspaceID1})
219+
// Verify nothing is removed (keyspace is still there because it's ENABLED)
220+
re.Contains(kg.Keyspaces, keyspaceID1)
221+
222+
// Test 2: Update keyspaces to ARCHIVED/TOMBSTONE state and batch remove
223+
// Set keyspace1 to ARCHIVED
224+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID1, keyspacepb.KeyspaceState_DISABLED, 0)
225+
re.NoError(err)
226+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID1, keyspacepb.KeyspaceState_ARCHIVED, 0)
227+
re.NoError(err)
228+
229+
// Set keyspace2 to TOMBSTONE
230+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_DISABLED, 0)
231+
re.NoError(err)
232+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_ARCHIVED, 0)
233+
re.NoError(err)
234+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID2, keyspacepb.KeyspaceState_TOMBSTONE, 0)
235+
re.NoError(err)
236+
237+
// Batch remove keyspace1 and keyspace2
238+
MustRemoveKeyspacesFromGroup(re, suite.server, constant.DefaultKeyspaceGroupID,
239+
[]uint32{keyspaceID1, keyspaceID2})
240+
241+
// Verify both keyspaces are removed
242+
kg = MustLoadKeyspaceGroupByID(re, suite.server, constant.DefaultKeyspaceGroupID)
243+
re.NotContains(kg.Keyspaces, keyspaceID1)
244+
re.NotContains(kg.Keyspaces, keyspaceID2)
245+
re.Contains(kg.Keyspaces, keyspaceID3) // keyspace3 should still be there
246+
247+
// Test 3: Mix valid and invalid keyspaces
248+
// Set keyspace3 to ARCHIVED
249+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID3, keyspacepb.KeyspaceState_DISABLED, 0)
250+
re.NoError(err)
251+
_, err = keyspaceManager.UpdateKeyspaceStateByID(keyspaceID3, keyspacepb.KeyspaceState_ARCHIVED, 0)
252+
re.NoError(err)
253+
254+
// Include: valid (keyspace3), already removed (keyspace1), non-existent (99999)
255+
// Should only remove keyspace3, others are skipped
256+
MustRemoveKeyspacesFromGroup(re, suite.server, constant.DefaultKeyspaceGroupID,
257+
[]uint32{keyspaceID3, keyspaceID1, 99999})
258+
259+
// Verify only keyspace3 is removed
260+
kg = MustLoadKeyspaceGroupByID(re, suite.server, constant.DefaultKeyspaceGroupID)
261+
re.NotContains(kg.Keyspaces, keyspaceID3)
262+
263+
// Test 4: Try to remove from non-existent group
264+
FailRemoveKeyspacesFromGroupWithCode(re, suite.server, 999,
265+
[]uint32{keyspaceID1}, http.StatusInternalServerError)
266+
267+
// Test 5: Try to remove with empty keyspace list (should fail - empty list)
268+
FailRemoveKeyspacesFromGroupWithCode(re, suite.server, constant.DefaultKeyspaceGroupID,
269+
[]uint32{}, http.StatusBadRequest)
270+
271+
// Test 6: All keyspaces in wrong state (should succeed but nothing removed)
272+
keyspaceMeta4, err := keyspaceManager.CreateKeyspace(&keyspace.CreateKeyspaceRequest{
273+
Name: "test_keyspace_4",
274+
CreateTime: 0,
275+
})
276+
re.NoError(err)
277+
keyspaceID4 := keyspaceMeta4.GetId()
278+
279+
kg = MustRemoveKeyspacesFromGroup(re, suite.server, constant.DefaultKeyspaceGroupID,
280+
[]uint32{keyspaceID4}) // ENABLED state, will be skipped
281+
// Verify keyspace4 is still there
282+
re.Contains(kg.Keyspaces, keyspaceID4)
283+
}

0 commit comments

Comments
 (0)
⚡