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

package schedule

import (
	"context"
	"errors"
	"github.com/zeromicro/go-zero/core/logx"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/collector"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/storeLink"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	ADAPTERID = "1777144940459986944" // 异构适配器id
)

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{}
	rus, err := l.queryResources(req.ClusterIDs)
	if err != nil {
		return nil, err
	}

	resp.Data = rus

	return resp, nil
}

func (l *QueryResourcesLogic) queryResources(clusterIds []string) ([]*collector.ResourceSpec, error) {
	var ulist []*collector.ResourceSpec

	var clusters []types.ClusterInfo
	if len(clusterIds) != 0 {
		for _, id := range clusterIds {
			cluster, err := l.svcCtx.Scheduler.AiStorages.GetClustersById(id)
			if err != nil {
				return nil, err
			}
			clusters = append(clusters, *cluster)
		}
	} else {
		cs, err := l.svcCtx.Scheduler.AiStorages.GetClustersByAdapterId(ADAPTERID)
		if err != nil {
			return nil, err
		}
		clusters = cs.List
	}

	if len(clusters) == 0 {
		return nil, errors.New("no clusters found ")
	}

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

	var wg sync.WaitGroup
	for _, cluster := range clusters {
		wg.Add(1)
		c := cluster
		go func() {
			defer wg.Done()
			done := make(chan bool)
			var u *collector.ResourceSpec
			var err error
			go func() {

				col, found := l.svcCtx.Scheduler.AiService.AiCollectorAdapterMap[strconv.FormatInt(c.AdapterId, 10)][c.Id]
				if !found {
					done <- true
					return
				}

				u, err = col.GetResourceSpecs(l.ctx)
				if err != nil {
					done <- true
					return
				}

				done <- true
			}()

			select {
			case <-done:
				ch <- u
			case <-time.After(10 * time.Second):
				return
			}

		}()
	}
	wg.Wait()
	close(ch)

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

	// handle empty usage
	return handleEmptyResourceUsage(clusters, ulist), nil
}

func handleEmptyResourceUsage(list []types.ClusterInfo, ulist []*collector.ResourceSpec) []*collector.ResourceSpec {
	var rus []*collector.ResourceSpec
	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.ResourceSpec{
				ClusterId: l.Id,
				Resources: nil,
				Msg:       "resources unavailable, please retry later",
			}
			rus = append(rus, ru)
		} else {
			if s == nil {
				ru := &collector.ResourceSpec{
					ClusterId: l.Id,
					Resources: nil,
					Msg:       "resources unavailable, please retry later",
				}
				rus = append(rus, ru)
			} else {
				r, ok := s.(*collector.ResourceSpec)
				if ok {
					if r.Resources == nil {
						ru := &collector.ResourceSpec{
							ClusterId: r.ClusterId,
							Resources: nil,
							Msg:       "resources unavailable, please retry later",
						}
						rus = append(rus, ru)
					} else {
						// add cluster type
						t, ok := storeLink.ClusterTypeMap[strings.Title(l.Name)]
						if ok {
							r.ClusterType = t
						}
						rus = append(rus, r)
					}
				}
			}
		}
	}

	return rus
}