JCS-pub/coordinator/internal/rpc/storage.go

package rpc

import (
	"context"

	"gitlink.org.cn/cloudream/common/consts/errorcode"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc"
	corrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/coordinator"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/db"
	cortypes "gitlink.org.cn/cloudream/jcs-pub/coordinator/types"
)

func (svc *Service) SelectStorageHub(ctx context.Context, msg *corrpc.SelectStorageHub) (*corrpc.SelectStorageHubResp, *rpc.CodeError) {
	d := svc.db
	resp, err := db.DoTx02(d, func(tx db.SQLContext) ([]*cortypes.Hub, error) {
		allLoc, err := d.HubLocation().GetAll(tx)
		if err != nil {
			return nil, err
		}

		stgHubIDs := make([]cortypes.HubID, 0, len(msg.Storages))
		for _, stg := range msg.Storages {
			stgLoc := stg.GetLocation()

			var matchedHubID cortypes.HubID
			var matchedScore int
			for _, loc := range allLoc {
				sc := matchLocation(stgLoc, loc)
				if sc > matchedScore {
					matchedScore = sc
					matchedHubID = loc.HubID
				}
			}

			stgHubIDs = append(stgHubIDs, matchedHubID)
		}

		hubs, err := d.Hub().BatchGetByID(tx, stgHubIDs)
		if err != nil {
			return nil, err
		}

		hubMap := make(map[cortypes.HubID]*cortypes.Hub)
		for _, hub := range hubs {
			h := hub
			hubMap[hub.HubID] = &h
		}

		resp := make([]*cortypes.Hub, len(msg.Storages))
		for i := range msg.Storages {
			resp[i] = hubMap[stgHubIDs[i]]
		}
		return resp, nil
	})

	if err != nil {
		logger.Warnf("select storage hubs: %s", err.Error())
		return nil, rpc.Failed(errorcode.OperationFailed, "%v", err)
	}

	return &corrpc.SelectStorageHubResp{
		Hubs: resp,
	}, nil
}

// 匹配规则：
// 1. 按照StorageName、Location顺序检查StorageLocation和HubLocation
// 2. "*"代表通配符，匹配任意值，如果匹配到了通配，那么就直接结束匹配
// 3. 匹配越精确，分数越高
func matchLocation(loc cortypes.Location, hubLoc cortypes.HubLocation) int {
	if hubLoc.StorageName == "*" {
		return 1
	}

	if hubLoc.StorageName != loc.StorageName {
		return 0
	}

	if hubLoc.Location == "*" {
		return 2
	}

	if hubLoc.Location == loc.Location {
		return 3
	}

	return 0
}