diff --git a/consul/prepared_query_endpoint.go b/consul/prepared_query_endpoint.go index 5f7d78894..47b6fe1a7 100644 --- a/consul/prepared_query_endpoint.go +++ b/consul/prepared_query_endpoint.go @@ -390,6 +390,23 @@ func (p *PreparedQuery) Execute(args *structs.PreparedQueryExecuteRequest, return err } + // If we applied a distance sort, make sure that the node queried for is in + // position 0, provided the results are from the same datacenter. + if qs.Node != "" && reply.Datacenter == qs.Datacenter { + for i, node := range reply.Nodes { + if node.Node.Node == qs.Node { + reply.Nodes[0], reply.Nodes[i] = reply.Nodes[i], reply.Nodes[0] + break + } + + // Put a cap on the depth of the search. The local agent should + // never be further in than this if distance sorting was applied. + if i == 9 { + break + } + } + } + // Apply the limit if given. if args.Limit > 0 && len(reply.Nodes) > args.Limit { reply.Nodes = reply.Nodes[:args.Limit] diff --git a/consul/prepared_query_endpoint_test.go b/consul/prepared_query_endpoint_test.go index 5630f26c6..49075ed92 100644 --- a/consul/prepared_query_endpoint_test.go +++ b/consul/prepared_query_endpoint_test.go @@ -1684,6 +1684,34 @@ func TestPreparedQuery_Execute(t *testing.T) { } } + // If the exact node we are sorting near appears in the list, make sure it + // gets popped to the front of the result. + { + req := structs.PreparedQueryExecuteRequest{ + Source: 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 != "node1" { + t.Fatalf("expect node1 first, got: %q", node) + } + } + } + // Bake the magic "_agent" flag into the query. query.Query.Service.Near = "_agent" if err := msgpackrpc.CallWithCodec(codec1, "PreparedQuery.Apply", &query, &query.Query.ID); err != nil {