mcs, tso: fix Nil pointer deference when (*AllocatorManager).GetMember

pkg/tso/keyspace_group_manager.go

@@ -130,13 +130,17 @@ func (s *state) getAMWithMembershipCheck(
 		return nil, kgid, genNotServedErr(errs.ErrGetAllocatorManager, keyspaceGroupID)
 	}
 
+	// The keyspace doesn't belong to any keyspace group but the keyspace has been assigned to a
+	// keyspace group before, which means the keyspace group hasn't initialized yet.
 	if keyspaceGroupID != mcsutils.DefaultKeyspaceGroupID {
 		return nil, keyspaceGroupID, errs.ErrKeyspaceNotAssigned.FastGenByArgs(keyspaceID)
 	}
 
-	// The keyspace doesn't belong to any keyspace group, so return the default keyspace group.
-	// It's for migrating the existing keyspaces which have no keyspace group assigned, so the
-	// the default keyspace group is used to serve the keyspaces.
+	// For migrating the existing keyspaces which have no keyspace group assigned as configured in the
+	// keyspace meta. All these keyspaces will be served by the default keyspace group.
+	if s.ams[mcsutils.DefaultKeyspaceGroupID] == nil {
+		return nil, mcsutils.DefaultKeyspaceGroupID, errs.ErrKeyspaceNotAssigned.FastGenByArgs(keyspaceID)
+	}
 	return s.ams[mcsutils.DefaultKeyspaceGroupID], mcsutils.DefaultKeyspaceGroupID, nil
 }
 

tests/integrations/mcs/tso/keyspace_group_manager_test.go

@@ -26,6 +26,7 @@ import (
 	pd "github.com/tikv/pd/client"
 	"github.com/tikv/pd/client/testutil"
 	"github.com/tikv/pd/pkg/election"
+	mcsutils "github.com/tikv/pd/pkg/mcs/utils"
 	"github.com/tikv/pd/pkg/member"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	tsopkg "github.com/tikv/pd/pkg/tso"
@@ -86,6 +87,47 @@ func cleanupKeyspaceGroups(re *require.Assertions, server *tests.TestServer) {
 	}
 }
 
+func (suite *tsoKeyspaceGroupManagerTestSuite) TestKeyspacesServedByDefaultKeyspaceGroup() {
+	// There is only default keyspace group. All keyspaces which have not been assigned to
+	// any keyspace group will be served by the default keyspace group; otherwise,
+	// they won't be served by any keyspace group.
+	re := suite.Require()
+	testutil.Eventually(re, func() bool {
+		for _, server := range suite.tsoCluster.GetServers() {
+			allServed := true
+			for _, keyspaceID := range []uint32{0, 1, 2} {
+				if server.IsKeyspaceServing(keyspaceID, mcsutils.DefaultKeyspaceGroupID) {
+					tam, err := server.GetTSOAllocatorManager(mcsutils.DefaultKeyspaceGroupID)
+					re.NoError(err)
+					re.NotNil(tam)
+				} else {
+					allServed = false
+				}
+			}
+			return allServed
+		}
+		return false
+	}, testutil.WithWaitFor(5*time.Second), testutil.WithTickInterval(50*time.Millisecond))
+
+	// All keyspaces which have been assigned to any keyspace group before, except default
+	// keyspace, won't be served at this time. Default keyspace will be served by default
+	// keyspace group all the time.
+	for _, server := range suite.tsoCluster.GetServers() {
+		server.IsKeyspaceServing(mcsutils.DefaultKeyspaceID, mcsutils.DefaultKeyspaceGroupID)
+		for _, keyspaceGroupID := range []uint32{1, 2, 3} {
+			server.IsKeyspaceServing(mcsutils.DefaultKeyspaceID, keyspaceGroupID)
+			server.IsKeyspaceServing(mcsutils.DefaultKeyspaceID, keyspaceGroupID)
+			for _, keyspaceID := range []uint32{1, 2, 3} {
+				if server.IsKeyspaceServing(keyspaceID, keyspaceGroupID) {
+					tam, err := server.GetTSOAllocatorManager(keyspaceGroupID)
+					re.NoError(err)
+					re.NotNil(tam)
+				}
+			}
+		}
+	}
+}
+
 func (suite *tsoKeyspaceGroupManagerTestSuite) TestTSOKeyspaceGroupSplit() {
 	re := suite.Require()
 	// Create the keyspace group 1 with keyspaces [111, 222, 333].