From 3797e6544c400ba37d95770a29cd221cca94abac Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Mon, 20 Jun 2016 14:07:08 -0700
Subject: [PATCH] consul: support PreferLocal in PQ's

---
 consul/prepared_query_endpoint.go | 11 +++++++++++
 consul/structs/prepared_query.go  |  9 +++++++++
 2 files changed, 20 insertions(+)

diff --git a/consul/prepared_query_endpoint.go b/consul/prepared_query_endpoint.go
index d30b6c10dd..970b459d56 100644
--- a/consul/prepared_query_endpoint.go
+++ b/consul/prepared_query_endpoint.go
@@ -372,6 +372,17 @@ func (p *PreparedQuery) Execute(args *structs.PreparedQueryExecuteRequest,
 		return err
 	}
 
+	// Prefer the local service if it exists and is in the set.
+	if query.Service.PreferLocal {
+		for i, node := range reply.Nodes {
+			if node.Node.Node == args.Origin {
+				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]
diff --git a/consul/structs/prepared_query.go b/consul/structs/prepared_query.go
index b1b20c9ed3..85439b6e0c 100644
--- a/consul/structs/prepared_query.go
+++ b/consul/structs/prepared_query.go
@@ -34,6 +34,11 @@ type ServiceQuery struct {
 	// discarded)
 	OnlyPassing bool
 
+	// If PreferLocal is true, the local agent will be checked for a local
+	// copy of the service before continuing to remote machines. This is
+	// useful to prefer colocated services but fall back when unavailable.
+	PreferLocal bool
+
 	// 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.
@@ -163,6 +168,10 @@ func (q *PreparedQuerySpecificRequest) RequestDatacenter() string {
 
 // PreparedQueryExecuteRequest is used to execute a prepared query.
 type PreparedQueryExecuteRequest struct {
+	// Origin is used to carry around a reference to the node which
+	// is executing the query on behalf of the client.
+	Origin string
+
 	// Datacenter is the target this request is intended for.
 	Datacenter string