Making the status command return the allocs of currently registered job

Author: diptanu

api/jobs.go

@@ -2,7 +2,9 @@ package api
 
 import (
 	"fmt"
+	"net/url"
 	"sort"
+	"strconv"
 	"time"
 )
 
@@ -89,9 +91,18 @@ func (j *Jobs) Info(jobID string, q *QueryOptions) (*Job, *QueryMeta, error) {
 }
 
 // Allocations is used to return the allocs for a given job ID.
-func (j *Jobs) Allocations(jobID string, q *QueryOptions) ([]*AllocationListStub, *QueryMeta, error) {
+func (j *Jobs) Allocations(jobID string, allAllocs bool, q *QueryOptions) ([]*AllocationListStub, *QueryMeta, error) {
 	var resp []*AllocationListStub
-	qm, err := j.client.query("/v1/job/"+jobID+"/allocations", &resp, q)
+	u, err := url.Parse("/v1/job/" + jobID + "/allocations")
+	if err != nil {
+		return nil, nil, err
+	}
+
+	v := u.Query()
+	v.Add("all", strconv.FormatBool(allAllocs))
+	u.RawQuery = v.Encode()
+
+	qm, err := j.client.query(u.String(), &resp, q)
 	if err != nil {
 		return nil, nil, err
 	}

command/agent/job_endpoint.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"net/http"
+	"strconv"
 	"strings"
 
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -130,8 +131,11 @@ func (s *HTTPServer) jobAllocations(resp http.ResponseWriter, req *http.Request,
 	if req.Method != "GET" {
 		return nil, CodedError(405, ErrInvalidMethod)
 	}
+	allAllocs, _ := strconv.ParseBool(req.URL.Query().Get("all"))
+
 	args := structs.JobSpecificRequest{
-		JobID: jobName,
+		JobID:     jobName,
+		AllAllocs: allAllocs,
 	}
 	if s.parse(resp, req, &args.Region, &args.QueryOptions) {
 		return nil, nil

command/fs.go

@@ -348,7 +348,7 @@ func (f *FSCommand) followFile(client *api.Client, alloc *api.Allocation,
 // but use a dead allocation if no running allocations are found
 func getRandomJobAlloc(client *api.Client, jobID string) (string, error) {
 	var runningAllocs []*api.AllocationListStub
-	allocs, _, err := client.Jobs().Allocations(jobID, nil)
+	allocs, _, err := client.Jobs().Allocations(jobID, false, nil)
 
 	// Check that the job actually has allocations
 	if len(allocs) == 0 {

command/status.go

@@ -20,9 +20,10 @@ const (
 
 type StatusCommand struct {
 	Meta
-	length  int
-	evals   bool
-	verbose bool
+	length    int
+	evals     bool
+	allAllocs bool
+	verbose   bool
 }
 
 func (c *StatusCommand) Help() string {
@@ -45,6 +46,9 @@ Status Options:
   -evals
     Display the evaluations associated with the job.
 
+  -all-allocs
+    Display the allocs that matches the job name.
+
   -verbose
     Display full information.
 `
@@ -62,6 +66,7 @@ func (c *StatusCommand) Run(args []string) int {
 	flags.Usage = func() { c.Ui.Output(c.Help()) }
 	flags.BoolVar(&short, "short", false, "")
 	flags.BoolVar(&c.evals, "evals", false, "")
+	flags.BoolVar(&c.allAllocs, "all-allocs", false, "")
 	flags.BoolVar(&c.verbose, "verbose", false, "")
 
 	if err := flags.Parse(args); err != nil {
@@ -218,7 +223,7 @@ func (c *StatusCommand) outputJobInfo(client *api.Client, job *api.Job) error {
 	var evals, allocs []string
 
 	// Query the allocations
-	jobAllocs, _, err := client.Jobs().Allocations(job.ID, nil)
+	jobAllocs, _, err := client.Jobs().Allocations(job.ID, c.allAllocs, nil)
 	if err != nil {
 		return fmt.Errorf("Error querying job allocations: %s", err)
 	}

nomad/job_endpoint.go

@@ -534,7 +534,7 @@ func (j *Job) Allocations(args *structs.JobSpecificRequest,
 			if err != nil {
 				return err
 			}
-			allocs, err := snap.AllocsByJob(args.JobID)
+			allocs, err := snap.AllocsByJob(args.JobID, args.AllAllocs)
 			if err != nil {
 				return err
 			}

nomad/state/state_store.go

@@ -1073,9 +1073,19 @@ func (s *StateStore) AllocsByNodeTerminal(node string, terminal bool) ([]*struct
 }
 
 // AllocsByJob returns all the allocations by job id
-func (s *StateStore) AllocsByJob(jobID string) ([]*structs.Allocation, error) {
+func (s *StateStore) AllocsByJob(jobID string, all bool) ([]*structs.Allocation, error) {
 	txn := s.db.Txn(false)
 
+	// Get the job
+	var job *structs.Job
+	rawJob, err := txn.First("jobs", "id", jobID)
+	if err != nil {
+		return nil, err
+	}
+	if rawJob != nil {
+		job = rawJob.(*structs.Job)
+	}
+
 	// Get an iterator over the node allocations
 	iter, err := txn.Get("allocs", "job", jobID)
 	if err != nil {
@@ -1088,6 +1098,14 @@ func (s *StateStore) AllocsByJob(jobID string) ([]*structs.Allocation, error) {
 		if raw == nil {
 			break
 		}
+
+		alloc := raw.(*structs.Allocation)
+		// If the allocation belongs to a job with the same ID but a diff create
+		// index and we are not getting all the allocations whose Jobs matches
+		// the same Job ID then we skip it
+		if !all && job != nil && alloc.Job.CreateIndex != job.CreateIndex {
+			continue
+		}
 		out = append(out, raw.(*structs.Allocation))
 	}
 	return out, nil

nomad/state/state_store_test.go

@@ -2362,7 +2362,7 @@ func TestStateStore_AllocsByJob(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
-	out, err := state.AllocsByJob("foo")
+	out, err := state.AllocsByJob("foo", false)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -2375,6 +2375,61 @@ func TestStateStore_AllocsByJob(t *testing.T) {
 	}
 }
 
+func TestStateStore_AllocsForRegisteredJob(t *testing.T) {
+	state := testStateStore(t)
+	var allocs []*structs.Allocation
+	var allocs1 []*structs.Allocation
+
+	job := mock.Job()
+	job.ID = "foo"
+	state.UpsertJob(100, job)
+	for i := 0; i < 3; i++ {
+		alloc := mock.Alloc()
+		alloc.Job = job
+		alloc.JobID = job.ID
+		allocs = append(allocs, alloc)
+	}
+	if err := state.UpsertAllocs(200, allocs); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if err := state.DeleteJob(250, job.ID); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	job1 := mock.Job()
+	job1.ID = "foo"
+	job1.CreateIndex = 50
+	state.UpsertJob(300, job1)
+	for i := 0; i < 4; i++ {
+		alloc := mock.Alloc()
+		alloc.Job = job1
+		alloc.JobID = job1.ID
+		allocs1 = append(allocs1, alloc)
+	}
+
+	if err := state.UpsertAllocs(1000, allocs1); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	out, err := state.AllocsByJob(job1.ID, true)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	expected := len(allocs) + len(allocs1)
+	if len(out) != expected {
+		t.Fatalf("expected: %v, actual: %v", expected, len(out))
+	}
+
+	out1, err := state.AllocsByJob(job1.ID, false)
+	expected = len(allocs1)
+	if len(out1) != expected {
+		t.Fatalf("expected: %v, actual: %v", expected, len(out1))
+	}
+
+}
+
 func TestStateStore_AllocsByIDPrefix(t *testing.T) {
 	state := testStateStore(t)
 	var allocs []*structs.Allocation

nomad/structs/structs.go

@@ -249,7 +249,8 @@ type JobEvaluateRequest struct {
 
 // JobSpecificRequest is used when we just need to specify a target job
 type JobSpecificRequest struct {
-	JobID string
+	JobID     string
+	AllAllocs bool
 	QueryOptions
 }
 

scheduler/generic_sched.go

@@ -354,7 +354,7 @@ func (s *GenericScheduler) computeJobAllocs() error {
 	}
 
 	// Lookup the allocations by JobID
-	allocs, err := s.state.AllocsByJob(s.eval.JobID)
+	allocs, err := s.state.AllocsByJob(s.eval.JobID, true)
 	if err != nil {
 		return fmt.Errorf("failed to get allocs for job '%s': %v",
 			s.eval.JobID, err)

scheduler/scheduler.go

@@ -66,7 +66,7 @@ type State interface {
 	Nodes() (memdb.ResultIterator, error)
 
 	// AllocsByJob returns the allocations by JobID
-	AllocsByJob(jobID string) ([]*structs.Allocation, error)
+	AllocsByJob(jobID string, all bool) ([]*structs.Allocation, error)
 
 	// AllocsByNode returns all the allocations by node
 	AllocsByNode(node string) ([]*structs.Allocation, error)

scheduler/system_sched.go

@@ -178,7 +178,7 @@ func (s *SystemScheduler) process() (bool, error) {
 // existing allocations and node status to update the allocations.
 func (s *SystemScheduler) computeJobAllocs() error {
 	// Lookup the allocations by JobID
-	allocs, err := s.state.AllocsByJob(s.eval.JobID)
+	allocs, err := s.state.AllocsByJob(s.eval.JobID, true)
 	if err != nil {
 		return fmt.Errorf("failed to get allocs for job '%s': %v",
 			s.eval.JobID, err)