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"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/db"
)

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

		stgHubIDs := make([]jcstypes.HubID, 0, len(msg.Storages))
		for _, stg := range msg.Storages {
			var matchedHubID jcstypes.HubID
			switch stg := stg.(type) {
			case *jcstypes.PubShardsType:
				matchedHubID = stg.MasterHub

			default:
				stgLoc := stg.GetLocation()
				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[jcstypes.HubID]*jcstypes.Hub)
		for _, hub := range hubs {
			h := hub
			hubMap[hub.HubID] = &h
		}

		resp := make([]*jcstypes.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 jcstypes.Location, hubLoc jcstypes.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
}