Add the ability to pass labels

FairwindsOps · Jun 23, 2020 · d76b69a · d76b69a
1 parent 88b73cb
commit d76b69a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -40,6 +40,7 @@ var (
 	demoCount     int
 	demoDelay     int
 	demoDuration  int
+	labelSelector map[string]string
 )
 
 func init() {
@@ -49,6 +50,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&host, "host", "", "", "The hostname or IP address of the allocator server")
 	rootCmd.PersistentFlags().StringVarP(&namespace, "namespace", "n", "", "The namespace of gameservers to request from")
 	rootCmd.PersistentFlags().BoolVarP(&multicluster, "multicluster", "m", false, "If true, multicluster allocation will be requested")
+	rootCmd.PersistentFlags().StringToStringVar(&labelSelector, "labels", nil, "A map of labels to match on the allocation.")
 
 	rootCmd.AddCommand(allocateCmd)
 	rootCmd.AddCommand(loadTestCmd)
@@ -106,7 +108,7 @@ var allocateCmd = &cobra.Command{
 	Long:    `Request an allocated server`,
 	PreRunE: argsValidator,
 	Run: func(cmd *cobra.Command, args []string) {
-		allocatorClient, err := allocator.NewClient(keyFile, certFile, caCertFile, host, namespace, multicluster)
+		allocatorClient, err := allocator.NewClient(keyFile, certFile, caCertFile, host, namespace, multicluster, labelSelector)
 		if err != nil {
 			klog.Error(err)
 		}
@@ -125,7 +127,7 @@ var loadTestCmd = &cobra.Command{
 	Long:    `Allocates a set of servers, communicates with them, and then closes the connection.`,
 	PreRunE: argsValidator,
 	Run: func(cmd *cobra.Command, args []string) {
-		allocatorClient, err := allocator.NewClient(keyFile, certFile, caCertFile, host, namespace, multicluster)
+		allocatorClient, err := allocator.NewClient(keyFile, certFile, caCertFile, host, namespace, multicluster, labelSelector)
 		if err != nil {
 			klog.Error(err)
 		}

diff --git a/pkg/allocator/allocator.go b/pkg/allocator/allocator.go
@@ -38,6 +38,7 @@ type Client struct {
 	Namespace    string
 	Multicluster bool
 	DialOpts     grpc.DialOption
+	MatchLabels  map[string]string
 }
 
 // Allocation is a game server allocation
@@ -47,7 +48,7 @@ type Allocation struct {
 }
 
 // NewClient builds a new client object
-func NewClient(keyFile string, certFile string, cacertFile string, externalIP string, namespace string, multiCluster bool) (*Client, error) {
+func NewClient(keyFile, certFile, cacertFile, externalIP, namespace string, multiCluster bool, labelSelector map[string]string) (*Client, error) {
 	endpoint := externalIP + ":443"
 	cert, err := ioutil.ReadFile(certFile)
 	if err != nil {
@@ -69,6 +70,7 @@ func NewClient(keyFile string, certFile string, cacertFile string, externalIP st
 		Endpoint:     endpoint,
 		Multicluster: multiCluster,
 		Namespace:    namespace,
+		MatchLabels:  labelSelector,
 	}
 	err = newClient.createRemoteClusterDialOption()
 	if err != nil {
@@ -107,6 +109,9 @@ func (c *Client) AllocateGameserver() (*Allocation, error) {
 		MultiClusterSetting: &pb.MultiClusterSetting{
 			Enabled: c.Multicluster,
 		},
+		RequiredGameServerSelector: &pb.LabelSelector{
+			MatchLabels: c.MatchLabels,
+		},
 	}
 
 	resp, err := c.makeRequest(request)