refactor: refacts the seed peer selection logic from grpc resolver to self pick

Signed-off-by: chlins <chlins.zhang@gmail.com>

Author: chlins

scheduler/job/job.go

@@ -40,12 +40,13 @@ import (
 	commonv1 "d7y.io/api/v2/pkg/apis/common/v1"
 	commonv2 "d7y.io/api/v2/pkg/apis/common/v2"
 	dfdaemonv2 "d7y.io/api/v2/pkg/apis/dfdaemon/v2"
-	managerv2 "d7y.io/api/v2/pkg/apis/manager/v2"
 
 	logger "d7y.io/dragonfly/v2/internal/dflog"
 	internaljob "d7y.io/dragonfly/v2/internal/job"
 	managertypes "d7y.io/dragonfly/v2/manager/types"
+	"d7y.io/dragonfly/v2/pkg/dfnet"
 	"d7y.io/dragonfly/v2/pkg/idgen"
+	cndsystemclient "d7y.io/dragonfly/v2/pkg/rpc/cdnsystem/client"
 	dfdaemonclient "d7y.io/dragonfly/v2/pkg/rpc/dfdaemon/client"
 	"d7y.io/dragonfly/v2/scheduler/config"
 	resource "d7y.io/dragonfly/v2/scheduler/resource/standard"
@@ -63,10 +64,11 @@ type job struct {
 	localJob     *internaljob.Job
 	resource     resource.Resource
 	config       *config.Config
+	dialOptions  []grpc.DialOption
 }
 
 // New creates a new Job.
-func New(cfg *config.Config, resource resource.Resource) (Job, error) {
+func New(cfg *config.Config, resource resource.Resource, dialOptions ...grpc.DialOption) (Job, error) {
 	redisConfig := &internaljob.Config{
 		Addrs:            cfg.Database.Redis.Addrs,
 		MasterName:       cfg.Database.Redis.MasterName,
@@ -111,6 +113,7 @@ func New(cfg *config.Config, resource resource.Resource) (Job, error) {
 		localJob:     localJob,
 		resource:     resource,
 		config:       cfg,
+		dialOptions:  dialOptions,
 	}
 
 	namedJobFuncs := map[string]any{
@@ -230,8 +233,8 @@ func (j *job) preheatSingleSeedPeer(ctx context.Context, req *internaljob.Prehea
 	}
 
 	// If scheduler has no available seed peer, return error.
-	if len(j.resource.SeedPeer().Client().Addrs()) == 0 {
-		return nil, fmt.Errorf("cluster %d scheduler %s has no available seed peer", j.config.Manager.SchedulerClusterID, j.config.Server.AdvertiseIP)
+	if _, err := j.resource.SeedPeer().SelectSeedPeer(ctx, req.TaskUUID); err != nil {
+		return nil, fmt.Errorf("cluster %d scheduler %s has no available seed peer: %w", j.config.Manager.SchedulerClusterID, j.config.Server.AdvertiseIP, err)
 	}
 
 	// Preheat by v2 grpc protocol. If seed peer does not support
@@ -262,8 +265,20 @@ func (j *job) preheatV1SingleSeedPeer(ctx context.Context, req *internaljob.Preh
 		Priority:    commonv1.Priority(req.Priority),
 	}
 
+	selectedSeedPeer, err := j.resource.SeedPeer().SelectSeedPeer(ctx, taskID)
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO: reuse the client if we encounter the performance issue in future.
+	client, err := cndsystemclient.GetClientByAddr(ctx, dfnet.NetAddr{Type: dfnet.TCP, Addr: fmt.Sprintf("%s:%d", selectedSeedPeer.IP, selectedSeedPeer.Port)}, j.dialOptions...)
+	if err != nil {
+		return nil, err
+	}
+	defer client.Close()
+
 	// Trigger seed peer download seeds.
-	stream, err := j.resource.SeedPeer().Client().ObtainSeeds(ctx, &cdnsystemv1.SeedRequest{
+	stream, err := client.ObtainSeeds(ctx, &cdnsystemv1.SeedRequest{
 		TaskId:  taskID,
 		Url:     req.URL,
 		UrlMeta: urlMeta,
@@ -330,7 +345,20 @@ func (j *job) preheatV2SingleSeedPeerByURL(ctx context.Context, url string, req
 	taskID := idgen.TaskIDV2ByURLBased(url, req.PieceLength, req.Tag, req.Application, strings.Split(req.FilteredQueryParams, idgen.FilteredQueryParamsSeparator))
 	filteredQueryParams := strings.Split(req.FilteredQueryParams, idgen.FilteredQueryParamsSeparator)
 	advertiseIP := j.config.Server.AdvertiseIP.String()
-	stream, err := j.resource.SeedPeer().Client().DownloadTask(ctx, taskID, &dfdaemonv2.DownloadTaskRequest{
+
+	selectedSeedPeer, err := j.resource.SeedPeer().SelectSeedPeer(ctx, taskID)
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO: reuse the client if we encounter the performance issue in future.
+	client, err := dfdaemonclient.GetV2ByAddr(ctx, fmt.Sprintf("%s:%d", selectedSeedPeer.IP, selectedSeedPeer.Port), j.dialOptions...)
+	if err != nil {
+		return nil, err
+	}
+	defer client.Close()
+
+	stream, err := client.DownloadTask(ctx, taskID, &dfdaemonv2.DownloadTaskRequest{
 		Download: &commonv2.Download{
 			Url:                 url,
 			PieceLength:         req.PieceLength,
@@ -398,7 +426,7 @@ func (j *job) preheatAllSeedPeers(ctx context.Context, req *internaljob.PreheatR
 			for _, seedPeer := range seedPeers {
 				var (
 					hostname = seedPeer.Hostname
-					ip       = seedPeer.Ip
+					ip       = seedPeer.IP
 					port     = seedPeer.Port
 				)
 
@@ -542,20 +570,20 @@ func (j *job) preheatAllSeedPeers(ctx context.Context, req *internaljob.PreheatR
 // 2. Count: If count is provided, selects up to the specified number of seed peers. If count exceeds the number of available seed peers, all seed peers are selected.
 // 3. Percentage: If percentage is provided, selects a proportional number of seed peers (rounded down). Ensures at least one seed peer is selected if percentage > 0.
 // Priority: IPs > Count > Percentage
-func (j *job) selectSeedPeers(ips []string, count *uint32, percentage *uint8, log *logger.SugaredLoggerOnWith) ([]*managerv2.SeedPeer, error) {
+func (j *job) selectSeedPeers(ips []string, count *uint32, percentage *uint8, log *logger.SugaredLoggerOnWith) ([]*resource.Host, error) {
 	if !j.config.SeedPeer.Enable {
 		return nil, fmt.Errorf("cluster %d scheduler %s has disabled seed peer", j.config.Manager.SchedulerClusterID, j.config.Server.AdvertiseIP)
 	}
 
-	seedPeers := j.resource.SeedPeer().Client().SeedPeers()
+	seedPeers := j.resource.HostManager().LoadAllNonNormals()
 	if len(seedPeers) == 0 {
 		return nil, fmt.Errorf("cluster %d scheduler %s has no available seed peer", j.config.Manager.SchedulerClusterID, j.config.Server.AdvertiseIP)
 	}
 
 	if len(ips) > 0 {
-		selectedSeedPeers := make([]*managerv2.SeedPeer, 0, len(ips))
+		selectedSeedPeers := make([]*resource.Host, 0, len(ips))
 		for _, seedPeer := range seedPeers {
-			if slices.Contains(ips, seedPeer.Ip) {
+			if slices.Contains(ips, seedPeer.IP) {
 				selectedSeedPeers = append(selectedSeedPeers, seedPeer)
 				continue
 			}

scheduler/resource/standard/host_manager.go

@@ -63,6 +63,9 @@ type HostManager interface {
 	// LoadAllNormals loads all normal hosts through the Range of sync.Map.
 	LoadAllNormals() []*Host
 
+	// LoadAllNonNormals loads all non-normal hosts through the Range of sync.Map.
+	LoadAllNonNormals() []*Host
+
 	// Try to reclaim host.
 	RunGC(context.Context) error
 }
@@ -163,6 +166,27 @@ func (h *hostManager) LoadAllNormals() []*Host {
 	return hosts
 }
 
+// LoadAllNonNormals loads all non-normal hosts through the Range of sync.Map.
+func (h *hostManager) LoadAllNonNormals() []*Host {
+	hosts := make([]*Host, 0)
+	h.Map.Range(func(key, value any) bool {
+		host, ok := value.(*Host)
+		if !ok {
+			host.Log.Error("invalid host")
+			return true
+		}
+
+		if host.Type == types.HostTypeNormal {
+			return true
+		}
+
+		hosts = append(hosts, host)
+		return true
+	})
+
+	return hosts
+}
+
 // LoadRandom loads host randomly through the Range of sync.Map.
 func (h *hostManager) LoadRandom(n int, blocklist set.SafeSet[string]) []*Host {
 	hosts := make([]*Host, 0, n)

scheduler/resource/standard/host_manager_mock.go

@@ -91,12 +91,7 @@ func New(cfg *config.Config, gc gc.GC, dynconfig config.DynconfigInterface, tran
 	// Initialize seed peer interface.
 	if cfg.SeedPeer.Enable {
 		dialOptions := []grpc.DialOption{grpc.WithStatsHandler(otelgrpc.NewClientHandler()), grpc.WithTransportCredentials(transportCredentials)}
-		client, err := newSeedPeerClient(dynconfig, hostManager, dialOptions...)
-		if err != nil {
-			return nil, err
-		}
-
-		resource.seedPeer = newSeedPeer(client, peerManager, hostManager)
+		resource.seedPeer = newSeedPeer(peerManager, hostManager, dialOptions...)
 	}
 
 	return resource, nil

scheduler/resource/standard/resource.go

@@ -26,17 +26,22 @@ import (
 	"time"
 
 	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc"
+	"stathat.com/c/consistent"
 
 	cdnsystemv1 "d7y.io/api/v2/pkg/apis/cdnsystem/v1"
 	commonv1 "d7y.io/api/v2/pkg/apis/common/v1"
 	commonv2 "d7y.io/api/v2/pkg/apis/common/v2"
 	dfdaemonv2 "d7y.io/api/v2/pkg/apis/dfdaemon/v2"
 	schedulerv1 "d7y.io/api/v2/pkg/apis/scheduler/v1"
 
+	"d7y.io/dragonfly/v2/pkg/dfnet"
 	"d7y.io/dragonfly/v2/pkg/digest"
 	"d7y.io/dragonfly/v2/pkg/idgen"
 	"d7y.io/dragonfly/v2/pkg/net/http"
+	cndsystemclient "d7y.io/dragonfly/v2/pkg/rpc/cdnsystem/client"
 	"d7y.io/dragonfly/v2/pkg/rpc/common"
+	dfdaemonclient "d7y.io/dragonfly/v2/pkg/rpc/dfdaemon/client"
 	"d7y.io/dragonfly/v2/scheduler/metrics"
 )
 
@@ -55,31 +60,31 @@ type SeedPeer interface {
 	// Used only in v1 version of the grpc.
 	TriggerTask(context.Context, *http.Range, *Task) (*Peer, *schedulerv1.PeerResult, error)
 
-	// Client returns grpc client of seed peer.
-	Client() SeedPeerClient
+	// SelectSeedPeer selects a seed peer by the task id.
+	SelectSeedPeer(context.Context, string) (*Host, error)
 
 	// Stop seed peer service.
 	Stop() error
 }
 
 // seedPeer contains content for seed peer.
 type seedPeer struct {
-	// client is the dynamic client of seed peer.
-	client SeedPeerClient
-
 	// peerManager is PeerManager interface.
 	peerManager PeerManager
 
 	// hostManager is HostManager interface.
 	hostManager HostManager
+
+	// dialOpts is the options for grpc dial.
+	dialOptions []grpc.DialOption
 }
 
 // New SeedPeer interface.
-func newSeedPeer(client SeedPeerClient, peerManager PeerManager, hostManager HostManager) SeedPeer {
+func newSeedPeer(peerManager PeerManager, hostManager HostManager, dialOptions ...grpc.DialOption) SeedPeer {
 	return &seedPeer{
-		client:      client,
 		peerManager: peerManager,
 		hostManager: hostManager,
+		dialOptions: dialOptions,
 	}
 }
 
@@ -89,7 +94,19 @@ func (s *seedPeer) TriggerDownloadTask(ctx context.Context, taskID string, req *
 	ctx, cancel := context.WithCancel(trace.ContextWithSpan(ctx, trace.SpanFromContext(ctx)))
 	defer cancel()
 
-	stream, err := s.client.DownloadTask(ctx, taskID, req)
+	selectedSeedPeer, err := s.SelectSeedPeer(ctx, taskID)
+	if err != nil {
+		return err
+	}
+
+	// TODO: reuse the client if we encounter the performance issue in future.
+	client, err := dfdaemonclient.GetV2ByAddr(ctx, fmt.Sprintf("%s:%d", selectedSeedPeer.IP, selectedSeedPeer.Port), s.dialOptions...)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	stream, err := client.DownloadTask(ctx, taskID, req)
 	if err != nil {
 		return err
 	}
@@ -126,7 +143,19 @@ func (s *seedPeer) TriggerTask(ctx context.Context, rg *http.Range, task *Task)
 		urlMeta.Range = rg.URLMetaString()
 	}
 
-	stream, err := s.client.ObtainSeeds(ctx, &cdnsystemv1.SeedRequest{
+	selectedSeedPeer, err := s.SelectSeedPeer(ctx, task.ID)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// TODO: reuse the client if we encounter the performance issue in future.
+	client, err := cndsystemclient.GetClientByAddr(ctx, dfnet.NetAddr{Type: dfnet.TCP, Addr: fmt.Sprintf("%s:%d", selectedSeedPeer.IP, selectedSeedPeer.Port)}, s.dialOptions...)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer client.Close()
+
+	stream, err := client.ObtainSeeds(ctx, &cdnsystemv1.SeedRequest{
 		TaskId:  task.ID,
 		Url:     task.URL,
 		UrlMeta: urlMeta,
@@ -222,6 +251,32 @@ func (s *seedPeer) TriggerTask(ctx context.Context, rg *http.Range, task *Task)
 	}
 }
 
+// SelectSeedPeer selects a seed peer by the task id.
+func (s *seedPeer) SelectSeedPeer(ctx context.Context, taskID string) (*Host, error) {
+	// Currently non normal host is seed peer.
+	seedPeers := s.hostManager.LoadAllNonNormals()
+	if len(seedPeers) == 0 {
+		return nil, fmt.Errorf("no seed peer found in host manager")
+	}
+
+	// Build the hash ring by adding all the seed peers.
+	hashring := consistent.New()
+	hosts := make(map[string]*Host, len(seedPeers))
+	for _, host := range seedPeers {
+		hashKey := fmt.Sprintf("%s:%d:%s", host.IP, host.Port, host.Hostname)
+		hashring.Add(hashKey)
+		hosts[hashKey] = host
+	}
+
+	// Pick one seed peer from the hash ring by the task id.
+	hashKey, err := hashring.Get(taskID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to pick the seed peer: %w", err)
+	}
+
+	return hosts[hashKey], nil
+}
+
 // Initialize seed peer.
 func (s *seedPeer) initSeedPeer(ctx context.Context, rg *http.Range, task *Task, hostID string, peerID string) (*Peer, error) {
 	// Load host from manager.
@@ -256,12 +311,7 @@ func (s *seedPeer) initSeedPeer(ctx context.Context, rg *http.Range, task *Task,
 	return peer, nil
 }
 
-// Client is seed peer grpc client.
-func (s *seedPeer) Client() SeedPeerClient {
-	return s.client
-}
-
 // Stop seed peer service.
 func (s *seedPeer) Stop() error {
-	return s.client.Close()
+	return nil
 }