From b52e1e8be357f80472ba0783443eb71c0f15e063 Mon Sep 17 00:00:00 2001
From: Xhofe <i@nn.ci>
Date: Thu, 31 Mar 2022 21:52:19 +0800
Subject: [PATCH] fix: load balance

---
 model/account.go                   | 25 +++++++++++++++++++------
 server/common/common.go            |  7 ++++++-
 server/common/files.go             | 12 ++++--------
 server/controllers/file/refresh.go |  4 ++++
 4 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/model/account.go b/model/account.go
index 5c69c6be..ffa66a47 100644
--- a/model/account.go
+++ b/model/account.go
@@ -4,6 +4,7 @@ import (
 	"github.com/Xhofe/alist/conf"
 	"github.com/Xhofe/alist/utils"
 	log "github.com/sirupsen/logrus"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -131,6 +132,7 @@ var balanceMap sync.Map
 // GetBalancedAccount 根据名称获取账号，负载均衡之后的
 func GetBalancedAccount(name string) (Account, bool) {
 	accounts := GetAccountsByPath(name)
+	log.Debugf("accounts: %+v", accounts)
 	accountNum := len(accounts)
 	switch accountNum {
 	case 0:
@@ -203,7 +205,7 @@ func GetAccountsByPath(path string) []Account {
 		name := utils.ParsePath(v.Name)
 		bIndex := strings.LastIndex(name, balance)
 		if bIndex != -1 {
-			name = v.Name[:bIndex]
+			name = name[:bIndex]
 		}
 		// 不是这个账号
 		if path != name && !strings.HasPrefix(path, name+"/") {
@@ -220,18 +222,29 @@ func GetAccountsByPath(path string) []Account {
 		}
 		accounts = append(accounts, v)
 	}
+	sort.Slice(accounts, func(i, j int) bool {
+		return accounts[i].Name < accounts[j].Name
+	})
 	return accounts
 }
 
 // GetAccountFilesByPath 根据路径获取账号虚拟文件
 // 如有账号： /a/b,/a/c,/a/d/e,/a/b.balance1,/av
 // GetAccountFilesByPath(/a) => b,c,d
-func GetAccountFilesByPath(prefix string) ([]File, error) {
+func GetAccountFilesByPath(prefix string) []File {
 	files := make([]File, 0)
-	var accounts []Account
-	if err := conf.DB.Order(columnName("index")).Find(&accounts).Error; err != nil {
-		return nil, err
+	accounts := make([]Account, AccountsCount())
+	i := 0
+	for _, v := range accountsMap {
+		accounts[i] = v
+		i += 1
 	}
+	sort.Slice(accounts, func(i, j int) bool {
+		if accounts[i].Index == accounts[j].Index {
+			return accounts[i].Name < accounts[j].Name
+		}
+		return accounts[i].Index < accounts[j].Index
+	})
 	prefix = utils.ParsePath(prefix)
 	set := make(map[string]interface{})
 	for _, v := range accounts {
@@ -257,5 +270,5 @@ func GetAccountFilesByPath(prefix string) ([]File, error) {
 		})
 		set[name] = nil
 	}
-	return files, nil
+	return files
 }
diff --git a/server/common/common.go b/server/common/common.go
index ff4c7ea4..e829a9ae 100644
--- a/server/common/common.go
+++ b/server/common/common.go
@@ -33,7 +33,12 @@ func ParsePath(rawPath string) (*model.Account, string, base.Driver, error) {
 	if !ok {
 		return nil, "", nil, fmt.Errorf("no [%s] driver", account.Type)
 	}
-	return &account, strings.TrimPrefix(rawPath, utils.ParsePath(account.Name)), driver, nil
+	name := utils.ParsePath(account.Name)
+	bIndex := strings.LastIndex(name, ".balance")
+	if bIndex != -1 {
+		name = name[:bIndex]
+	}
+	return &account, strings.TrimPrefix(rawPath, name), driver, nil
 }
 
 func ErrorResp(c *gin.Context, err error, code int) {
diff --git a/server/common/files.go b/server/common/files.go
index e051732b..cd22e980 100644
--- a/server/common/files.go
+++ b/server/common/files.go
@@ -4,22 +4,21 @@ import (
 	"github.com/Xhofe/alist/drivers/base"
 	"github.com/Xhofe/alist/drivers/operate"
 	"github.com/Xhofe/alist/model"
+	log "github.com/sirupsen/logrus"
 )
 
 func Path(rawPath string) (*model.File, []model.File, *model.Account, base.Driver, string, error) {
 	account, path, driver, err := ParsePath(rawPath)
 	if err != nil {
 		if err.Error() == "path not found" {
-			accountFiles, err := model.GetAccountFilesByPath(rawPath)
-			if err != nil {
-				return nil, nil, nil, nil, "", err
-			}
+			accountFiles := model.GetAccountFilesByPath(rawPath)
 			if len(accountFiles) != 0 {
 				return nil, accountFiles, nil, nil, path, nil
 			}
 		}
 		return nil, nil, nil, nil, "", err
 	}
+	log.Debugln("use account: ", account.Name)
 	file, files, err := operate.Path(driver, account, path)
 	if err != nil {
 		return nil, nil, nil, nil, "", err
@@ -27,10 +26,7 @@ func Path(rawPath string) (*model.File, []model.File, *model.Account, base.Drive
 	if file != nil {
 		return file, nil, account, driver, path, nil
 	} else {
-		accountFiles, err := model.GetAccountFilesByPath(rawPath)
-		if err != nil {
-			return nil, nil, nil, nil, "", err
-		}
+		accountFiles := model.GetAccountFilesByPath(rawPath)
 		files = append(files, accountFiles...)
 		return nil, files, account, driver, path, nil
 	}
diff --git a/server/controllers/file/refresh.go b/server/controllers/file/refresh.go
index 84787cc4..3216c9a2 100644
--- a/server/controllers/file/refresh.go
+++ b/server/controllers/file/refresh.go
@@ -18,6 +18,10 @@ func RefreshFolder(c *gin.Context) {
 	}
 	account, path_, _, err := common.ParsePath(req.Path)
 	if err != nil {
+		if err.Error() == "path not found" && req.Path == "/" {
+			common.SuccessResp(c)
+			return
+		}
 		common.ErrorResp(c, err, 500)
 		return
 	}