Refactor determining the handler for a node id call

command/agent/fs_endpoint.go

@@ -45,32 +45,6 @@ func (s *HTTPServer) FsRequest(resp http.ResponseWriter, req *http.Request) (int
 	}
 }
 
-// rpcHandlerForAlloc is a helper that given an allocation ID returns whether to
-// use the local clients RPC, the local clients remote RPC or the server on the
-// agent.
-func (s *HTTPServer) rpcHandlerForAlloc(allocID string) (localClient, remoteClient, server bool) {
-	c := s.agent.Client()
-	srv := s.agent.Server()
-
-	// See if the local client can handle the request.
-	localAlloc := false
-	if c != nil {
-		_, err := c.GetClientAlloc(allocID)
-		if err == nil {
-			localAlloc = true
-		}
-	}
-
-	// Only use the client RPC to server if we don't have a server and the local
-	// client can't handle the call.
-	useClientRPC := c != nil && !localAlloc && srv == nil
-
-	// Use the server as a last case.
-	useServerRPC := !localAlloc && !useClientRPC && srv != nil
-
-	return localAlloc, useClientRPC, useServerRPC
-}
-
 func (s *HTTPServer) DirectoryListRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
 	var allocID, path string
 

command/agent/fs_endpoint_test.go

@@ -11,7 +11,6 @@ import (
 	"testing"
 
 	cstructs "github.com/hashicorp/nomad/client/structs"
-	"github.com/hashicorp/nomad/helper/uuid"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/hashicorp/nomad/testutil"
@@ -108,49 +107,6 @@ func mockFSAlloc(nodeID string, config map[string]interface{}) *structs.Allocati
 	return a
 }
 
