JCS-pub/hub/internal/ticktock/test_hub_connectivities.go

package ticktock

import (
	"context"
	"sync"
	"time"

	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/common/utils/reflect2"
	stgglb "gitlink.org.cn/cloudream/jcs-pub/common/globals"
	corrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/coordinator"
	hubrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/hub"
	cortypes "gitlink.org.cn/cloudream/jcs-pub/coordinator/types"
)

type TestHubConnectivities struct {
	myHubID cortypes.HubID
}

func (t *TestHubConnectivities) Name() string {
	return reflect2.TypeNameOf[TestHubConnectivities]()
}
func (j *TestHubConnectivities) Execute(t *TickTock) {
	log := logger.WithType[TestHubConnectivities]("TickTock")
	startTime := time.Now()
	log.Infof("job start")
	defer func() {
		log.Infof("job end, time: %v", time.Since(startTime))
	}()

	coorCli := stgglb.CoordinatorRPCPool.Get()
	defer coorCli.Release()

	getHubs, cerr := coorCli.GetHubs(context.Background(), corrpc.NewGetHubs(nil))
	if cerr != nil {
		log.Warnf("get all hubs: %v", cerr)
		return
	}

	tests := make([]cortypes.HubConnectivity, len(getHubs.Hubs))

	wg := sync.WaitGroup{}
	for i, hub := range getHubs.Hubs {
		wg.Add(1)

		go func(hub *cortypes.Hub, i int) {
			defer wg.Done()

			tests[i] = j.testOne(hub)
		}(hub, i)
	}
	wg.Wait()

	_, cerr = coorCli.ReportHubConnectivity(context.Background(), &corrpc.ReportHubConnectivity{
		Connecttivities: tests,
	})
	if cerr != nil {
		log.Warnf("report hub connectivities: %v", cerr)
	}
}

func (j *TestHubConnectivities) testOne(hub *cortypes.Hub) cortypes.HubConnectivity {
	log := logger.WithType[TestHubConnectivities]("TickTock")

	rpcAddr, ok := hub.Address.(*cortypes.GRPCAddressInfo)
	if !ok {
		return cortypes.HubConnectivity{
			FromHubID: j.myHubID,
			ToHubID:   hub.HubID,
			Latency:   nil,
			TestTime:  time.Now(),
		}
	}
	hubCli := stgglb.HubRPCPool.Get(stgglb.SelectGRPCAddress(hub, rpcAddr))
	defer hubCli.Release()

	// 后几次ping计算延迟
	var avgLatency time.Duration
	for i := 0; i < 3; i++ {
		start := time.Now()
		_, cerr := hubCli.Ping(context.Background(), &hubrpc.Ping{})
		if cerr != nil {
			log.Warnf("ping %v: %v", hub.String(), cerr)
			return cortypes.HubConnectivity{
				FromHubID: j.myHubID,
				ToHubID:   hub.HubID,
				Latency:   nil,
				TestTime:  time.Now(),
			}
		}

		latency := time.Since(start)
		avgLatency += latency

		// 每次ping之间间隔1秒
		<-time.After(time.Second)
	}
	latency := avgLatency / 3
	latencyMs := float32(latency.Microseconds()) / 1000

	return cortypes.HubConnectivity{
		FromHubID: j.myHubID,
		ToHubID:   hub.HubID,
		Latency:   &latencyMs,
		TestTime:  time.Now(),
	}
}