client: hookup service wrapper for use in clients

client/allocrunner/alloc_runner.go

@@ -6,8 +6,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/hashicorp/nomad/client/lib/cgutil"
-
 	log "github.com/hashicorp/go-hclog"
 	multierror "github.com/hashicorp/go-multierror"
 	"github.com/hashicorp/nomad/client/allocdir"
@@ -20,9 +18,11 @@ import (
 	"github.com/hashicorp/nomad/client/devicemanager"
 	"github.com/hashicorp/nomad/client/dynamicplugins"
 	cinterfaces "github.com/hashicorp/nomad/client/interfaces"
+	"github.com/hashicorp/nomad/client/lib/cgutil"
 	"github.com/hashicorp/nomad/client/pluginmanager/csimanager"
 	"github.com/hashicorp/nomad/client/pluginmanager/drivermanager"
 	"github.com/hashicorp/nomad/client/serviceregistration"
+	"github.com/hashicorp/nomad/client/serviceregistration/wrapper"
 	cstate "github.com/hashicorp/nomad/client/state"
 	cstructs "github.com/hashicorp/nomad/client/structs"
 	"github.com/hashicorp/nomad/client/vaultclient"
@@ -178,6 +178,10 @@ type allocRunner struct {
 	// rpcClient is the RPC Client that should be used by the allocrunner and its
 	// hooks to communicate with Nomad Servers.
 	rpcClient RPCer
+
+	// serviceRegWrapper is the handler wrapper that is used by service hooks
+	// to perform service and check registration and deregistration.
+	serviceRegWrapper *wrapper.HandlerWrapper
 }
 
 // RPCer is the interface needed by hooks to make RPC calls.
@@ -221,6 +225,7 @@ func NewAllocRunner(config *Config) (*allocRunner, error) {
 		driverManager:            config.DriverManager,
 		serversContactedCh:       config.ServersContactedCh,
 		rpcClient:                config.RPCClient,
+		serviceRegWrapper:        config.ServiceRegWrapper,
 	}
 
 	// Create the logger based on the allocation ID
@@ -274,6 +279,7 @@ func (ar *allocRunner) initTaskRunners(tasks []*structs.Task) error {
 			ServersContactedCh:   ar.serversContactedCh,
 			StartConditionMetCtx: ar.taskHookCoordinator.startConditionForTask(task),
 			ShutdownDelayCtx:     ar.shutdownDelayCtx,
+			ServiceRegWrapper:    ar.serviceRegWrapper,
 		}
 
 		if ar.cpusetManager != nil {

client/allocrunner/alloc_runner_hooks.go

@@ -153,8 +153,8 @@ func (ar *allocRunner) initRunnerHooks(config *clientconfig.Config) error {
 		newNetworkHook(hookLogger, ns, alloc, nm, nc, ar, builtTaskEnv),
 		newGroupServiceHook(groupServiceHookConfig{
 			alloc:               alloc,
-			consul:              ar.consulClient,
-			consulNamespace:     alloc.ConsulNamespace(),
+			namespace:           alloc.ServiceProviderNamespace(),
+			serviceRegWrapper:   ar.serviceRegWrapper,
 			restarter:           ar,
 			taskEnvBuilder:      envBuilder,
 			networkStatusGetter: ar,

client/allocrunner/alloc_runner_test.go

@@ -489,7 +489,8 @@ func TestAllocRunner_TaskGroup_ShutdownDelay(t *testing.T) {
 	tg := alloc.Job.TaskGroups[0]
 	tg.Services = []*structs.Service{
 		{
-			Name: "shutdown_service",
+			Name:     "shutdown_service",
+			Provider: structs.ServiceProviderConsul,
 		},
 	}
 
@@ -1314,6 +1315,7 @@ func TestAllocRunner_TaskFailed_KillTG(t *testing.T) {
 		{
 			Name:      "fakservice",
 			PortLabel: "http",
+			Provider:  structs.ServiceProviderConsul,
 			Checks: []*structs.ServiceCheck{
 				{
 					Name:     "fakecheck",

client/allocrunner/alloc_runner_unix_test.go

@@ -38,6 +38,7 @@ func TestAllocRunner_Restore_RunningTerminal(t *testing.T) {
 		{
 			Name:      "foo",
 			PortLabel: "8888",
+			Provider:  structs.ServiceProviderConsul,
 		},
 	}
 	task := alloc.Job.TaskGroups[0].Tasks[0]

client/allocrunner/config.go

@@ -12,6 +12,7 @@ import (
 	"github.com/hashicorp/nomad/client/pluginmanager/csimanager"
 	"github.com/hashicorp/nomad/client/pluginmanager/drivermanager"
 	"github.com/hashicorp/nomad/client/serviceregistration"
+	"github.com/hashicorp/nomad/client/serviceregistration/wrapper"
 	cstate "github.com/hashicorp/nomad/client/state"
 	"github.com/hashicorp/nomad/client/vaultclient"
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -81,4 +82,8 @@ type Config struct {
 	// RPCClient is the RPC Client that should be used by the allocrunner and its
 	// hooks to communicate with Nomad Servers.
 	RPCClient RPCer
+
+	// ServiceRegWrapper is the handler wrapper that is used by service hooks
+	// to perform service and check registration and deregistration.
+	ServiceRegWrapper *wrapper.HandlerWrapper
 }

client/allocrunner/groupservice_hook.go

@@ -8,6 +8,7 @@ import (
 	log "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
 	"github.com/hashicorp/nomad/client/serviceregistration"
+	"github.com/hashicorp/nomad/client/serviceregistration/wrapper"
 	"github.com/hashicorp/nomad/client/taskenv"
 	agentconsul "github.com/hashicorp/nomad/command/agent/consul"
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -25,15 +26,22 @@ type networkStatusGetter interface {
 // deregistration.
 type groupServiceHook struct {
 	allocID             string
+	jobID               string
 	group               string
 	restarter           agentconsul.WorkloadRestarter
-	consulClient        serviceregistration.Handler
-	consulNamespace     string
 	prerun              bool
 	deregistered        bool
 	networkStatusGetter networkStatusGetter
 	shutdownDelayCtx    context.Context
 
+	// namespace is the Nomad or Consul namespace in which service
+	// registrations will be made.
+	namespace string
+
+	// serviceRegWrapper is the handler wrapper that is used to perform service
+	// and check registration and deregistration.
+	serviceRegWrapper *wrapper.HandlerWrapper
+
 	logger log.Logger
 
 	// The following fields may be updated
@@ -51,13 +59,19 @@ type groupServiceHook struct {
 
 type groupServiceHookConfig struct {
 	alloc               *structs.Allocation
-	consul              serviceregistration.Handler
-	consulNamespace     string
 	restarter           agentconsul.WorkloadRestarter
 	taskEnvBuilder      *taskenv.Builder
 	networkStatusGetter networkStatusGetter
 	shutdownDelayCtx    context.Context
 	logger              log.Logger
+
+	// namespace is the Nomad or Consul namespace in which service
+	// registrations will be made.
+	namespace string
+
+	// serviceRegWrapper is the handler wrapper that is used to perform service
+	// and check registration and deregistration.
+	serviceRegWrapper *wrapper.HandlerWrapper
 }
 
 func newGroupServiceHook(cfg groupServiceHookConfig) *groupServiceHook {
@@ -70,15 +84,16 @@ func newGroupServiceHook(cfg groupServiceHookConfig) *groupServiceHook {
 
 	h := &groupServiceHook{
 		allocID:             cfg.alloc.ID,
+		jobID:               cfg.alloc.JobID,
 		group:               cfg.alloc.TaskGroup,
 		restarter:           cfg.restarter,
-		consulClient:        cfg.consul,
-		consulNamespace:     cfg.consulNamespace,
+		namespace:           cfg.namespace,
 		taskEnvBuilder:      cfg.taskEnvBuilder,
 		delay:               shutdownDelay,
 		networkStatusGetter: cfg.networkStatusGetter,
 		logger:              cfg.logger.Named(groupServiceHookName),
 		services:            cfg.alloc.Job.LookupTaskGroup(cfg.alloc.TaskGroup).Services,
+		serviceRegWrapper:   cfg.serviceRegWrapper,
 		shutdownDelayCtx:    cfg.shutdownDelayCtx,
 	}
 
@@ -114,7 +129,7 @@ func (h *groupServiceHook) prerunLocked() error {
 	}
 
 	services := h.getWorkloadServices()
-	return h.consulClient.RegisterWorkload(services)
+	return h.serviceRegWrapper.RegisterWorkload(services)
 }
 
 func (h *groupServiceHook) Update(req *interfaces.RunnerUpdateRequest) error {
@@ -157,7 +172,7 @@ func (h *groupServiceHook) Update(req *interfaces.RunnerUpdateRequest) error {
 		return nil
 	}
 
-	return h.consulClient.UpdateWorkload(oldWorkloadServices, newWorkloadServices)
+	return h.serviceRegWrapper.UpdateWorkload(oldWorkloadServices, newWorkloadServices)
 }
 
 func (h *groupServiceHook) PreTaskRestart() error {
@@ -213,7 +228,7 @@ func (h *groupServiceHook) Postrun() error {
 func (h *groupServiceHook) deregister() {
 	if len(h.services) > 0 {
 		workloadServices := h.getWorkloadServices()
-		h.consulClient.RemoveWorkload(workloadServices)
+		h.serviceRegWrapper.RemoveWorkload(workloadServices)
 	}
 }
 
@@ -229,8 +244,9 @@ func (h *groupServiceHook) getWorkloadServices() *serviceregistration.WorkloadSe
 	// Create task services struct with request's driver metadata
 	return &serviceregistration.WorkloadServices{
 		AllocID:       h.allocID,
+		JobID:         h.jobID,
 		Group:         h.group,
-		Namespace:     h.consulNamespace,
+		Namespace:     h.namespace,
 		Restarter:     h.restarter,
 		Services:      interpolatedServices,
 		Networks:      h.networks,