-func TestHTTP_FS_rpcHandlerForAlloc(t *testing.T) {
-	t.Parallel()
-	require := require.New(t)
-	agent := NewTestAgent(t, t.Name(), nil)
-
-	a := mockFSAlloc(agent.client.NodeID(), nil)
-	addAllocToClient(agent, a, terminalClientAlloc)
-
-	// Case 1: Client has allocation
-	// Outcome: Use local client
-	lc, rc, s := agent.Server.rpcHandlerForAlloc(a.ID)
-	require.True(lc)
-	require.False(rc)
-	require.False(s)
-
-	// Case 2: Client doesn't have allocation and there is a server
-	// Outcome: Use server
-	lc, rc, s = agent.Server.rpcHandlerForAlloc(uuid.Generate())
-	require.False(lc)
-	require.False(rc)
-	require.True(s)
-
-	// Case 3: Client doesn't have allocation and there is no server
-	// Outcome: Use client RPC to server
-	srv := agent.server
-	agent.server = nil
-	lc, rc, s = agent.Server.rpcHandlerForAlloc(uuid.Generate())
-	require.False(lc)
-	require.True(rc)
-	require.False(s)
-	agent.server = srv
-
-	// Case 4: No client
-	// Outcome: Use server
-	client := agent.client
-	agent.client = nil
-	lc, rc, s = agent.Server.rpcHandlerForAlloc(uuid.Generate())
-	require.False(lc)
-	require.False(rc)
-	require.True(s)
-	agent.client = client
-}
-
 func TestHTTP_FS_List_MissingParams(t *testing.T) {
 	t.Parallel()
 	require := require.New(t)

command/agent/helpers.go

@@ -0,0 +1,50 @@
+package agent
+
+// rpcHandlerForAlloc is a helper that given an allocation ID returns whether to
+// use the local clients RPC, the local clients remote RPC or the server on the
+// agent.
+func (s *HTTPServer) rpcHandlerForAlloc(allocID string) (localClient, remoteClient, server bool) {
+	c := s.agent.Client()
+	srv := s.agent.Server()
+
+	// See if the local client can handle the request.
+	localAlloc := false
+	if c != nil {
+		_, err := c.GetClientAlloc(allocID)
+		if err == nil {
+			localAlloc = true
+		}
+	}
+
+	// Only use the client RPC to server if we don't have a server and the local
+	// client can't handle the call.
+	useClientRPC := c != nil && !localAlloc && srv == nil
+
+	// Use the server as a last case.
+	useServerRPC := !localAlloc && !useClientRPC && srv != nil
+
+	return localAlloc, useClientRPC, useServerRPC
+}
+
+// rpcHandlerForNode is a helper that given a node ID returns whether to
+// use the local clients RPC, the local clients remote RPC or the server on the
+// agent. If there is a local node and no node id is given, it is assumed the
+// local node is being targed.
+func (s *HTTPServer) rpcHandlerForNode(nodeID string) (localClient, remoteClient, server bool) {
+	c := s.agent.Client()
+	srv := s.agent.Server()
+
+	// See if the local client can handle the request.
+	localClient = c != nil && // Must have a client
+		(nodeID == "" || // If no node ID is given
+			nodeID == c.NodeID()) // Requested node is the local node.
+
+	// Only use the client RPC to server if we don't have a server and the local
+	// client can't handle the call.
+	useClientRPC := c != nil && !localClient && srv == nil
+
+	// Use the server as a last case.
+	useServerRPC := !localClient && !useClientRPC && srv != nil && nodeID != ""
+
+	return localClient, useClientRPC, useServerRPC
+}

command/agent/helpers_test.go

@@ -0,0 +1,92 @@
+package agent
+
+import (
+	"testing"
+
+	"github.com/hashicorp/nomad/helper/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHTTP_rpcHandlerForAlloc(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+	agent := NewTestAgent(t, t.Name(), nil)
+
+	a := mockFSAlloc(agent.client.NodeID(), nil)
+	addAllocToClient(agent, a, terminalClientAlloc)
+
+	// Case 1: Client has allocation
+	// Outcome: Use local client
+	lc, rc, s := agent.Server.rpcHandlerForAlloc(a.ID)
+	require.True(lc)
+	require.False(rc)
+	require.False(s)
+
+	// Case 2: Client doesn't have allocation and there is a server
+	// Outcome: Use server
+	lc, rc, s = agent.Server.rpcHandlerForAlloc(uuid.Generate())
+	require.False(lc)
+	require.False(rc)
+	require.True(s)
+
+	// Case 3: Client doesn't have allocation and there is no server
+	// Outcome: Use client RPC to server
+	srv := agent.server
+	agent.server = nil
+	lc, rc, s = agent.Server.rpcHandlerForAlloc(uuid.Generate())
+	require.False(lc)
+	require.True(rc)
+	require.False(s)
+	agent.server = srv
+
+	// Case 4: No client
+	// Outcome: Use server
+	client := agent.client
+	agent.client = nil
+	lc, rc, s = agent.Server.rpcHandlerForAlloc(uuid.Generate())
+	require.False(lc)
+	require.False(rc)
+	require.True(s)
+	agent.client = client
+}
+
+func TestHTTP_rpcHandlerForNode(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+	agent := NewTestAgent(t, t.Name(), nil)
+	cID := agent.client.NodeID()
+
+	// Case 1: Node running, no node ID given
+	// Outcome: Use local node
+	lc, rc, s := agent.Server.rpcHandlerForNode("")
+	require.True(lc)
+	require.False(rc)
+	require.False(s)
+
+	// Case 2: Node running, it's ID given
+	// Outcome: Use local node
+	lc, rc, s = agent.Server.rpcHandlerForNode(cID)
+	require.True(lc)
+	require.False(rc)
+	require.False(s)
+
+	// Case 3: Local node but wrong ID and there is no server
+	// Outcome: Use client RPC to server
+	srv := agent.server
+	agent.server = nil
+	lc, rc, s = agent.Server.rpcHandlerForNode(uuid.Generate())
+	require.False(lc)
+	require.True(rc)
+	require.False(s)
+	agent.server = srv
+
+	// Case 4: No client
+	// Outcome: Use server
+	client := agent.client
+	agent.client = nil
+	lc, rc, s = agent.Server.rpcHandlerForNode(uuid.Generate())
+	require.False(lc)
+	require.False(rc)
+	require.True(s)
+	agent.client = client
+}

command/agent/stats_endpoint.go

@@ -12,27 +12,15 @@ func (s *HTTPServer) ClientStatsRequest(resp http.ResponseWriter, req *http.Requ
 	// Get the requested Node ID
 	requestedNode := req.URL.Query().Get("node_id")
 
-	localClient := s.agent.Client()
-	localServer := s.agent.Server()
-
-	// See if the local client can handle the request.
-	useLocalClient := localClient != nil && // Must have a client
-		(requestedNode == "" || // If no node ID is given
-			requestedNode == localClient.NodeID()) // Requested node is the local node.
-
-	// Only use the client RPC to server if we don't have a server and the local
-	// client can't handle the call.
-	useClientRPC := localClient != nil && !useLocalClient && localServer == nil
-
-	// Use the server as a last case.
-	useServerRPC := localServer != nil && requestedNode != ""
-
 	// Build the request and parse the ACL token
 	args := cstructs.ClientStatsRequest{
 		NodeID: requestedNode,
 	}
 	s.parse(resp, req, &args.QueryOptions.Region, &args.QueryOptions)
 
+	// Determine the handler to use
+	useLocalClient, useClientRPC, useServerRPC := s.rpcHandlerForNode(requestedNode)
+
 	// Make the RPC
 	var reply cstructs.ClientStatsResponse
 	var rpcErr error