feat: add context for FindSchedulerCluster (#841)

Signed-off-by: Gaius <gaius.qi@gmail.com>
dragonflyoss · Nov 30, 2021 · de30f0a · de30f0a
1 parent 956b4fb
commit de30f0a
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 6 deletions.
diff --git a/manager/searcher/searcher.go b/manager/searcher/searcher.go
@@ -17,6 +17,7 @@
 package searcher
 
 import (
+	"context"
 	"strings"
 
 	"github.com/mitchellh/mapstructure"
@@ -72,7 +73,7 @@ type Scopes struct {
 }
 
 type Searcher interface {
-	FindSchedulerCluster([]model.SchedulerCluster, *manager.ListSchedulersRequest) (model.SchedulerCluster, bool)
+	FindSchedulerCluster(context.Context, []model.SchedulerCluster, *manager.ListSchedulersRequest) (model.SchedulerCluster, bool)
 }
 
 type searcher struct{}
@@ -86,7 +87,7 @@ func New() Searcher {
 	return s
 }
 
-func (s *searcher) FindSchedulerCluster(schedulerClusters []model.SchedulerCluster, client *manager.ListSchedulersRequest) (model.SchedulerCluster, bool) {
+func (s *searcher) FindSchedulerCluster(ctx context.Context, schedulerClusters []model.SchedulerCluster, client *manager.ListSchedulersRequest) (model.SchedulerCluster, bool) {
 	conditions := client.HostInfo
 	if len(schedulerClusters) <= 0 || len(conditions) <= 0 {
 		return model.SchedulerCluster{}, false

diff --git a/manager/searcher/searcher_test.go b/manager/searcher/searcher_test.go
@@ -17,6 +17,7 @@
 package searcher
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -345,7 +346,7 @@ func TestSchedulerCluster(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			searcher := New()
-			clusters, ok := searcher.FindSchedulerCluster(tc.schedulerClusters, &manager.ListSchedulersRequest{
+			clusters, ok := searcher.FindSchedulerCluster(context.Background(), tc.schedulerClusters, &manager.ListSchedulersRequest{
 				HostName: "foo",
 				Ip:       "127.0.0.1",
 				HostInfo: tc.conditions,

diff --git a/manager/searcher/testdata/main.go b/manager/searcher/testdata/main.go
@@ -17,6 +17,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"os"
@@ -40,7 +41,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	cluster, ok := s.FindSchedulerCluster([]model.SchedulerCluster{}, &manager.ListSchedulersRequest{})
+	cluster, ok := s.FindSchedulerCluster(context.Background(), []model.SchedulerCluster{}, &manager.ListSchedulersRequest{})
 	if !ok {
 		fmt.Println("scheduler cluster not found")
 		os.Exit(1)

diff --git a/manager/searcher/testdata/plugin/searcher.go b/manager/searcher/testdata/plugin/searcher.go
@@ -17,13 +17,15 @@
 package main
 
 import (
+	"context"
+
 	"d7y.io/dragonfly/v2/manager/model"
 	"d7y.io/dragonfly/v2/pkg/rpc/manager"
 )
 
 type searcher struct{}
 
-func (s *searcher) FindSchedulerCluster(schedulerClusters []model.SchedulerCluster, client *manager.ListSchedulersRequest) (model.SchedulerCluster, bool) {
+func (s *searcher) FindSchedulerCluster(ctx context.Context, schedulerClusters []model.SchedulerCluster, client *manager.ListSchedulersRequest) (model.SchedulerCluster, bool) {
 	return model.SchedulerCluster{Name: "foo"}, true
 }
 

diff --git a/manager/service/service_grpc.go b/manager/service/service_grpc.go
@@ -382,7 +382,7 @@ func (s *GRPC) ListSchedulers(ctx context.Context, req *manager.ListSchedulersRe
 	}
 
 	// Search optimal scheduler cluster
-	schedulerCluster, ok := s.searcher.FindSchedulerCluster(schedulerClusters, req)
+	schedulerCluster, ok := s.searcher.FindSchedulerCluster(ctx, schedulerClusters, req)
 	if !ok {
 		return nil, status.Error(codes.NotFound, "scheduler cluster not found")
 	}