api: support to query whether pd has loaded region

pkg/storage/storage.go

@@ -73,13 +73,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 +98,7 @@ func NewCoreStorage(defaultStorage Storage, regionStorage endpoint.RegionStorage
 	return &coreStorage{
 		Storage:       defaultStorage,
 		regionStorage: regionStorage,
+		regionLoaded:  unloaded,
 	}
 }
 
@@ -116,12 +125,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 +142,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 +211,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
+}

pkg/versioninfo/versioninfo.go

@@ -31,6 +31,7 @@ type Status struct {
 	Version        string `json:"version"`
 	GitHash        string `json:"git_hash"`
 	StartTimestamp int64  `json:"start_timestamp"`
+	RegionLoaded   bool   `json:"region_loaded"`
 }
 
 const (

server/api/status.go

@@ -17,6 +17,7 @@ package api
 import (
 	"net/http"
 
+	"github.com/tikv/pd/pkg/storage"
 	"github.com/tikv/pd/pkg/versioninfo"
 	"github.com/tikv/pd/server"
 	"github.com/unrolled/render"
@@ -44,6 +45,7 @@ func (h *statusHandler) GetPDStatus(w http.ResponseWriter, _ *http.Request) {
 		GitHash:        versioninfo.PDGitHash,
 		Version:        versioninfo.PDReleaseVersion,
 		StartTimestamp: h.svr.StartTimestamp(),
+		RegionLoaded:   storage.AreRegionsLoaded(h.svr.GetStorage()),
 	}
 
 	h.rd.JSON(w, http.StatusOK, version)

server/api/status_test.go

@@ -23,12 +23,11 @@ import (
 	"github.com/tikv/pd/pkg/versioninfo"
 )
 
-func checkStatusResponse(re *require.Assertions, body []byte) {
-	got := versioninfo.Status{}
-	re.NoError(json.Unmarshal(body, &got))
-	re.Equal(versioninfo.PDBuildTS, got.BuildTS)
-	re.Equal(versioninfo.PDGitHash, got.GitHash)
-	re.Equal(versioninfo.PDReleaseVersion, got.Version)
+type oldStatus struct {
+	BuildTS        string `json:"build_ts"`
+	Version        string `json:"version"`
+	GitHash        string `json:"git_hash"`
+	StartTimestamp int64  `json:"start_timestamp"`
 }
 
 func TestStatus(t *testing.T) {
@@ -46,3 +45,17 @@ func TestStatus(t *testing.T) {
 		resp.Body.Close()
 	}
 }
+
+func checkStatusResponse(re *require.Assertions, body []byte) {
+	got := versioninfo.Status{}
+	re.NoError(json.Unmarshal(body, &got))
+	re.Equal(versioninfo.PDBuildTS, got.BuildTS)
+	re.Equal(versioninfo.PDGitHash, got.GitHash)
+	re.Equal(versioninfo.PDReleaseVersion, got.Version)
+	re.False(got.RegionLoaded)
+	gotWithOldStatus := oldStatus{}
+	re.NoError(json.Unmarshal(body, &gotWithOldStatus))
+	re.Equal(versioninfo.PDBuildTS, gotWithOldStatus.BuildTS)
+	re.Equal(versioninfo.PDGitHash, gotWithOldStatus.GitHash)
+	re.Equal(versioninfo.PDReleaseVersion, gotWithOldStatus.Version)
+}