@@ -3,30 +3,63 @@ package allocrunner
33import (
44 "sync"
55
6- hclog "github.com/hashicorp/go-hclog"
76 log "github.com/hashicorp/go-hclog"
87 "github.com/hashicorp/nomad/client/allocrunner/interfaces"
98 "github.com/hashicorp/nomad/client/consul"
9+ "github.com/hashicorp/nomad/client/taskenv"
10+ agentconsul "github.com/hashicorp/nomad/command/agent/consul"
1011 "github.com/hashicorp/nomad/nomad/structs"
12+ "github.com/hashicorp/nomad/plugins/drivers"
1113)
1214
1315// groupServiceHook manages task group Consul service registration and
1416// deregistration.
1517type groupServiceHook struct {
16- alloc * structs.Allocation
18+ allocID string
19+ group string
20+ restarter agentconsul.WorkloadRestarter
1721 consulClient consul.ConsulServiceAPI
1822 prerun bool
19- mu sync.Mutex
2023
2124 logger log.Logger
25+
26+ // The following fields may be updated
27+ canary bool
28+ services []* structs.Service
29+ networks structs.Networks
30+ taskEnvBuilder * taskenv.Builder
31+
32+ // Since Update() may be called concurrently with any other hook all
33+ // hook methods must be fully serialized
34+ mu sync.Mutex
35+ }
36+
37+ type groupServiceHookConfig struct {
38+ alloc * structs.Allocation
39+ consul consul.ConsulServiceAPI
40+ restarter agentconsul.WorkloadRestarter
41+ taskEnvBuilder * taskenv.Builder
42+ logger log.Logger
2243}
2344
24- func newGroupServiceHook (logger hclog. Logger , alloc * structs. Allocation , consulClient consul. ConsulServiceAPI ) * groupServiceHook {
45+ func newGroupServiceHook (cfg groupServiceHookConfig ) * groupServiceHook {
2546 h := & groupServiceHook {
26- alloc : alloc ,
27- consulClient : consulClient ,
47+ allocID : cfg .alloc .ID ,
48+ group : cfg .alloc .TaskGroup ,
49+ restarter : cfg .restarter ,
50+ consulClient : cfg .consul ,
51+ taskEnvBuilder : cfg .taskEnvBuilder ,
52+ }
53+ h .logger = cfg .logger .Named (h .Name ())
54+ h .services = cfg .alloc .Job .LookupTaskGroup (h .group ).Services
55+
56+ if cfg .alloc .AllocatedResources != nil {
57+ h .networks = cfg .alloc .AllocatedResources .Shared .Networks
58+ }
59+
60+ if cfg .alloc .DeploymentStatus != nil {
61+ h .canary = cfg .alloc .DeploymentStatus .Canary
2862 }
29- h .logger = logger .Named (h .Name ())
3063 return h
3164}
3265
@@ -41,26 +74,97 @@ func (h *groupServiceHook) Prerun() error {
4174 h .prerun = true
4275 h .mu .Unlock ()
4376 }()
44- return h .consulClient .RegisterGroup (h .alloc )
77+
78+ if len (h .services ) == 0 {
79+ return nil
80+ }
81+
82+ services := h .getWorkloadServices ()
83+ return h .consulClient .RegisterWorkload (services )
4584}
4685
4786func (h * groupServiceHook ) Update (req * interfaces.RunnerUpdateRequest ) error {
4887 h .mu .Lock ()
4988 defer h .mu .Unlock ()
50- oldAlloc := h .alloc
51- h .alloc = req .Alloc
89+ oldWorkloadServices := h .getWorkloadServices ()
90+
91+ // Store new updated values out of request
92+ canary := false
93+ if req .Alloc .DeploymentStatus != nil {
94+ canary = req .Alloc .DeploymentStatus .Canary
95+ }
96+
97+ var networks structs.Networks
98+ if req .Alloc .AllocatedResources != nil {
99+ networks = req .Alloc .AllocatedResources .Shared .Networks
100+ }
101+
102+ // Update group service hook fields
103+ h .networks = networks
104+ h .services = req .Alloc .Job .LookupTaskGroup (h .group ).Services
105+ h .canary = canary
106+ h .taskEnvBuilder .UpdateTask (req .Alloc , nil )
107+
108+ // Create new task services struct with those new values
109+ newWorkloadServices := h .getWorkloadServices ()
52110
53111 if ! h .prerun {
54112 // Update called before Prerun. Update alloc and exit to allow
55113 // Prerun to do initial registration.
56114 return nil
57115 }
58116
59- return h .consulClient .UpdateGroup ( oldAlloc , h . alloc )
117+ return h .consulClient .UpdateWorkload ( oldWorkloadServices , newWorkloadServices )
60118}
61119
62120func (h * groupServiceHook ) Postrun () error {
63121 h .mu .Lock ()
64122 defer h .mu .Unlock ()
65- return h .consulClient .RemoveGroup (h .alloc )
123+ h .deregister ()
124+ return nil
125+ }
126+
127+ func (h * groupServiceHook ) driverNet () * drivers.DriverNetwork {
128+ if len (h .networks ) == 0 {
129+ return nil
130+ }
131+
132+ //TODO(schmichael) only support one network for now
133+ net := h .networks [0 ]
134+ //TODO(schmichael) there's probably a better way than hacking driver network
135+ return & drivers.DriverNetwork {
136+ AutoAdvertise : true ,
137+ IP : net .IP ,
138+ // Copy PortLabels from group network
139+ PortMap : net .PortLabels (),
140+ }
141+ }
142+
143+ // deregister services from Consul.
144+ func (h * groupServiceHook ) deregister () {
145+ if len (h .services ) > 0 {
146+ workloadServices := h .getWorkloadServices ()
147+ h .consulClient .RemoveWorkload (workloadServices )
148+
149+ // Canary flag may be getting flipped when the alloc is being
150+ // destroyed, so remove both variations of the service
151+ workloadServices .Canary = ! workloadServices .Canary
152+ h .consulClient .RemoveWorkload (workloadServices )
153+ }
154+ }
155+
156+ func (h * groupServiceHook ) getWorkloadServices () * agentconsul.WorkloadServices {
157+ // Interpolate with the task's environment
158+ interpolatedServices := taskenv .InterpolateServices (h .taskEnvBuilder .Build (), h .services )
159+
160+ // Create task services struct with request's driver metadata
161+ return & agentconsul.WorkloadServices {
162+ AllocID : h .allocID ,
163+ Group : h .group ,
164+ Restarter : h .restarter ,
165+ Services : interpolatedServices ,
166+ DriverNetwork : h .driverNet (),
167+ Networks : h .networks ,
168+ Canary : h .canary ,
169+ }
66170}
0 commit comments