From e54d5402e2a65158a19ee7da1ceab537ef5f0faa Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sun, 6 Sep 2015 15:46:45 -0700
Subject: [PATCH] nomad: adding alloc lookup method

---
 nomad/alloc_endpoint.go      | 36 ++++++++++++++++++++++++++++++++++++
 nomad/alloc_endpoint_test.go | 33 +++++++++++++++++++++++++++++++++
 nomad/structs/structs.go     | 12 ++++++++++++
 3 files changed, 81 insertions(+)

diff --git a/nomad/alloc_endpoint.go b/nomad/alloc_endpoint.go
index 7d811adab..62e10c57a 100644
--- a/nomad/alloc_endpoint.go
+++ b/nomad/alloc_endpoint.go
@@ -49,3 +49,39 @@ func (a *Alloc) List(args *structs.AllocListRequest, reply *structs.AllocListRes
 	a.srv.setQueryMeta(&reply.QueryMeta)
 	return nil
 }
+
+// GetAlloc is used to lookup a particular allocation
+func (a *Alloc) GetAlloc(args *structs.AllocSpecificRequest,
+	reply *structs.SingleAllocResponse) error {
+	if done, err := a.srv.forward("Alloc.GetAlloc", args, args, reply); done {
+		return err
+	}
+	defer metrics.MeasureSince([]string{"nomad", "alloc", "get_alloc"}, time.Now())
+
+	// Lookup the allocation
+	snap, err := a.srv.fsm.State().Snapshot()
+	if err != nil {
+		return err
+	}
+	out, err := snap.GetAllocByID(args.AllocID)
+	if err != nil {
+		return err
+	}
+
+	// Setup the output
+	if out != nil {
+		reply.Alloc = out
+		reply.Index = out.ModifyIndex
+	} else {
+		// Use the last index that affected the nodes table
+		index, err := snap.GetIndex("allocs")
+		if err != nil {
+			return err
+		}
+		reply.Index = index
+	}
+
+	// Set the query response
+	a.srv.setQueryMeta(&reply.QueryMeta)
+	return nil
+}
diff --git a/nomad/alloc_endpoint_test.go b/nomad/alloc_endpoint_test.go
index 314b9d897..697fb4bf2 100644
--- a/nomad/alloc_endpoint_test.go
+++ b/nomad/alloc_endpoint_test.go
@@ -1,6 +1,7 @@
 package nomad
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/hashicorp/net-rpc-msgpackrpc"
@@ -42,3 +43,35 @@ func TestAllocEndpoint_List(t *testing.T) {
 		t.Fatalf("bad: %#v", resp.Allocations[0])
 	}
 }
+
+func TestAllocEndpoint_GetAlloc(t *testing.T) {
+	s1 := testServer(t, nil)
+	defer s1.Shutdown()
+	codec := rpcClient(t, s1)
+	testutil.WaitForLeader(t, s1.RPC)
+
+	// Create the register request
+	alloc := mock.Alloc()
+	state := s1.fsm.State()
+	err := state.UpdateAllocations(1000, []*structs.Allocation{alloc})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Lookup the jobs
+	get := &structs.AllocSpecificRequest{
+		AllocID:      alloc.ID,
+		QueryOptions: structs.QueryOptions{Region: "region1"},
+	}
+	var resp structs.SingleAllocResponse
+	if err := msgpackrpc.CallWithCodec(codec, "Alloc.GetAlloc", get, &resp); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if resp.Index != 1000 {
+		t.Fatalf("Bad index: %d %d", resp.Index, 1000)
+	}
+
+	if !reflect.DeepEqual(alloc, resp.Alloc) {
+		t.Fatalf("bad: %#v", resp.Alloc)
+	}
+}
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index bb831c3af..1f8053d57 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -238,6 +238,12 @@ type AllocListRequest struct {
 	QueryOptions
 }
 
+// AllocSpecificRequest is used to query a specific allocation
+type AllocSpecificRequest struct {
+	AllocID string
+	QueryOptions
+}
+
 // GenericRequest is used to request where no
 // specific information is needed.
 type GenericRequest struct {
@@ -318,6 +324,12 @@ type JobListResponse struct {
 	QueryMeta
 }
 
+// SingleAllocResponse is used to return a single allocation
+type SingleAllocResponse struct {
+	Alloc *Allocation
+	QueryMeta
+}
+
 // JobAllocationsResponse is used to return the allocations for a job
 type JobAllocationsResponse struct {
 	Allocations []*Allocation