fix(nomad): keep the counts a job runs with on an edit, unless the file changes them

docs/jobs.md

@@ -114,6 +114,15 @@ the cluster for a plan and shows it. If the cluster cannot plan the file, for
 example because of a syntax error, urga opens the editor again with your edit
 and the reason at the top. See [Editor](configuration.md#editor).
 
+An edit keeps the count a group runs with. If the autoscaler or `s` scaled a
+group to 5 and you leave its `count` in the file alone, the job is submitted
+with 5, not with the number in the file. If you change the `count` in the
+editor, the number you wrote is used. A group the job does not run yet gets
+the count of the file. The plan lists the counts it kept above what will
+happen. When the cluster cannot read the counts of the file as it was
+opened, for example a file that calls `file()`, the editor still opens, and
+the plan says that every count is the one in the file.
+
 ## Creating a job
 
 `n` opens a template of a job in your editor: a service with one group `web`

internal/nomad/edit.go

@@ -21,9 +21,10 @@ import (
 // makes, and so are the values of its variables: the next edit opens the same
 // file, and the job runs with the values it ran with before. It is submitted
 // at the index of its plan: a job that changed since is refused rather than
-// overwritten.
-func (c *Client) SubmitJob(ctx context.Context, namespace, source string, vars JobVariables, index uint64) error {
-	job, kept, err := c.jobOf(ctx, namespace, source, vars)
+// overwritten. The counts it opened with decide which counts the job keeps,
+// see keepCounts.
+func (c *Client) SubmitJob(ctx context.Context, namespace, source string, vars JobVariables, index uint64, opened map[string]int) error {
+	job, kept, _, err := c.jobOf(ctx, namespace, source, vars, opened)
 	if err != nil {
 		return err
 	}
@@ -38,18 +39,95 @@ func (c *Client) SubmitJob(ctx context.Context, namespace, source string, vars J
 	return err
 }
 
-// jobOf reads a job file and puts the job in the namespace it is sent to.
-func (c *Client) jobOf(ctx context.Context, namespace, source string, vars JobVariables) (*api.Job, *api.JobSubmission, error) {
-	job, kept, err := c.parseJob(ctx, namespace, source, vars, true)
+// jobOf reads a job file and puts the job in the namespace it is sent to,
+// with the counts it keeps from the job that runs.
+func (c *Client) jobOf(ctx context.Context, namespace, source string, vars JobVariables, opened map[string]int) (*api.Job, *api.JobSubmission, []KeptCount, error) {
+	job, submission, err := c.parseJob(ctx, namespace, source, vars, true)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, nil, err
 	}
 
 	if namespace != "" && namespace != AllNamespaces {
 		job.Namespace = &namespace
 	}
 
-	return job, kept, nil
+	kept, err := c.keepCounts(ctx, namespace, job, opened)
+
+	return job, submission, kept, err
+}
+
+// KeptCount is a group that keeps the count its job runs with, though the
+// file says another.
+type KeptCount struct {
+	Group string
+	Kept  int
+	File  int
+}
+
+// keepCounts sets the count of each group of an edited job. A group whose
+// count is as in the file when it was opened keeps the count the job runs it
+// with, which the autoscaler or a person may have changed since. A group
+// whose count was edited, a new group, and a job that is not running take
+// the count of the file. Nil opened counts keep nothing.
+func (c *Client) keepCounts(ctx context.Context, namespace string, job *api.Job, opened map[string]int) ([]KeptCount, error) {
+	if opened == nil {
+		return nil, nil
+	}
+
+	running, _, err := c.api.Jobs().Info(valueOf(job.ID), c.query(ctx, namespace))
+	if answered(err, http.StatusNotFound) {
+		return nil, nil
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	var kept []KeptCount
+
+	for _, group := range job.TaskGroups {
+		name := valueOf(group.Name)
+		file := countOf(group)
+
+		was, wasOpened := opened[name]
+		now := slices.IndexFunc(running.TaskGroups, func(g *api.TaskGroup) bool { return valueOf(g.Name) == name })
+
+		if !wasOpened || was != file || now < 0 {
+			continue
+		}
+
+		if count := countOf(running.TaskGroups[now]); count != file {
+			group.Count = &count
+			kept = append(kept, KeptCount{Group: name, Kept: count, File: file})
+		}
+	}
+
+	return kept, nil
+}
+
+// countOf is the count of a group, which is 1 when the file leaves it out.
+func countOf(group *api.TaskGroup) int {
+	if group.Count == nil {
+		return 1
+	}
+
+	return *group.Count
+}
+
+// JobCounts reads the count of each group of a job file, the way a plan of
+// it would.
+func (c *Client) JobCounts(ctx context.Context, namespace, source string, vars JobVariables) (map[string]int, error) {
+	job, _, _, err := c.jobOf(ctx, namespace, source, vars, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	counts := make(map[string]int, len(job.TaskGroups))
+	for _, group := range job.TaskGroups {
+		counts[valueOf(group.Name)] = countOf(group)
+	}
+
+	return counts, nil
 }
 
 // parseJob reads a job file the way the cluster does, and returns what of it

internal/nomad/edit_test.go

@@ -81,7 +81,7 @@ func TestSubmitJob_JSON(t *testing.T) {
 	client, asked := clusterServer(t, map[string]string{"/v1/jobs": `{"EvalID": "eval-1"}`})
 
 	source := `{"ID": "web", "Name": "web"}`
-	r.NoError(client.SubmitJob(context.Background(), "production", source, nomad.JobVariables{}, 0))
+	r.NoError(client.SubmitJob(context.Background(), "production", source, nomad.JobVariables{}, 0, nil))
 
 	// A file that is already JSON goes straight to the cluster, and is saved
 	// with the new version: the next edit opens it again.
@@ -98,7 +98,7 @@ func TestSubmitJob_JSONUnderAJobKey(t *testing.T) {
 	client, asked := clusterServer(t, map[string]string{"/v1/jobs": `{"EvalID": "eval-1"}`})
 
 	source := `{"Job": {"ID": "web", "Name": "web"}}`
-	r.NoError(client.SubmitJob(context.Background(), "production", source, nomad.JobVariables{}, 0))
+	r.NoError(client.SubmitJob(context.Background(), "production", source, nomad.JobVariables{}, 0, nil))
 
 	// A JSON job file holds the job or wraps it in a Job key, and the nomad
 	// command keeps either one as it was written. Read as the job itself,
@@ -118,7 +118,7 @@ func TestSubmitJob_HCLGoesThroughTheCluster(t *testing.T) {
 		"/v1/jobs":       `{"EvalID": "eval-1"}`,
 	})
 
-	r.NoError(client.SubmitJob(context.Background(), "production", "job \"web\" {}", nomad.JobVariables{}, 0))
+	r.NoError(client.SubmitJob(context.Background(), "production", "job \"web\" {}", nomad.JobVariables{}, 0, nil))
 
 	// HCL is read by Nomad itself, urga does not parse job files.
 	r.Len(*asked, 2)
@@ -134,7 +134,7 @@ func TestSubmitJob_ParsesInTheNamespace(t *testing.T) {
 		"/v1/jobs":       `{"EvalID": "eval-1"}`,
 	})
 
-	r.NoError(client.SubmitJob(context.Background(), "production", "job \"web\" {}", nomad.JobVariables{}, 0))
+	r.NoError(client.SubmitJob(context.Background(), "production", "job \"web\" {}", nomad.JobVariables{}, 0, nil))
 
 	// A token may parse jobs in its own namespace only. Asking without one
 	// asks in the default namespace, where it is refused.
@@ -155,7 +155,7 @@ func TestSubmitJob_HCLKeepsTheSourceAndVariables(t *testing.T) {
 		File:  "count = 3\n",
 	}
 
-	r.NoError(client.SubmitJob(context.Background(), "production", source, vars, 0))
+	r.NoError(client.SubmitJob(context.Background(), "production", source, vars, 0, nil))
 
 	// The cluster reads variables only as a file, so the flags are written
 	// into one next to the file the job was run with.
@@ -181,7 +181,7 @@ func TestSubmitJob_FlagValuesStayText(t *testing.T) {
 
 	vars := nomad.JobVariables{Flags: map[string]string{"greeting": "say \"hi\"\n${name} %{ok}\\"}}
 
-	r.NoError(client.SubmitJob(context.Background(), "production", "job \"web\" {}", vars, 0))
+	r.NoError(client.SubmitJob(context.Background(), "production", "job \"web\" {}", vars, 0, nil))
 
 	// A flag was typed as text. Written into HCL as it is, a quote ends the
 	// string early and ${ reads another variable instead of the characters.
@@ -193,7 +193,7 @@ func TestSubmitJob_BrokenJSON(t *testing.T) {
 
 	client, _ := recorder(t, `{}`)
 
-	err := client.SubmitJob(context.Background(), "production", "{not json", nomad.JobVariables{}, 0)
+	err := client.SubmitJob(context.Background(), "production", "{not json", nomad.JobVariables{}, 0, nil)
 
 	r.ErrorContains(err, "not valid JSON")
 }
@@ -308,7 +308,7 @@ func variablesSent(t *testing.T, vars nomad.JobVariables) string {
 		"/v1/jobs":       `{"EvalID": "eval-1"}`,
 	})
 
-	require.NoError(t, client.SubmitJob(context.Background(), "production", "job \"web\" {}", vars, 0))
+	require.NoError(t, client.SubmitJob(context.Background(), "production", "job \"web\" {}", vars, 0, nil))
 
 	text, ok := (*asked)[0].body["Variables"].(string)
 	require.True(t, ok)
@@ -386,7 +386,7 @@ func TestSubmitJob_VariablesFileThatDoesNotParse(t *testing.T) {
 	client, _ := recorder(t, `{}`)
 
 	err := client.SubmitJob(context.Background(), "production", "job \"web\" {}",
-		nomad.JobVariables{Flags: map[string]string{"a": "b"}, File: "count = = 3"}, 0)
+		nomad.JobVariables{Flags: map[string]string{"a": "b"}, File: "count = = 3"}, 0, nil)
 
 	r.ErrorContains(err, "the variables file of the job does not parse")
 }
@@ -407,7 +407,7 @@ func TestSubmitJob_JSONVariablesFileThatDoesNotParse(t *testing.T) {
 		nomad.JobVariables{
 			Flags: map[string]string{"a": "b"},
 			File:  "{\"image\": \"v1\"}\n{\"count\": 3}\n",
-		}, 0)
+		}, 0, nil)
 
 	require.ErrorContains(t, err, "the variables file of the job does not parse")
 }
@@ -417,3 +417,20 @@ func TestSubmitJob_FileAloneIsSentAsWrittenWhateverItHolds(t *testing.T) {
 		require.Equal(t, file, variablesSent(t, nomad.JobVariables{File: file}))
 	}
 }
+
+func TestJobCounts_AGroupWithoutACountHasOne(t *testing.T) {
+	r := require.New(t)
+
+	client, asked := clusterServer(t, map[string]string{"/v1/jobs/parse": `{"ID": "web", "TaskGroups": [
+		{"Name": "web", "Count": 3},
+		{"Name": "api"}
+	]}`})
+
+	counts, err := client.JobCounts(context.Background(), "production", "job \"web\" {}", nomad.JobVariables{})
+	r.NoError(err)
+
+	r.Equal(map[string]int{"web": 3, "api": 1}, counts)
+
+	// Read the way the plan reads it: the cluster fills in what is left out.
+	r.Equal(true, (*asked)[0].body["Canonicalize"])
+}

internal/nomad/plan.go

@@ -29,6 +29,9 @@ type Plan struct {
 
 	Warnings string
 
+	// Kept are the groups that keep the count the job runs with.
+	Kept []KeptCount
+
 	// Index is the one the job had when it was planned. Submitting at it
 	// fails when the job changed in between.
 	Index uint64
@@ -55,9 +58,11 @@ type PlanGroup struct {
 	Ignore int
 }
 
-// PlanJob asks the cluster what submitting the job file would do.
-func (c *Client) PlanJob(ctx context.Context, namespace, source string, vars JobVariables) (Plan, error) {
-	job, _, err := c.jobOf(ctx, namespace, source, vars)
+// PlanJob asks the cluster what submitting the job file would do. The counts
+// the file had when it was opened decide which counts the job keeps, as they
+// do for SubmitJob; nil keeps none.
+func (c *Client) PlanJob(ctx context.Context, namespace, source string, vars JobVariables, opened map[string]int) (Plan, error) {
+	job, _, kept, err := c.jobOf(ctx, namespace, source, vars, opened)
 	if err != nil {
 		return Plan{}, err
 	}
@@ -67,7 +72,10 @@ func (c *Client) PlanJob(ctx context.Context, namespace, source string, vars Job
 		return Plan{}, err
 	}
 
-	return newPlan(answer), nil
+	plan := newPlan(answer)
+	plan.Kept = kept
+
+	return plan, nil
 }
 
 func newPlan(answer *api.JobPlanResponse) Plan {