PR Code Review Changes

cmd/kill.go

@@ -15,6 +15,8 @@
 package cmd
 
 import (
+	"strconv"
+
 	realis "github.com/aurora-scheduler/gorealis/v2"
 	"github.com/spf13/cobra"
 )
@@ -26,6 +28,7 @@ func init() {
 
 	// Kill Job
 	killCmd.AddCommand(killJobCmd)
+	killCmd.AddCommand(killTaskCmd)
 
 	killJobCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment")
 	killJobCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role")
@@ -34,6 +37,17 @@ func init() {
 	killJobCmd.MarkFlagRequired("environment")
 	killJobCmd.MarkFlagRequired("role")
 	killJobCmd.MarkFlagRequired("name")
+
+	//Set flags for killTask sub-command
+	killTaskCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment")
+	killTaskCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role")
+	killTaskCmd.Flags().StringVarP(name, "name", "n", "", "Aurora Name")
+	killTaskCmd.Flags().StringVarP(instance, "instance", "i", "", "Instance Number")
+	killTaskCmd.Flags().BoolVarP(&monitor, "monitor", "m", true, "monitor the result after sending the command")
+	killTaskCmd.MarkFlagRequired("environment")
+	killTaskCmd.MarkFlagRequired("role")
+	killTaskCmd.MarkFlagRequired("name")
+	killTaskCmd.MarkFlagRequired("instance")
 }
 
 var killCmd = &cobra.Command{
@@ -47,6 +61,12 @@ var killJobCmd = &cobra.Command{
 	Run:   killJob,
 }
 
+var killTaskCmd = &cobra.Command{
+	Use:   "task",
+	Short: "Kill an Aurora Task",
+	Run:   killTask,
+}
+
 func killJob(cmd *cobra.Command, args []string) {
 	log.Infof("Killing job [Env:%s Role:%s Name:%s]\n", *env, *role, *name)
 
@@ -64,3 +84,26 @@ func killJob(cmd *cobra.Command, args []string) {
 		}
 	}
 }
+
+func killTask(cmd *cobra.Command, args []string) {
+	log.Infof("Killing task [Env:%s Role:%s Name:%s Instance:%s]\n", *env, *role, *name, *instance)
+
+	//Set jobKey for the task to be killed.
+	task := realis.NewTask().
+		Environment(*env).
+		Role(*role).
+		Name(*name)
+
+	//Convert instance from string type to int32 type, and call the killtasks function
+	instanceNumber, _ := strconv.Atoi(*instance)
+	_, err := client.KillInstances(task.JobKey(), (int32)(instanceNumber))
+
+	if err != nil {
+		log.Fatalln(err)
+	}
+	if monitor {
+		if ok, err := client.MonitorInstances(task.JobKey(), 0, 5, 50); !ok || err != nil {
+			log.Fatalln("Unable to kill the given task")
+		}
+	}
+}

cmd/killTask.go

@@ -1,74 +0,0 @@
-/**
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package cmd
-
-import (
-	"strconv"
-
-	realis "github.com/aurora-scheduler/gorealis/v2"
-	"github.com/spf13/cobra"
-)
-
-func init() {
-	rootCmd.AddCommand(killTask)
-
-	/* Sub-Commands */
-
-	// Kill Job
-	killTask.AddCommand(killTaskCmd)
-
-	killTaskCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment")
-	killTaskCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role")
-	killTaskCmd.Flags().StringVarP(name, "name", "n", "", "Aurora Name")
-	killTaskCmd.Flags().StringVarP(instance, "instance", "i", "", "Instance number")
-	killTaskCmd.Flags().BoolVarP(&monitor, "monitor", "m", true, "monitor the result after sending the command")
-	killTaskCmd.MarkFlagRequired("environment")
-	killTaskCmd.MarkFlagRequired("role")
-	killTaskCmd.MarkFlagRequired("name")
-	killTaskCmd.MarkFlagRequired("instance")
-}
-
-var killTask = &cobra.Command{
-	Use:   "killTask",
-	Short: "Kill an Aurora Task",
-}
-
-var killTaskCmd = &cobra.Command{
-	Use:   "task",
-	Short: "Kill an Aurora Task",
-	Run:   killTaskFunc,
-}
-
-func killTaskFunc(cmd *cobra.Command, args []string) {
-	log.Infof("Killing task [Env:%s Role:%s Name:%s Instance:%s]\n", *env, *role, *name, *instance)
-
-	//Set jobKey for the task to be killed.
-	task := realis.NewTask().
-		Environment(*env).
-		Role(*role).
-		Name(*name)
-
-	//Convert instance from string type to int32 type, and call the killtasks function
-	instanceNumber, _ := strconv.Atoi(*instance)
-	_, err := client.KillInstances(task.JobKey(), (int32)(instanceNumber))
-
-	if err != nil {
-		log.Fatalln(err)
-	}
-	if monitor {
-		if ok, err := client.MonitorInstances(task.JobKey(), 0, 5, 50); !ok || err != nil {
-			log.Fatalln("Unable to kill the given task")
-		}
-	}
-}

cmd/root.go

@@ -28,7 +28,7 @@ import (
 )
 
 var username, password, zkAddr, schedAddr string
-var env, role, name, instance = new(string), new(string), new(string), new(string)
+var env, role, name = new(string), new(string), new(string)
 var dedicated string
 var ram, disk, gpu, port int64
 var cpu float64
@@ -51,6 +51,7 @@ var monitor bool
 var timeout time.Duration
 var log = logrus.New()
 var taskStatus = new(string)
+var instance = new(string)
 
 const australisVer = "v1.0.4"