From df4ac32e4757a37263c43aeb4440ed4c200fd26e Mon Sep 17 00:00:00 2001
From: lhy1024 <admin@liudos.us>
Date: Mon, 18 Nov 2024 12:00:50 +0800
Subject: [PATCH] refactor

Signed-off-by: lhy1024 <admin@liudos.us>
---
 pkg/storage/storage.go | 46 ++++++++++++++++++++++++++----------------
 1 file changed, 29 insertions(+), 17 deletions(-)

diff --git a/pkg/storage/storage.go b/pkg/storage/storage.go
index 0bb5f7c76ba..9679732af17 100644
--- a/pkg/storage/storage.go
+++ b/pkg/storage/storage.go
@@ -73,14 +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
-	regionLoadedFromDefault bool
-	regionLoadedFromStorage bool
-	mu                      syncutil.RWMutex
+	useRegionStorageFlag int32
+	regionLoaded         regionSource
+	mu                   syncutil.RWMutex
 }
 
 // NewCoreStorage creates a new core storage with the given default and region storage.
@@ -91,6 +98,7 @@ func NewCoreStorage(defaultStorage Storage, regionStorage endpoint.RegionStorage
 	return &coreStorage{
 		Storage:       defaultStorage,
 		regionStorage: regionStorage,
+		regionLoaded:  unloaded,
 	}
 }
 
@@ -117,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)
+		atomic.StoreInt32(&ps.useRegionStorageFlag, 1)
 		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)
+	atomic.StoreInt32(&ps.useRegionStorageFlag, 0)
 	return ps.Storage
 }
 
@@ -137,26 +145,26 @@ func TryLoadRegionsOnce(ctx context.Context, s Storage, f func(region *core.Regi
 	ps.mu.Lock()
 	defer ps.mu.Unlock()
 
-	if atomic.LoadInt32(&ps.useRegionStorage) == 0 {
+	if !ps.useRegionStorage() {
 		err := ps.Storage.LoadRegions(ctx, f)
 		if err == nil {
-			ps.regionLoadedFromDefault = true
+			ps.regionLoaded = fromEtcd
 		}
 		return err
 	}
 
-	if !ps.regionLoadedFromStorage {
+	if ps.regionLoaded == unloaded {
 		if err := ps.regionStorage.LoadRegions(ctx, f); err != nil {
 			return err
 		}
-		ps.regionLoadedFromStorage = 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() {
 		return ps.regionStorage.LoadRegion(regionID, region)
 	}
 	return ps.Storage.LoadRegion(regionID, region)
@@ -164,7 +172,7 @@ func (ps *coreStorage) LoadRegion(regionID uint64, region *metapb.Region) (ok bo
 
 // 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() {
 		return ps.regionStorage.LoadRegions(ctx, f)
 	}
 	return ps.Storage.LoadRegions(ctx, f)
@@ -172,7 +180,7 @@ func (ps *coreStorage) LoadRegions(ctx context.Context, f func(region *core.Regi
 
 // SaveRegion saves one region to storage.
 func (ps *coreStorage) SaveRegion(region *metapb.Region) error {
-	if atomic.LoadInt32(&ps.useRegionStorage) > 0 {
+	if ps.useRegionStorage() {
 		return ps.regionStorage.SaveRegion(region)
 	}
 	return ps.Storage.SaveRegion(region)
@@ -180,7 +188,7 @@ func (ps *coreStorage) SaveRegion(region *metapb.Region) error {
 
 // DeleteRegion deletes one region from storage.
 func (ps *coreStorage) DeleteRegion(region *metapb.Region) error {
-	if atomic.LoadInt32(&ps.useRegionStorage) > 0 {
+	if ps.useRegionStorage() {
 		return ps.regionStorage.DeleteRegion(region)
 	}
 	return ps.Storage.DeleteRegion(region)
@@ -204,13 +212,17 @@ func (ps *coreStorage) Close() error {
 	return nil
 }
 
+func (ps *coreStorage) useRegionStorage() bool {
+	return atomic.LoadInt32(&ps.useRegionStorageFlag) > 0
+}
+
 // AreRegionsLoaded returns whether the regions are loaded.
 func AreRegionsLoaded(s Storage) bool {
 	ps := s.(*coreStorage)
 	ps.mu.RLock()
 	defer ps.mu.RUnlock()
-	if atomic.LoadInt32(&ps.useRegionStorage) == 0 {
-		return ps.regionLoadedFromDefault
+	if ps.useRegionStorage() {
+		return ps.regionLoaded == fromLeveldb
 	}
-	return ps.regionLoadedFromStorage
+	return ps.regionLoaded == fromEtcd
 }