pcm-coordinator/internal/logic/schedule/queryresourceslogic.go

package schedule

import (
	"context"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/collector"
	"sync"
	"time"

	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"

	"github.com/zeromicro/go-zero/core/logx"
)

type QueryResourcesLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewQueryResourcesLogic(ctx context.Context, svcCtx *svc.ServiceContext) *QueryResourcesLogic {
	return &QueryResourcesLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *QueryResourcesLogic) QueryResources(req *types.QueryResourcesReq) (resp *types.QueryResourcesResp, err error) {
	resp = &types.QueryResourcesResp{}
	var ulist []*collector.ResourceUsage

	clusters, err := l.svcCtx.Scheduler.AiStorages.GetClustersByAdapterId(req.AdapterId)
	if err != nil {
		return nil, err
	}

	var ch = make(chan *collector.ResourceUsage, len(clusters.List))

	var wg sync.WaitGroup
	for _, cluster := range clusters.List {
		wg.Add(1)
		c := cluster
		go func() {
			defer wg.Done()

			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()

			col, found := l.svcCtx.Scheduler.AiService.AiCollectorAdapterMap[req.AdapterId][c.Id]
			if !found {
				return
			}
			usage, err := col.GetResourceUsage(ctx)
			if err != nil {
				return
			}
			ch <- usage
		}()
	}
	wg.Wait()
	close(ch)

	for v := range ch {
		ulist = append(ulist, v)
	}

	// handle empty usage
	rus := handleEmptyResourceUsage(clusters.List, ulist)
	resp.Data = rus

	return resp, nil
}

func handleEmptyResourceUsage(list []types.ClusterInfo, ulist []*collector.ResourceUsage) []*collector.ResourceUsage {
	var rus []*collector.ResourceUsage
	m := make(map[string]interface{})
	for _, u := range ulist {
		if u == nil {
			continue
		}
		m[u.ClusterId] = u
	}

	for _, l := range list {
		s, ok := m[l.Id]
		if !ok {
			ru := &collector.ResourceUsage{
				ClusterId: l.Id,
				Usages:    nil,
				Msg:       "resources unavailable, please retry later",
			}
			rus = append(rus, ru)
		} else {
			if s == nil {
				ru := &collector.ResourceUsage{
					ClusterId: l.Id,
					Usages:    nil,
					Msg:       "resources unavailable, please retry later",
				}
				rus = append(rus, ru)
			} else {
				r, ok := s.(*collector.ResourceUsage)
				if ok {
					if r.Usages == nil {
						ru := &collector.ResourceUsage{
							ClusterId: r.ClusterId,
							Usages:    nil,
							Msg:       "resources unavailable, please retry later",
						}
						rus = append(rus, ru)
					} else {
						rus = append(rus, r)
					}
				}
			}
		}
	}

	return rus
}