Add force-grace-period argument to step cancel command

buildkite · Nov 11, 2024 · ba910f9 · ba910f9
1 parent f29fc84
commit ba910f9
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 10 deletions.
diff --git a/api/steps.go b/api/steps.go
@@ -56,8 +56,9 @@ func (c *Client) StepUpdate(ctx context.Context, stepIdOrKey string, stepUpdate
 }
 
 type StepCancel struct {
-	Build string `json:"build_id"`
-	Force bool   `json:"force,omitempty"`
+	Build            string `json:"build_id"`
+	Force            bool   `json:"force"`
+	ForceGracePeriod int64  `json:"force_grace_period"`
 }
 
 type StepCancelResponse struct {

diff --git a/clicommand/step_cancel.go b/clicommand/step_cancel.go
@@ -22,12 +22,15 @@ Cancel all unfinished jobs for a step
 Example:
 
     $ buildkite-agent step cancel --step "key"
-    $ buildkite-agent step cancel --step "key" --force`
+    $ buildkite-agent step cancel --step "key" --force
+    $ buildkite-agent step cancel --step "key" --force --force-grace-period 30
+		`
 
 type StepCancelConfig struct {
-	StepOrKey string `cli:"step" validate:"required"`
-	Force     bool   `cli:"force"`
-	Build     string `cli:"build"`
+	StepOrKey        string `cli:"step" validate:"required"`
+	Force            bool   `cli:"force"`
+	ForceGracePeriod int64  `cli:"force-grace-period"`
+	Build            string `cli:"build"`
 
 	// Global flags
 	Debug       bool     `cli:"debug"`
@@ -63,10 +66,17 @@ var StepCancelCommand = cli.Command{
 		},
 		cli.BoolFlag{
 			Name:   "force",
-			Usage:  "Don't wait for the agent to finish before cancelling the jobs",
+			Usage:  "Don't wait for the agent to finish before cancelling the jobs. Used in conjunction with `--force-grace-period` to immediately cancel the job",
 			EnvVar: "BUILDKITE_STEP_CANCEL_FORCE",
 		},
 
+		cli.Int64Flag{
+			Name:   "force-grace-period",
+			Value:  defaultCancelGracePeriod,
+			Usage:  "The number of seconds to wait for agents to finish uploading artifacts before cancelling their jobs",
+			EnvVar: "BUILDKITE_FORCE_GRACE_PERIOD,BUILDKITE_CANCEL_GRACE_PERIOD",
+		},
+
 		// API Flags
 		AgentAccessTokenFlag,
 		EndpointFlag,
@@ -84,6 +94,10 @@ var StepCancelCommand = cli.Command{
 		ctx, cfg, l, _, done := setupLoggerAndConfig[StepCancelConfig](context.Background(), c)
 		defer done()
 
+		if cfg.ForceGracePeriod < 0 {
+			return fmt.Errorf("The value of `--force-grace-period` must be greater than or equal to 0")
+		}
+
 		return cancelStep(ctx, cfg, l)
 	},
 }
@@ -92,10 +106,10 @@ func cancelStep(ctx context.Context, cfg StepCancelConfig, l logger.Logger) erro
 	// Create the API client
 	client := api.NewClient(l, loadAPIClientConfig(cfg, "AgentAccessToken"))
 
-	// Create the value to cancel
 	cancel := &api.StepCancel{
-		Build: cfg.Build,
-		Force: cfg.Force,
+		Build:            cfg.Build,
+		Force:            cfg.Force,
+		ForceGracePeriod: cfg.ForceGracePeriod,
 	}
 
 	// Post the change

diff --git a/clicommand/step_cancel_test.go b/clicommand/step_cancel_test.go
@@ -21,6 +21,7 @@ func TestStepCancel(t *testing.T) {
 		}))
 
 		cfg := StepCancelConfig{
+			ForceGracePeriod: 10,
 			Force:            true,
 			Build:            "1",
 			StepOrKey:        "some-random-key",
@@ -40,6 +41,7 @@ func TestStepCancel(t *testing.T) {
 		}))
 
 		cfg := StepCancelConfig{
+			ForceGracePeriod: 10,
 			Force:            true,
 			StepOrKey:        "some-random-key",
 			AgentAccessToken: "agentaccesstoken",