api: support to query whether pd has loaded region (#8749) (#9319)

close #8426, close #8748

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>
Signed-off-by: lhy1024 <admin@liudos.us>
Signed-off-by: artem_danilov <artem.danilov@airbnb.com>

Co-authored-by: lhy1024 <liuhanyang@pingcap.com>
Co-authored-by: lhy1024 <admin@liudos.us>
Co-authored-by: Artem Danilov <artem.danilov@airbnb.com>

Author: 3 people

pkg/storage/storage.go

@@ -18,13 +18,14 @@ import (
 	"context"
 	"sync/atomic"
 
+	clientv3 "go.etcd.io/etcd/client/v3"
+
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/encryption"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/syncutil"
-	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
 // Storage is the interface for the backend storage of the PD.
@@ -73,13 +74,21 @@ func NewRegionStorageWithLevelDBBackend(
 
 // TODO: support other KV storage backends like BadgerDB in the future.
 
+type regionSource int
+
+const (
+	unloaded regionSource = iota
+	fromEtcd
+	fromLeveldb
+)
+
 type coreStorage struct {
 	Storage
 	regionStorage endpoint.RegionStorage
 
-	useRegionStorage int32
-	regionLoaded     bool
-	mu               syncutil.Mutex
+	useRegionStorage atomic.Bool
+	regionLoaded     regionSource
+	mu               syncutil.RWMutex
 }
 
 // NewCoreStorage creates a new core storage with the given default and region storage.
@@ -90,6 +99,7 @@ func NewCoreStorage(defaultStorage Storage, regionStorage endpoint.RegionStorage
 	return &coreStorage{
 		Storage:       defaultStorage,
 		regionStorage: regionStorage,
+		regionLoaded:  unloaded,
 	}
 }
 
@@ -116,12 +126,12 @@ func TrySwitchRegionStorage(s Storage, useLocalRegionStorage bool) endpoint.Regi
 	if useLocalRegionStorage {
 		// Switch the region storage to regionStorage, all region info will be read/saved by the internal
 		// regionStorage, and in most cases it's LevelDB-backend.
-		atomic.StoreInt32(&ps.useRegionStorage, 1)
+		ps.useRegionStorage.Store(true)
 		return ps.regionStorage
 	}
 	// Switch the region storage to defaultStorage, all region info will be read/saved by the internal
 	// defaultStorage, and in most cases it's etcd-backend.
-	atomic.StoreInt32(&ps.useRegionStorage, 0)
+	ps.useRegionStorage.Store(false)
 	return ps.Storage
 }
 
@@ -133,48 +143,53 @@ func TryLoadRegionsOnce(ctx context.Context, s Storage, f func(region *core.Regi
 		return s.LoadRegions(ctx, f)
 	}
 
-	if atomic.LoadInt32(&ps.useRegionStorage) == 0 {
-		return ps.Storage.LoadRegions(ctx, f)
-	}
-
 	ps.mu.Lock()
 	defer ps.mu.Unlock()
-	if !ps.regionLoaded {
+
+	if !ps.useRegionStorage.Load() {
+		err := ps.Storage.LoadRegions(ctx, f)
+		if err == nil {
+			ps.regionLoaded = fromEtcd
+		}
+		return err
+	}
+
+	if ps.regionLoaded == unloaded {
 		if err := ps.regionStorage.LoadRegions(ctx, f); err != nil {
 			return err
 		}
-		ps.regionLoaded = true
+		ps.regionLoaded = fromLeveldb
 	}
 	return nil
 }
 
 // LoadRegion loads one region from storage.
 func (ps *coreStorage) LoadRegion(regionID uint64, region *metapb.Region) (ok bool, err error) {
-	if atomic.LoadInt32(&ps.useRegionStorage) > 0 {
+	if ps.useRegionStorage.Load() {
 		return ps.regionStorage.LoadRegion(regionID, region)
 	}
 	return ps.Storage.LoadRegion(regionID, region)
 }
 
 // LoadRegions loads all regions from storage to RegionsInfo.
 func (ps *coreStorage) LoadRegions(ctx context.Context, f func(region *core.RegionInfo) []*core.RegionInfo) error {
-	if atomic.LoadInt32(&ps.useRegionStorage) > 0 {
+	if ps.useRegionStorage.Load() {
 		return ps.regionStorage.LoadRegions(ctx, f)
 	}
 	return ps.Storage.LoadRegions(ctx, f)
 }
 
 // SaveRegion saves one region to storage.
 func (ps *coreStorage) SaveRegion(region *metapb.Region) error {
-	if atomic.LoadInt32(&ps.useRegionStorage) > 0 {
+	if ps.useRegionStorage.Load() {
 		return ps.regionStorage.SaveRegion(region)
 	}
 	return ps.Storage.SaveRegion(region)
 }
 
 // DeleteRegion deletes one region from storage.
 func (ps *coreStorage) DeleteRegion(region *metapb.Region) error {
-	if atomic.LoadInt32(&ps.useRegionStorage) > 0 {
+	if ps.useRegionStorage.Load() {
 		return ps.regionStorage.DeleteRegion(region)
 	}
 	return ps.Storage.DeleteRegion(region)
@@ -197,3 +212,14 @@ func (ps *coreStorage) Close() error {
 	}
 	return nil
 }
+
+// AreRegionsLoaded returns whether the regions are loaded.
+func AreRegionsLoaded(s Storage) bool {
+	ps := s.(*coreStorage)
+	ps.mu.RLock()
+	defer ps.mu.RUnlock()
+	if ps.useRegionStorage.Load() {
+		return ps.regionLoaded == fromLeveldb
+	}
+	return ps.regionLoaded == fromEtcd
+}

server/api/health_test.go

@@ -49,22 +49,22 @@ func TestHealthSlice(t *testing.T) {
 	re := require.New(t)
 	cfgs, svrs, clean := mustNewCluster(re, 3)
 	defer clean()
-	var leader, follow *server.Server
+	var leader, follower *server.Server
 
 	for _, svr := range svrs {
 		if !svr.IsClosed() && svr.GetMember().IsLeader() {
 			leader = svr
 		} else {
-			follow = svr
+			follower = svr
 		}
 	}
 	mustBootstrapCluster(re, leader)
 	addr := leader.GetConfig().ClientUrls + apiPrefix + "/api/v1/health"
-	follow.Close()
+	follower.Close()
 	resp, err := testDialClient.Get(addr)
 	re.NoError(err)
 	defer resp.Body.Close()
 	buf, err := io.ReadAll(resp.Body)
 	re.NoError(err)
-	checkSliceResponse(re, buf, cfgs, follow.GetConfig().Name)
+	checkSliceResponse(re, buf, cfgs, follower.GetConfig().Name)
 }

server/apiv2/handlers/ready.go

@@ -0,0 +1,69 @@
+// Copyright 2024 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package handlers
+
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/pingcap/failpoint"
+
+	"github.com/tikv/pd/pkg/storage"
+	"github.com/tikv/pd/server"
+	"github.com/tikv/pd/server/apiv2/middlewares"
+)
+
+// ReadyStatus reflects the cluster's ready status.
+// NOTE: This type is exported by HTTP API. Please pay more attention when modifying it.
+type ReadyStatus struct {
+	RegionLoaded bool `json:"region_loaded"`
+}
+
+// Ready checks if the region is loaded.
+// @Summary  It will return whether pd follower is ready to became leader. This request is always served by the instance that receives it and never forwarded to the leader.
+// @Router   /ready [get]
+// @Param    verbose query  bool    false  "Whether to return details."
+// @Success  200
+// @Failure  500
+func Ready(c *gin.Context) {
+	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
+	s := svr.GetStorage()
+	regionLoaded := storage.AreRegionsLoaded(s)
+	failpoint.Inject("loadRegionSlow", func(val failpoint.Value) {
+		if s, ok := val.(string); ok {
+			if svr.GetAddr() == s {
+				regionLoaded = false
+			}
+		}
+	})
+	if regionLoaded {
+		c.Status(http.StatusOK)
+	} else {
+		c.Status(http.StatusInternalServerError)
+	}
+
+	if _, ok := c.GetQuery("verbose"); !ok {
+		return
+	}
+	resp := &ReadyStatus{
+		RegionLoaded: regionLoaded,
+	}
+	if regionLoaded {
+		c.IndentedJSON(http.StatusOK, resp)
+	} else {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, resp)
+	}
+}

server/apiv2/router.go

@@ -50,8 +50,10 @@ func NewV2Handler(_ context.Context, svr *server.Server) (http.Handler, apiutil.
 		c.Set(middlewares.ServerContextKey, svr)
 		c.Next()
 	})
-	router.Use(middlewares.Redirector())
 	root := router.Group(apiV2Prefix)
+	// add ready handler before Redirector to avoid redirecting it to the leader
+	root.GET("ready", handlers.Ready)
+	root.Use(middlewares.Redirector())
 	handlers.RegisterKeyspace(root)
 	handlers.RegisterTSOKeyspaceGroup(root)
 	handlers.RegisterMicroService(root)

tests/server/apiv2/handlers/ready_test.go

@@ -0,0 +1,89 @@
+// Copyright 2024 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package handlers
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/pingcap/failpoint"
+
+	"github.com/tikv/pd/server/apiv2/handlers"
+	"github.com/tikv/pd/tests"
+)
+
+func TestReadyAPI(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cluster, err := tests.NewTestCluster(ctx, 3)
+	re.NoError(err)
+	defer cluster.Destroy()
+	re.NoError(cluster.RunInitialServers())
+	re.NotEmpty(cluster.WaitLeader())
+	leader := cluster.GetLeaderServer()
+	re.NoError(leader.BootstrapCluster())
+	leaderURL := leader.GetConfig().ClientUrls + v2Prefix + "/ready"
+	followerServer := cluster.GetServer(cluster.GetFollower())
+	followerURL := followerServer.GetConfig().ClientUrls + v2Prefix + "/ready"
+	checkReadyAPI(re, leaderURL, true)
+	checkReadyAPI(re, followerURL, true)
+
+	// check ready status when region is not loaded for leader
+	failpoint.Enable("github.com/tikv/pd/server/apiv2/handlers/loadRegionSlow", `return("`+leader.GetAddr()+`")`)
+	checkReadyAPI(re, leaderURL, false)
+	checkReadyAPI(re, followerURL, true)
+
+	// check ready status when region is not loaded for follower
+	failpoint.Enable("github.com/tikv/pd/server/apiv2/handlers/loadRegionSlow", `return("`+followerServer.GetAddr()+`")`)
+	checkReadyAPI(re, leaderURL, true)
+	checkReadyAPI(re, followerURL, false)
+
+	failpoint.Disable("github.com/tikv/pd/server/apiv2/handlers/loadRegionSlow")
+}
+
+func checkReadyAPI(re *require.Assertions, url string, isReady bool) {
+	expectCode := http.StatusOK
+	if !isReady {
+		expectCode = http.StatusInternalServerError
+	}
+	// check ready status
+	req, err := http.NewRequest(http.MethodGet, url, http.NoBody)
+	re.NoError(err)
+	resp, err := tests.TestDialClient.Do(req)
+	re.NoError(err)
+	defer resp.Body.Close()
+	buf, err := io.ReadAll(resp.Body)
+	re.NoError(err)
+	re.Empty(buf)
+	re.Equal(expectCode, resp.StatusCode)
+	// check ready status with verbose
+	req, err = http.NewRequest(http.MethodGet, url+"?verbose", http.NoBody)
+	re.NoError(err)
+	resp, err = tests.TestDialClient.Do(req)
+	re.NoError(err)
+	defer resp.Body.Close()
+	buf, err = io.ReadAll(resp.Body)
+	re.NoError(err)
+	r := &handlers.ReadyStatus{}
+	re.NoError(json.Unmarshal(buf, &r))
+	re.Equal(expectCode, resp.StatusCode)
+	re.Equal(isReady, r.RegionLoaded)
+}

tests/server/apiv2/handlers/testutil.go

@@ -30,6 +30,7 @@ import (
 )
 
 const (
+	v2Prefix             = "/pd/api/v2"
 	keyspacesPrefix      = "/pd/api/v2/keyspaces"
 	keyspaceGroupsPrefix = "/pd/api/v2/tso/keyspace-groups"
 )