feat(scheduler): add support for filtering hosts by type in ListHosts

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

Author: chlins

go.mod

@@ -4,7 +4,7 @@ go 1.23.8
 
 require (
 	cloud.google.com/go/storage v1.50.0
-	d7y.io/api/v2 v2.1.60
+	d7y.io/api/v2 v2.1.69
 	github.com/MysteriousPotato/go-lockable v1.0.0
 	github.com/Showmax/go-fqdn v1.0.0
 	github.com/VividCortex/mysqlerr v1.0.0
@@ -96,7 +96,7 @@ require (
 	golang.org/x/time v0.12.0
 	google.golang.org/api v0.248.0
 	google.golang.org/grpc v1.75.0
-	google.golang.org/protobuf v1.36.8
+	google.golang.org/protobuf v1.36.9
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/yaml.v3 v3.0.1
 	gorm.io/driver/mysql v1.4.7

go.sum

@@ -63,8 +63,8 @@ cloud.google.com/go/storage v1.50.0 h1:3TbVkzTooBvnZsk7WaAQfOsNrdoM8QHusXA1cpk6Q
 cloud.google.com/go/storage v1.50.0/go.mod h1:l7XeiD//vx5lfqE3RavfmU9yvk5Pp0Zhcv482poyafY=
 cloud.google.com/go/trace v1.11.6 h1:2O2zjPzqPYAHrn3OKl029qlqG6W8ZdYaOWRyr8NgMT4=
 cloud.google.com/go/trace v1.11.6/go.mod h1:GA855OeDEBiBMzcckLPE2kDunIpC72N+Pq8WFieFjnI=
-d7y.io/api/v2 v2.1.60 h1:Imk1mw30wmTJNdqKFF5VO/eqI/t1H5PfIPPK4ryKX2U=
-d7y.io/api/v2 v2.1.60/go.mod h1:WzakPywEgs27gr/TwrlRarPRRK2o5nN0YrTzHDdFii8=
+d7y.io/api/v2 v2.1.69 h1:F0G9Vdi5Nk9IOp7hGSrw0o/mEr01La/bRSNQP6e/KSk=
+d7y.io/api/v2 v2.1.69/go.mod h1:Zd/KrYRZ1saG0c+n+G5/6HqIFKTHHLKrZzOpqi3nmLw=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
@@ -2186,8 +2186,8 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc=
-google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
+google.golang.org/protobuf v1.36.9 h1:w2gp2mA27hUeUzj9Ex9FBjsBm40zfaDtEWow293U7Iw=
+google.golang.org/protobuf v1.36.9/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
 gopkg.in/airbrake/gobrake.v2 v2.0.9/go.mod h1:/h5ZAUhDkGaJfjzjKLSjv6zCL6O0LLBxU4K+aSYdM/U=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

pkg/rpc/scheduler/client/client_v2.go

@@ -31,7 +31,6 @@ import (
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
-	"google.golang.org/protobuf/types/known/emptypb"
 
 	commonv2 "d7y.io/api/v2/pkg/apis/common/v2"
 	schedulerv2 "d7y.io/api/v2/pkg/apis/scheduler/v2"
@@ -155,7 +154,7 @@ type V2 interface {
 	AnnounceHost(context.Context, *schedulerv2.AnnounceHostRequest, ...grpc.CallOption) error
 
 	// ListHosts lists hosts in scheduler.
-	ListHosts(ctx context.Context, taskID string, opts ...grpc.CallOption) (*schedulerv2.ListHostsResponse, error)
+	ListHosts(ctx context.Context, taskID string, req *schedulerv2.ListHostsRequest, opts ...grpc.CallOption) (*schedulerv2.ListHostsResponse, error)
 
 	// DeleteHost releases host in scheduler.
 	DeleteHost(context.Context, *schedulerv2.DeleteHostRequest, ...grpc.CallOption) error
@@ -263,13 +262,13 @@ func (v *v2) AnnounceHost(ctx context.Context, req *schedulerv2.AnnounceHostRequ
 }
 
 // ListHosts lists host in all schedulers.
-func (v *v2) ListHosts(ctx context.Context, taskID string, opts ...grpc.CallOption) (*schedulerv2.ListHostsResponse, error) {
+func (v *v2) ListHosts(ctx context.Context, taskID string, req *schedulerv2.ListHostsRequest, opts ...grpc.CallOption) (*schedulerv2.ListHostsResponse, error) {
 	ctx, cancel := context.WithTimeout(ctx, contextTimeout)
 	defer cancel()
 
 	return v.SchedulerClient.ListHosts(
 		context.WithValue(ctx, pkgbalancer.ContextKey, taskID),
-		new(emptypb.Empty),
+		req,
 		opts...,
 	)
 }

scheduler/rpcserver/scheduler_server_v2.go

@@ -139,10 +139,10 @@ func (s *schedulerServerV2) AnnounceHost(ctx context.Context, req *schedulerv2.A
 }
 
 // ListHosts lists hosts in scheduler.
-func (s *schedulerServerV2) ListHosts(ctx context.Context, _ *emptypb.Empty) (*schedulerv2.ListHostsResponse, error) {
+func (s *schedulerServerV2) ListHosts(ctx context.Context, req *schedulerv2.ListHostsRequest) (*schedulerv2.ListHostsResponse, error) {
 	// Collect ListHostsCount metrics.
 	metrics.ListHostsCount.Inc()
-	resp, err := s.service.ListHosts(ctx)
+	resp, err := s.service.ListHosts(ctx, req)
 	if err != nil {
 		// Collect ListHostsFailureCount metrics.
 		metrics.ListHostsCountFailureCount.Inc()

scheduler/service/service_v2.go

@@ -949,9 +949,9 @@ func (v *V2) AnnounceHost(ctx context.Context, req *schedulerv2.AnnounceHostRequ
 }
 
 // ListHosts lists hosts in scheduler.
-func (v *V2) ListHosts(ctx context.Context) (*schedulerv2.ListHostsResponse, error) {
+func (v *V2) ListHosts(ctx context.Context, req *schedulerv2.ListHostsRequest) (*schedulerv2.ListHostsResponse, error) {
 	hosts := []*commonv2.Host{}
-	v.resource.HostManager().Range(func(_ any, value any) bool {
+	constructHosts := func(value any) bool {
 		host, ok := value.(*standard.Host)
 		if !ok {
 			// Continue to next host.
@@ -1031,7 +1031,30 @@ func (v *V2) ListHosts(ctx context.Context) (*schedulerv2.ListHostsResponse, err
 		})
 
 		return true
-	})
+	}
+
+	// Return all hosts if no type specified.
+	if req.Type == nil {
+		v.resource.HostManager().Range(func(_ any, value any) bool {
+			return constructHosts(value)
+		})
+
+		return &schedulerv2.ListHostsResponse{
+			Hosts: hosts,
+		}, nil
+	}
+
+	// Filter hosts by type.
+	switch types.HostType(*req.Type) {
+	case types.HostTypeNormal:
+		v.resource.HostManager().RangeNormals(func(_ any, value any) bool {
+			return constructHosts(value)
+		})
+	case types.HostTypeSuperSeed:
+		v.resource.HostManager().RangeSeeds(func(_ any, value any) bool {
+			return constructHosts(value)
+		})
+	}
 
 	return &schedulerv2.ListHostsResponse{
 		Hosts: hosts,

scheduler/service/service_v2_test.go

@@ -1558,7 +1558,7 @@ func TestServiceV2_ListHosts(t *testing.T) {
 			svc := NewV2(&config.Config{Scheduler: mockSchedulerConfig, Metrics: config.MetricsConfig{EnableHost: true}}, resource, persistentCacheResource, scheduling, job, internalJobImage, dynconfig)
 
 			tc.mock(host, hostManager, resource.EXPECT(), hostManager.EXPECT())
-			resp, err := svc.ListHosts(context.Background())
+			resp, err := svc.ListHosts(context.Background(), &schedulerv2.ListHostsRequest{})
 			tc.expect(t, host, resp.Hosts, err)
 		})
 	}

test/e2e/host_test.go

@@ -26,6 +26,8 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 
+	schedulerv2 "d7y.io/api/v2/pkg/apis/scheduler/v2"
+
 	schedulerclient "d7y.io/dragonfly/v2/pkg/rpc/scheduler/client"
 	"d7y.io/dragonfly/v2/pkg/types"
 	"d7y.io/dragonfly/v2/test/e2e/util"
@@ -104,7 +106,7 @@ var _ = Describe("Clients Leaving", func() {
 })
 
 func calculateNormalHostCountFromScheduler(schedulerClient schedulerclient.V2) (hostCount int) {
-	resp, err := schedulerClient.ListHosts(context.Background(), "")
+	resp, err := schedulerClient.ListHosts(context.Background(), "", &schedulerv2.ListHostsRequest{})
 	fmt.Println(resp, err)
 	Expect(err).NotTo(HaveOccurred())