nomad: add support for querying consistent subset of services

dependency/dependency.go

@@ -53,6 +53,7 @@ type QueryOptions struct {
 	Datacenter        string
 	Region            string
 	Near              string
+	Choose            string
 	RequireConsistent bool
 	VaultGrace        time.Duration
 	WaitIndex         uint64
@@ -92,6 +93,10 @@ func (q *QueryOptions) Merge(o *QueryOptions) *QueryOptions {
 		r.Near = o.Near
 	}
 
+	if o.Choose != "" {
+		r.Choose = o.Choose
+	}
+
 	if o.RequireConsistent != false {
 		r.RequireConsistent = o.RequireConsistent
 	}
@@ -119,9 +124,16 @@ func (q *QueryOptions) ToConsulOpts() *consulapi.QueryOptions {
 }
 
 func (q *QueryOptions) ToNomadOpts() *nomadapi.QueryOptions {
+	var params map[string]string
+	if q.Choose != "" {
+		params = map[string]string{
+			"choose": q.Choose,
+		}
+	}
 	return &nomadapi.QueryOptions{
 		AllowStale: q.AllowStale,
 		Region:     q.Region,
+		Params:     params,
 		WaitIndex:  q.WaitIndex,
 		WaitTime:   q.WaitTime,
 	}
@@ -146,6 +158,10 @@ func (q *QueryOptions) String() string {
 		u.Add("near", q.Near)
 	}
 
+	if q.Choose != "" {
+		u.Add("choose", q.Choose)
+	}
+
 	if q.RequireConsistent {
 		u.Add("consistent", strconv.FormatBool(q.RequireConsistent))
 	}

dependency/nomad_service.go

@@ -17,6 +17,8 @@ var (
 
 	// NomadServiceQueryRe is the regex that is used to understand a service
 	// specific Nomad query.
+	//
+	// e.g. "<tag=value>.<name>@<region>"
 	NomadServiceQueryRe = regexp.MustCompile(`\A` + tagRe + serviceNameRe + regionRe + `\z`)
 )
 
@@ -46,6 +48,7 @@ type NomadServiceQuery struct {
 	region string
 	name   string
 	tag    string
+	choose string
 }
 
 // NewNomadServiceQuery parses a string into a NomadServiceQuery which is
@@ -56,6 +59,7 @@ func NewNomadServiceQuery(s string) (*NomadServiceQuery, error) {
 	}
 
 	m := regexpMatch(NomadServiceQueryRe, s)
+
 	return &NomadServiceQuery{
 		stopCh: make(chan struct{}, 1),
 		region: m["region"],
@@ -64,6 +68,21 @@ func NewNomadServiceQuery(s string) (*NomadServiceQuery, error) {
 	}, nil
 }
 
+// NewNomadServiceChooseQuery parses s using NewNomadServiceQuery, and then also
+// configures the resulting query with the choose parameter set according to the
+// count and key arguments.
+func NewNomadServiceChooseQuery(count int, key, s string) (*NomadServiceQuery, error) {
+	query, err := NewNomadServiceQuery(s)
+	if err != nil {
+		return nil, err
+	}
+
+	choose := fmt.Sprintf("%d|%s", count, key)
+	query.choose = choose
+
+	return query, nil
+}
+
 // Fetch queries the Nomad API defined by the given client and returns a slice
 // of NomadService objects.
 func (d *NomadServiceQuery) Fetch(client *ClientSet, opts *QueryOptions) (interface{}, *ResponseMetadata, error) {
@@ -75,6 +94,7 @@ func (d *NomadServiceQuery) Fetch(client *ClientSet, opts *QueryOptions) (interf
 
 	opts = opts.Merge(&QueryOptions{
 		Region: d.region,
+		Choose: d.choose,
 	})
 
 	u := &url.URL{
@@ -138,6 +158,9 @@ func (d *NomadServiceQuery) String() string {
 	if d.region != "" {
 		name = name + "@" + d.region
 	}
+	if d.choose != "" {
+		name = name + ":" + d.choose
+	}
 	return fmt.Sprintf("nomad.service(%s)", name)
 }
 

dependency/nomad_service_test.go

@@ -9,7 +9,6 @@ import (
 )
 
 func TestNewNomadServiceQuery(t *testing.T) {
-
 	cases := []struct {
 		name string
 		i    string
@@ -97,8 +96,20 @@ func TestNewNomadServiceQuery(t *testing.T) {
 	}
 }
 
-func TestNomadServiceQuery_Fetch(t *testing.T) {
+func TestNomadServiceChooseQuery(t *testing.T) {
+	query, err := NewNomadServiceChooseQuery(4, "abc123", "tag.name@us-east-1")
+	require.NoError(t, err)
+
+	query.stopCh = nil
+	require.Equal(t, &NomadServiceQuery{
+		region: "us-east-1",
+		name:   "name",
+		tag:    "tag",
+		choose: "4|abc123",
+	}, query)
+}
 
+func TestNomadServiceQuery_Fetch(t *testing.T) {
 	cases := []struct {
 		name string
 		i    string
@@ -185,8 +196,59 @@ func TestNomadServiceQuery_Fetch(t *testing.T) {
 	}
 }
 
-func TestNomadServiceQuery_String(t *testing.T) {
+func TestNomadServicesQuery_Fetch_3arg(t *testing.T) {
+	cases := []struct {
+		name    string
+		service string
+		count   int
+		key     string
+		exp     []*NomadService
+	}{
+		{
+			name:    "choose one",
+			service: "example-cache",
+			count:   1,
+			key:     "abc123",
+			exp: []*NomadService{
+				&NomadService{
+					Name:       "example-cache",
+					Address:    "127.0.0.1",
+					Datacenter: "dc1",
+					Tags:       ServiceTags([]string{"tag1", "tag2"}),
+					JobID:      "example",
+				},
+			},
+		},
+	}
+
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
+			d, err := NewNomadServiceChooseQuery(tc.count, tc.key, tc.service)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			actI, _, err := d.Fetch(testClients, nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			act := actI.([]*NomadService)
+
+			for _, s := range act {
+				// clear random fields
+				s.ID = ""
+				s.Node = ""
+				s.Port = 0
+				s.AllocID = ""
+			}
+
+			require.Equal(t, tc.exp, act)
+		})
+	}
+}
 
+func TestNomadServiceQuery_String(t *testing.T) {
 	cases := []struct {
 		name string
 		i    string
@@ -224,3 +286,31 @@ func TestNomadServiceQuery_String(t *testing.T) {
 		})
 	}
 }
+
+func TestNomadServiceQuery_String_3arg(t *testing.T) {
+	cases := []struct {
+		name  string
+		i     string
+		count int
+		key   string
+		exp   string
+	}{
+		{
+			"choose",
+			"redis",
+			3,
+			"abc123",
+			"nomad.service(redis:3|abc123)",
+		},
+	}
+
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
+			d, err := NewNomadServiceChooseQuery(3, "abc123", tc.i)
+			if err != nil {
+				t.Fatal(err)
+			}
+			require.Equal(t, tc.exp, d.String())
+		})
+	}
+}

dependency/nomad_services_test.go

@@ -53,17 +53,16 @@ func TestNewNomadServicesQueryQuery(t *testing.T) {
 	}
 }
 
-func TestNomadServicesQuery_Fetch(t *testing.T) {
-
+func TestNomadServicesQuery_Fetch_1arg(t *testing.T) {
 	cases := []struct {
-		name string
-		i    string
-		exp  []*NomadServicesSnippet
+		name    string
+		service string
+		exp     []*NomadServicesSnippet
 	}{
 		{
-			"all",
-			"",
-			[]*NomadServicesSnippet{
+			name:    "all",
+			service: "",
+			exp: []*NomadServicesSnippet{
 				&NomadServicesSnippet{
 					Name: "example-cache",
 					Tags: ServiceTags([]string{"tag1", "tag2"}),
@@ -74,7 +73,7 @@ func TestNomadServicesQuery_Fetch(t *testing.T) {
 
 	for i, tc := range cases {
 		t.Run(fmt.Sprintf("%d_%s", i, tc.name), func(t *testing.T) {
-			d, err := NewNomadServicesQuery(tc.i)
+			d, err := NewNomadServicesQuery(tc.service)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -88,6 +87,7 @@ func TestNomadServicesQuery_Fetch(t *testing.T) {
 		})
 	}
 }
+
 func TestNomadServicesQuery_String(t *testing.T) {
 
 	cases := []struct {

docs/templating-language.md

@@ -95,6 +95,9 @@ provides the following functions:
   - [modulo](#modulo)
   - [minimum](#minimum)
   - [maximum](#maximum)
+- [Nomad Functions](#nomad-functions)
+  - [nomadServices](#nomadservices)
+  - [nomadService](#nomadservice)
 - [Debugging Functions](#debugging)
   - [spew_dump](#spew_dump)
   - [spew_sdump](#spew_sdump)
@@ -1780,6 +1783,44 @@ This can also be used with a pipe function.
 {{ 5 | maximum 2 }} // 2
 ```
 
+## Nomad Functions
+
+Nomad service registrations can be queried using the `nomadServices` and `nomadService` functions.
+Typically these will be used from within a Nomad [template](https://www.nomadproject.io/docs/job-specification/template#nomad-services) configuration.
+
+### `nomadServices`
+
+This can be used to query the names of services registered in Nomad.
+
+```nomadServices
+{{ range nomadServices }}
+  {{ .Name .Tags }}
+{{ end }}
+```
+
+### `nomadService`
+
+This can be used to query for additional information about each instance of a service registered in Nomad.
+
+```nomadService
+{{ range nomadService "my-app" }}
+  {{ .Address }} {{ .Port }}
+{{ end}}
+```
+
+The `nomadService` function also supports basic load-balancing via a [rendezvous hashing](https://en.wikipedia.org/wiki/Rendezvous_hashing)
+algorithm implemented in Nomad's API. To activate this behavior, the function requires three arguments in this order:
+the number of instances desired, a unique but consistent identifier associated with the requester, and the service name.
+
+Typically the unique identifier would be the allocation ID in a Nomad job.
+
+```nomadService
+{{$allocID := env "NOMAD_ALLOC_ID" -}}
+{{range nomadService 3 $allocID "redis"}}
+  {{.Address}} {{.Port}} | {{.Tags}} @ {{.Datacenter}}
+{{- end}}
+```
+
 ## Debugging Functions
 
 Debugging functions help template developers understand the current context of a template block. These