From 61eaa104b68635a7c915af96ea259517881ba5e2 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Fri, 12 Jan 2018 16:52:24 -0800
Subject: [PATCH] Add RPC for querying for Node connections

---
 nomad/status_endpoint.go      | 20 +++++++++++++++++++
 nomad/status_endpoint_test.go | 36 +++++++++++++++++++++++++++++++++++
 nomad/structs/structs.go      | 12 ++++++++++++
 3 files changed, 68 insertions(+)

diff --git a/nomad/status_endpoint.go b/nomad/status_endpoint.go
index baa700ff5..a79bbd220 100644
--- a/nomad/status_endpoint.go
+++ b/nomad/status_endpoint.go
@@ -4,7 +4,10 @@ import (
 	"fmt"
 	"strconv"
 
+	"errors"
+
 	"github.com/hashicorp/consul/agent/consul/autopilot"
+
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -126,3 +129,20 @@ func (s *Status) RaftStats(args struct{}, reply *autopilot.ServerStats) error {
 
 	return nil
 }
+
+// HasNodeConn returns whether the server has a connection to the requested
+// Node.
+func (s *Status) HasNodeConn(args *structs.NodeSpecificRequest, reply *structs.NodeConnQueryResponse) error {
+	// Validate the args
+	if args.NodeID == "" {
+		return errors.New("Must provide the NodeID")
+	}
+
+	state, ok := s.srv.getNodeConn(args.NodeID)
+	if ok {
+		reply.Connected = true
+		reply.Established = state.Established
+	}
+
+	return nil
+}
diff --git a/nomad/status_endpoint_test.go b/nomad/status_endpoint_test.go
index 9dfe5ecd7..eb968b673 100644
--- a/nomad/status_endpoint_test.go
+++ b/nomad/status_endpoint_test.go
@@ -5,10 +5,12 @@ import (
 
 	"github.com/hashicorp/net-rpc-msgpackrpc"
 	"github.com/hashicorp/nomad/acl"
+	"github.com/hashicorp/nomad/helper/uuid"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/hashicorp/nomad/testutil"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestStatusVersion(t *testing.T) {
@@ -169,3 +171,37 @@ func TestStatusMembers_ACL(t *testing.T) {
 		assert.Len(out.Members, 1)
 	}
 }
+
+func TestStatus_HasClientConn(t *testing.T) {
+	t.Parallel()
+	s1 := TestServer(t, nil)
+	defer s1.Shutdown()
+	codec := rpcClient(t, s1)
+	require := require.New(t)
+
+	arg := &structs.NodeSpecificRequest{
+		QueryOptions: structs.QueryOptions{
+			Region:     "global",
+			AllowStale: true,
+		},
+	}
+
+	// Try without setting a node id
+	var out structs.NodeConnQueryResponse
+	require.NotNil(msgpackrpc.CallWithCodec(codec, "Status.HasNodeConn", arg, &out))
+
+	// Set a bad node id
+	arg.NodeID = uuid.Generate()
+	var out2 structs.NodeConnQueryResponse
+	require.Nil(msgpackrpc.CallWithCodec(codec, "Status.HasNodeConn", arg, &out2))
+	require.False(out2.Connected)
+
+	// Create a connection on that node
+	s1.addNodeConn(&RPCContext{
+		NodeID: arg.NodeID,
+	})
+	var out3 structs.NodeConnQueryResponse
+	require.Nil(msgpackrpc.CallWithCodec(codec, "Status.HasNodeConn", arg, &out3))
+	require.True(out3.Connected)
+	require.NotZero(out3.Established)
+}
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index bcc408074..75b8af4f9 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -1037,6 +1037,18 @@ type DeploymentUpdateResponse struct {
 	WriteMeta
 }
 
+// NodeConnQueryResponse is used to respond to a query of whether a server has
+// a connection to a specific Node
+type NodeConnQueryResponse struct {
+	// Connected indicates whether a connection to the Client exists
+	Connected bool
+
+	// Established marks the time at which the connection was established
+	Established time.Time
+
+	QueryMeta
+}
+
 const (
 	NodeStatusInit  = "initializing"
 	NodeStatusReady = "ready"