Support "near" parameter in prepared query service block

command/agent/dns.go

@@ -592,6 +592,10 @@ RPC:
 func (d *DNSServer) preparedQueryLookup(network, datacenter, query string, req, resp *dns.Msg) {
 	// Execute the prepared query.
 	args := structs.PreparedQueryExecuteRequest{
+		Origin: structs.QuerySource{
+			Datacenter: d.agent.config.Datacenter,
+			Node:       d.agent.config.NodeName,
+		},
 		Datacenter:    datacenter,
 		QueryIDOrName: query,
 		QueryOptions: structs.QueryOptions{

command/agent/http.go

@@ -531,13 +531,7 @@ func (s *HTTPServer) parseToken(req *http.Request, token *string) {
 // DC in the request, if given, or else the agent's DC.
 func (s *HTTPServer) parseSource(req *http.Request, source *structs.QuerySource) {
 	s.parseDC(req, &source.Datacenter)
-	if node := req.URL.Query().Get("near"); node != "" {
-		if node == "_agent" {
-			source.Node = s.agent.config.NodeName
-		} else {
-			source.Node = node
-		}
-	}
+	source.Node = req.URL.Query().Get("near")
 }
 
 // parse is a convenience method for endpoints that need

command/agent/http_test.go

@@ -382,18 +382,6 @@ func TestParseSource(t *testing.T) {
 	if source.Datacenter != "foo" || source.Node != "bob" {
 		t.Fatalf("bad: %v", source)
 	}
-
-	// The magic "_agent" node name will use the agent's local node name.
-	req, err = http.NewRequest("GET",
-		"/v1/catalog/nodes?near=_agent", nil)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
-	source = structs.QuerySource{}
-	srv.parseSource(req, &source)
-	if source.Datacenter != "dc1" || source.Node != srv.agent.config.NodeName {
-		t.Fatalf("bad: %v", source)
-	}
 }
 
 func TestParseWait(t *testing.T) {

command/agent/prepared_query_endpoint.go

@@ -95,6 +95,10 @@ func parseLimit(req *http.Request, limit *int) error {
 // preparedQueryExecute executes a prepared query.
 func (s *HTTPServer) preparedQueryExecute(id string, resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	args := structs.PreparedQueryExecuteRequest{
+		Origin: structs.QuerySource{
+			Datacenter: s.agent.config.Datacenter,
+			Node:       s.agent.config.NodeName,
+		},
 		QueryIDOrName: id,
 	}
 	s.parseSource(req, &args.Source)

command/agent/prepared_query_endpoint_test.go

@@ -279,6 +279,10 @@ func TestPreparedQuery_Execute(t *testing.T) {
 
 		m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
 			expected := &structs.PreparedQueryExecuteRequest{
+				Origin: structs.QuerySource{
+					Datacenter: srv.agent.config.Datacenter,
+					Node:       srv.agent.config.NodeName,
+				},
 				Datacenter:    "dc1",
 				QueryIDOrName: "my-id",
 				Limit:         5,

consul/prepared_query/walk_test.go

@@ -20,6 +20,7 @@ func TestWalk_ServiceQuery(t *testing.T) {
 		Failover: structs.QueryDatacenterOptions{
 			Datacenters: []string{"dc1", "dc2"},
 		},
+		Near: "_agent",
 		Tags: []string{"tag1", "tag2", "tag3"},
 	}
 	if err := walk(service, fn); err != nil {
@@ -30,6 +31,7 @@ func TestWalk_ServiceQuery(t *testing.T) {
 		".Service:the-service",
 		".Failover.Datacenters[0]:dc1",
 		".Failover.Datacenters[1]:dc2",
+		".Near:_agent",
 		".Tags[0]:tag1",
 		".Tags[1]:tag2",
 		".Tags[2]:tag3",

consul/prepared_query_endpoint.go

@@ -368,10 +368,44 @@ func (p *PreparedQuery) Execute(args *structs.PreparedQueryExecuteRequest,
 	// Shuffle the results in case coordinates are not available if they
 	// requested an RTT sort.
 	reply.Nodes.Shuffle()
-	if err := p.srv.sortNodesByDistanceFrom(args.Source, reply.Nodes); err != nil {
+
+	// Get the source to sort by. This can be passed in by the requestor, or
+	// pre-defined using the Near parameter in the prepared query. If the
+	// near parameter was defined, that will be preferred.
+	sortFrom := args.Source
+	if query.Service.Near != "" {
+		sortFrom = structs.QuerySource{
+			Datacenter: args.Datacenter,
+			Node:       query.Service.Near,
+		}
+	}
+
+	// Respect the magic "_agent" flag.
+	preferLocal := false
+	if sortFrom.Node == "_agent" {
+		preferLocal = true
+		sortFrom = args.Origin
+	}
+
+	// Perform the distance sort
+	if err := p.srv.sortNodesByDistanceFrom(sortFrom, reply.Nodes); err != nil {
 		return err
 	}
 
+	// Nodes cannot be any "closer" than localhost, so this special case ensures
+	// the local node is returned first if it is present in the result. This
+	// allows the local agent to be preferred even when network coordinates are
+	// not enabled. Only works if the results come from the request origin DC.
+	if preferLocal && reply.Datacenter == args.Origin.Datacenter {
+		for i, node := range reply.Nodes {
+			if node.Node.Node == args.Origin.Node {
+				remote := append(reply.Nodes[:i], reply.Nodes[i+1:]...)
+				reply.Nodes = append([]structs.CheckServiceNode{node}, remote...)
+				break
+			}
+		}
+	}
+
 	// Apply the limit if given.
 	if args.Limit > 0 && len(reply.Nodes) > args.Limit {
 		reply.Nodes = reply.Nodes[:args.Limit]

consul/prepared_query_endpoint_test.go

@@ -1607,6 +1607,137 @@ func TestPreparedQuery_Execute(t *testing.T) {
 		t.Fatalf("unique shuffle ratio too low: %d/100", len(uniques))
 	}
 
+	// Set the query to prefer a colocated service using the magic _agent token
+	query.Op = structs.PreparedQueryUpdate
+	query.Query.Service.Near = "_agent"
+	if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Apply", &query, &query.Query.ID); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Now try querying and make sure the local node is preferred.
+	{
+		req := structs.PreparedQueryExecuteRequest{
+			Origin: structs.QuerySource{
+				Datacenter: "dc1",
+				Node:       "node1",
+			},
+			Datacenter:    "dc1",
+			QueryIDOrName: query.Query.ID,
+			QueryOptions:  structs.QueryOptions{Token: execToken},
+		}
+
+		var reply structs.PreparedQueryExecuteResponse
+
+		// Repeat this a few times to make sure we don't just get lucky.
+		for i := 0; i < 10; i++ {
+			if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Execute", &req, &reply); err != nil {
+				t.Fatalf("err: %v", err)
+			}
+			if n := len(reply.Nodes); n != 10 {
+				t.Fatalf("expect 10 nodes, got: %d", n)
+			}
+			if node := reply.Nodes[0].Node.Node; node != "node1" {
+				t.Fatalf("expect node1 first, got: %q", node)
+			}
+		}
+	}
+
+	// Falls back to remote nodes if service is not available locally
+	{
+		req := structs.PreparedQueryExecuteRequest{
+			Origin: structs.QuerySource{
+				Datacenter: "dc1",
+				Node:       "node1",
+			},
+			Datacenter:    "dc1",
+			QueryIDOrName: query.Query.ID,
+			QueryOptions:  structs.QueryOptions{Token: execToken},
+		}
+
+		var reply structs.PreparedQueryExecuteResponse
+
+		if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Execute", &req, &reply); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if n := len(reply.Nodes); n != 10 {
+			t.Fatalf("expect 10 nodes, got: %d", n)
+		}
+	}
+
+	// Shuffles if the response comes from a non-local DC. We may
+	// need to try multiple times if at first we get a match by chance.
+	{
+		req := structs.PreparedQueryExecuteRequest{
+			Origin: structs.QuerySource{
+				Datacenter: "dc2",
+				Node:       "node1",
+			},
+			Datacenter:    "dc1",
+			QueryIDOrName: query.Query.ID,
+			QueryOptions:  structs.QueryOptions{Token: execToken},
+		}
+
+		var reply structs.PreparedQueryExecuteResponse
+
+		shuffled := false
+		for i := 0; i < 10; i++ {
+			if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Execute", &req, &reply); err != nil {
+				t.Fatalf("err: %v", err)
+			}
+			if n := len(reply.Nodes); n != 10 {
+				t.Fatalf("expect 10 nodes, got: %d", n)
+			}
+			if reply.Nodes[0].Node.Node != "node1" {
+				shuffled = true
+				break
+			}
+		}
+
+		if !shuffled {
+			t.Fatal("expect node shuffle for remote results")
+		}
+	}
+
+	// Bake a non-local node name into Near parameter of the query. This
+	// node was seeded with a coordinate above so distance sort works.
+	query.Query.Service.Near = "node3"
+	if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Apply", &query, &query.Query.ID); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Try the distance sort again to ensure the nearest node is returned
+	{
+		req := structs.PreparedQueryExecuteRequest{
+			Origin: structs.QuerySource{
+				Datacenter: "dc1",
+				Node:       "node1",
+			},
+			Datacenter:    "dc1",
+			QueryIDOrName: query.Query.ID,
+			QueryOptions:  structs.QueryOptions{Token: execToken},
+		}
+
+		var reply structs.PreparedQueryExecuteResponse
+
+		for i := 0; i < 10; i++ {
+			if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Execute", &req, &reply); err != nil {
+				t.Fatalf("err: %v", err)
+			}
+			if n := len(reply.Nodes); n != 10 {
+				t.Fatalf("expect 10 nodes, got: %d", n)
+			}
+			if node := reply.Nodes[0].Node.Node; node != "node3" {
+				t.Fatalf("expect node3, got: %q", node)
+			}
+		}
+	}
+
+	// Un-bake the Near parameter.
+	query.Query.Service.Near = ""
+	if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Apply", &query, &query.Query.ID); err != nil {
+		t.Fatalf("err:% v", err)
+	}
+
 	// Update the health of a node to mark it critical.
 	setHealth := func(node string, health string) {
 		req := structs.RegisterRequest{
@@ -1683,7 +1814,6 @@ func TestPreparedQuery_Execute(t *testing.T) {
 	}
 
 	// Make the query more picky so it excludes warning nodes.
-	query.Op = structs.PreparedQueryUpdate
 	query.Query.Service.OnlyPassing = true
 	if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Apply", &query, &query.Query.ID); err != nil {
 		t.Fatalf("err: %v", err)

consul/structs/prepared_query.go

@@ -34,6 +34,11 @@ type ServiceQuery struct {
 	// discarded)
 	OnlyPassing bool
 
+	// Near allows the query to always prefer the node nearest the given
+	// node. If the node does not exist, results are returned in their
+	// normal randomly-shuffled order.
+	Near string
+
 	// Tags are a set of required and/or disallowed tags. If a tag is in
 	// this list it must be present. If the tag is preceded with "!" then
 	// it is disallowed.
@@ -173,6 +178,12 @@ type PreparedQueryExecuteRequest struct {
 	// Limit will trim the resulting list down to the given limit.
 	Limit int
 
+	// Origin is used to carry around a reference to the node which
+	// is executing the query on behalf of the client. It is taken
+	// as a QuerySource so that it can be used directly for queries
+	// relating to the agent servicing the request.
+	Origin QuerySource
+
 	// Source is used to sort the results relative to a given node using
 	// network coordinates.
 	Source QuerySource