From 29333e86033a8ca6597360b981d36f62a5429366 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Fri, 1 Jul 2016 14:28:58 -0700
Subject: [PATCH] consul: sort source node first if at position <= 10 in PQ's

---
 consul/prepared_query_endpoint.go      | 17 ++++++++++++++++
 consul/prepared_query_endpoint_test.go | 28 ++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/consul/prepared_query_endpoint.go b/consul/prepared_query_endpoint.go
index 5f7d788943..47b6fe1a77 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 5630f26c61..49075ed928 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 {