From afa6a1ae4960adee7b7f3fabd4d36ac4edc8a899 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Fri, 18 Apr 2014 17:37:19 -0700
Subject: [PATCH] consul: Adding support for QueryMeta on all queries

---
 consul/catalog_endpoint.go |  4 ++++
 consul/health_endpoint.go  |  4 ++++
 consul/kvs_endpoint.go     |  2 ++
 consul/rpc.go              | 11 +++++++++++
 4 files changed, 21 insertions(+)

diff --git a/consul/catalog_endpoint.go b/consul/catalog_endpoint.go
index 52fbb46de3..6a65f1534a 100644
--- a/consul/catalog_endpoint.go
+++ b/consul/catalog_endpoint.go
@@ -100,6 +100,7 @@ func (c *Catalog) ListNodes(args *structs.DCSpecificRequest, reply *structs.Inde
 	return c.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("Nodes"),
 		func() (uint64, error) {
+			c.srv.setQueryMeta(&reply.QueryMeta)
 			reply.Index, reply.Nodes = state.Nodes()
 			return reply.Index, nil
 		})
@@ -116,6 +117,7 @@ func (c *Catalog) ListServices(args *structs.DCSpecificRequest, reply *structs.I
 	return c.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("Services"),
 		func() (uint64, error) {
+			c.srv.setQueryMeta(&reply.QueryMeta)
 			reply.Index, reply.Services = state.Services()
 			return reply.Index, nil
 		})
@@ -137,6 +139,7 @@ func (c *Catalog) ServiceNodes(args *structs.ServiceSpecificRequest, reply *stru
 	err := c.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("ServiceNodes"),
 		func() (uint64, error) {
+			c.srv.setQueryMeta(&reply.QueryMeta)
 			if args.TagFilter {
 				reply.Index, reply.ServiceNodes = state.ServiceTagNodes(args.ServiceName, args.ServiceTag)
 			} else {
@@ -174,6 +177,7 @@ func (c *Catalog) NodeServices(args *structs.NodeSpecificRequest, reply *structs
 	return c.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("NodeServices"),
 		func() (uint64, error) {
+			c.srv.setQueryMeta(&reply.QueryMeta)
 			reply.Index, reply.NodeServices = state.NodeServices(args.Node)
 			return reply.Index, nil
 		})
diff --git a/consul/health_endpoint.go b/consul/health_endpoint.go
index 421d4b8862..fee6ef9952 100644
--- a/consul/health_endpoint.go
+++ b/consul/health_endpoint.go
@@ -23,6 +23,7 @@ func (h *Health) ChecksInState(args *structs.ChecksInStateRequest,
 	return h.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("ChecksInState"),
 		func() (uint64, error) {
+			h.srv.setQueryMeta(&reply.QueryMeta)
 			reply.Index, reply.HealthChecks = state.ChecksInState(args.State)
 			return reply.Index, nil
 		})
@@ -40,6 +41,7 @@ func (h *Health) NodeChecks(args *structs.NodeSpecificRequest,
 	return h.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("NodeChecks"),
 		func() (uint64, error) {
+			h.srv.setQueryMeta(&reply.QueryMeta)
 			reply.Index, reply.HealthChecks = state.NodeChecks(args.Node)
 			return reply.Index, nil
 		})
@@ -63,6 +65,7 @@ func (h *Health) ServiceChecks(args *structs.ServiceSpecificRequest,
 	return h.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("ServiceChecks"),
 		func() (uint64, error) {
+			h.srv.setQueryMeta(&reply.QueryMeta)
 			reply.Index, reply.HealthChecks = state.ServiceChecks(args.ServiceName)
 			return reply.Index, nil
 		})
@@ -84,6 +87,7 @@ func (h *Health) ServiceNodes(args *structs.ServiceSpecificRequest, reply *struc
 	err := h.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("CheckServiceNodes"),
 		func() (uint64, error) {
+			h.srv.setQueryMeta(&reply.QueryMeta)
 			if args.TagFilter {
 				reply.Index, reply.Nodes = state.CheckServiceTagNodes(args.ServiceName, args.ServiceTag)
 			} else {
diff --git a/consul/kvs_endpoint.go b/consul/kvs_endpoint.go
index 57f3e49956..fa99e3c0e6 100644
--- a/consul/kvs_endpoint.go
+++ b/consul/kvs_endpoint.go
@@ -53,6 +53,7 @@ func (k *KVS) Get(args *structs.KeyRequest, reply *structs.IndexedDirEntries) er
 	return k.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("KVSGet"),
 		func() (uint64, error) {
+			k.srv.setQueryMeta(&reply.QueryMeta)
 			index, ent, err := state.KVSGet(args.Key)
 			if err != nil {
 				return 0, err
@@ -85,6 +86,7 @@ func (k *KVS) List(args *structs.KeyRequest, reply *structs.IndexedDirEntries) e
 	return k.srv.blockingRPC(&args.BlockingQuery,
 		state.QueryTables("KVSList"),
 		func() (uint64, error) {
+			k.srv.setQueryMeta(&reply.QueryMeta)
 			index, ent, err := state.KVSList(args.Key)
 			if err != nil {
 				return 0, err
diff --git a/consul/rpc.go b/consul/rpc.go
index c950f25d52..46f01b4166 100644
--- a/consul/rpc.go
+++ b/consul/rpc.go
@@ -253,3 +253,14 @@ RUN_QUERY:
 	}
 	return err
 }
+
+// setQueryMeta is used to populate the QueryMeta data for an RPC call
+func (s *Server) setQueryMeta(m *structs.QueryMeta) {
+	if s.IsLeader() {
+		m.LastContact = 0
+		m.KnownLeader = true
+	} else {
+		m.LastContact = time.Now().Sub(s.raft.LastContact())
+		m.KnownLeader = (s.raft.Leader() != nil)
+	}
+}