[swarming] Implement task cancel command to swarming

Bug: 1135017
Change-Id: I80d046009f0acec89b67ac48b11e14df7dc328ed
Reviewed-on: https://chromium-review.googlesource.com/c/infra/luci/luci-go/+/3246953
Commit-Queue: Junji Watanabe <jwata@google.com>
Commit-Queue: Takuto Ikuta <tikuta@chromium.org>
Auto-Submit: Junji Watanabe <jwata@google.com>
Reviewed-by: Takuto Ikuta <tikuta@chromium.org>
diff --git a/client/cmd/swarming/lib/cancel.go b/client/cmd/swarming/lib/cancel.go
new file mode 100644
index 0000000..84dd7ad
--- /dev/null
+++ b/client/cmd/swarming/lib/cancel.go
@@ -0,0 +1,106 @@
+// Copyright 2021 The LUCI Authors.
+//
+// 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 lib
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/maruel/subcommands"
+
+	"go.chromium.org/luci/common/api/swarming/swarming/v1"
+	"go.chromium.org/luci/common/errors"
+	"go.chromium.org/luci/common/system/signals"
+)
+
+// CmdCancelTask returns an object for the `cancel` subcommand.
+func CmdCancelTask(authFlags AuthFlags) *subcommands.Command {
+	return &subcommands.Command{
+		UsageLine: "cancel <options> <taskID>",
+		ShortDesc: "cancel a task",
+		LongDesc:  "Cancels the task specified by the taskID",
+		CommandRun: func() subcommands.CommandRun {
+			r := &cancelRun{}
+			r.Init(authFlags)
+			return r
+		},
+	}
+}
+
+type cancelRun struct {
+	commonFlags
+	killRunning bool
+}
+
+func (t *cancelRun) Init(authFlags AuthFlags) {
+	t.commonFlags.Init(authFlags)
+	t.Flags.BoolVar(&t.killRunning, "kill-running", false, "Kill the task even if it's running")
+}
+
+func (t *cancelRun) parse(taskIDs []string) error {
+	if err := t.commonFlags.Parse(); err != nil {
+		return err
+	}
+
+	if len(taskIDs) == 0 {
+		return errors.New("must specify a swarming task ID")
+	}
+
+	if len(taskIDs) > 1 {
+		return errors.New("please specify only one swarming task ID")
+	}
+
+	return nil
+}
+
+func (t *cancelRun) cancelTask(ctx context.Context, taskID string, service swarmingService) error {
+	req := &swarming.SwarmingRpcsTaskCancelRequest{
+		KillRunning: t.killRunning,
+	}
+	res, err := service.CancelTask(ctx, taskID, req)
+	if res != nil && !res.Ok {
+		err = errors.Reason("response was not OK. running=%v\n", res.WasRunning).Err()
+	}
+	if err != nil {
+		return errors.Annotate(err, "failed to cancel task %s\n", taskID).Err()
+	}
+
+	fmt.Printf("Cancelled %s\n", taskID)
+	return nil
+
+}
+
+func (t *cancelRun) main(_ subcommands.Application, taskID string) error {
+	ctx, cancel := context.WithCancel(t.defaultFlags.MakeLoggingContext(os.Stderr))
+	defer signals.HandleInterrupt(cancel)()
+	service, err := t.createSwarmingClient(ctx)
+	if err != nil {
+		return err
+	}
+	return t.cancelTask(ctx, taskID, service)
+}
+
+func (t *cancelRun) Run(a subcommands.Application, taskIDs []string, _ subcommands.Env) int {
+	if err := t.parse(taskIDs); err != nil {
+		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
+		return 1
+	}
+	if err := t.main(a, taskIDs[0]); err != nil {
+		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
+		return 1
+	}
+	return 0
+}
diff --git a/client/cmd/swarming/lib/cancel_test.go b/client/cmd/swarming/lib/cancel_test.go
new file mode 100644
index 0000000..8a9dd9f
--- /dev/null
+++ b/client/cmd/swarming/lib/cancel_test.go
@@ -0,0 +1,102 @@
+// Copyright 2021 The LUCI Authors.
+//
+// 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 lib
+
+import (
+	"context"
+	"testing"
+
+	. "github.com/smartystreets/goconvey/convey"
+
+	"go.chromium.org/luci/common/api/swarming/swarming/v1"
+	. "go.chromium.org/luci/common/testing/assertions"
+)
+
+func TestCancelTaskParse(t *testing.T) {
+	t.Parallel()
+
+	Convey(`Test CancelTaskParse when there's no input or too many inputs`, t, func() {
+
+		t := cancelRun{}
+		t.Init(&testAuthFlags{})
+
+		err := t.GetFlags().Parse([]string{"-server", "http://localhost:9050"})
+		So(err, ShouldBeNil)
+
+		Convey(`Test when one task ID is given.`, func() {
+			err = t.parse([]string{"onetaskid111"})
+			So(err, ShouldBeNil)
+		})
+
+		Convey(`Make sure that Parse handles when no task ID is given.`, func() {
+			err = t.parse([]string{})
+			So(err, ShouldErrLike, "must specify a swarming task ID")
+		})
+
+		Convey(`Make sure that Parse handles when too many task ID is given.`, func() {
+			err = t.parse([]string{"toomany234", "taskids567"})
+			So(err, ShouldErrLike, "specify only one")
+		})
+	})
+}
+
+func TestCancelTask(t *testing.T) {
+	t.Parallel()
+
+	Convey(`Cancel`, t, func() {
+		ctx := context.Background()
+		t := cancelRun{}
+
+		var givenTaskID string
+		var givenKillRunning bool
+		failTaskID := "failtask"
+
+		service := &testService{
+			cancelTask: func(ctx context.Context, taskID string, req *swarming.SwarmingRpcsTaskCancelRequest) (*swarming.SwarmingRpcsCancelResponse, error) {
+				givenTaskID = taskID
+				givenKillRunning = req.KillRunning
+				res := &swarming.SwarmingRpcsCancelResponse{
+					Ok:         true,
+					WasRunning: givenKillRunning,
+				}
+				if givenTaskID == failTaskID {
+					res.Ok = false
+					return res, nil
+				}
+				return res, nil
+			},
+		}
+
+		Convey(`Cancel task`, func() {
+			err := t.cancelTask(ctx, "task", service)
+			So(err, ShouldBeNil)
+			So(givenTaskID, ShouldEqual, "task")
+		})
+
+		Convey(`Cancel running task `, func() {
+			t.killRunning = true
+			err := t.cancelTask(ctx, "runningtask", service)
+			So(err, ShouldBeNil)
+			So(givenTaskID, ShouldEqual, "runningtask")
+			So(givenKillRunning, ShouldEqual, true)
+		})
+
+		Convey(`Cancel task was not OK`, func() {
+			err := t.cancelTask(ctx, failTaskID, service)
+			So(err, ShouldErrLike, "failed to cancel task")
+			So(givenTaskID, ShouldEqual, failTaskID)
+		})
+	})
+}
diff --git a/client/cmd/swarming/main.go b/client/cmd/swarming/main.go
index d62f6f3..1d7e85f 100644
--- a/client/cmd/swarming/main.go
+++ b/client/cmd/swarming/main.go
@@ -84,6 +84,7 @@
 		// Keep in alphabetical order of their name.
 		Commands: []*subcommands.Command{
 			lib.CmdBots(af),
+			lib.CmdCancelTask(af),
 			lib.CmdCollect(af),
 			lib.CmdDeleteBots(af),
 			lib.CmdReproduce(af),