/
wrapper.go
139 lines (118 loc) · 4.92 KB
/
wrapper.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package wrapper
import (
"fmt"
"github.com/hashicorp/go-hclog"
"github.com/hashicorp/nomad/client/serviceregistration"
"github.com/hashicorp/nomad/nomad/structs"
)
// HandlerWrapper is used to wrap service registration implementations of the
// Handler interface. We do not use a map or similar to store the handlers, so
// we can avoid having to use a lock. This may need to be updated if we ever
// support additional registration providers.
type HandlerWrapper struct {
log hclog.Logger
// consulServiceProvider is the handler for services where Consul is the
// provider. This provider is always created and available.
consulServiceProvider serviceregistration.Handler
// nomadServiceProvider is the handler for services where Nomad is the
// provider.
nomadServiceProvider serviceregistration.Handler
}
// NewHandlerWrapper configures and returns a HandlerWrapper for use within
// client hooks that need to interact with service and check registrations. It
// mimics the serviceregistration.Handler interface, but returns the
// implementation to allow future flexibility and is initially only intended
// for use with the alloc and task runner service hooks.
func NewHandlerWrapper(
log hclog.Logger, consulProvider, nomadProvider serviceregistration.Handler) *HandlerWrapper {
return &HandlerWrapper{
log: log,
nomadServiceProvider: nomadProvider,
consulServiceProvider: consulProvider,
}
}
// RegisterWorkload wraps the serviceregistration.Handler RegisterWorkload
// function. It determines which backend provider to call and passes the
// workload unless the provider is unknown, in which case an error will be
// returned.
func (h *HandlerWrapper) RegisterWorkload(workload *serviceregistration.WorkloadServices) error {
// Don't rely on callers to check there are no services to register.
if len(workload.Services) == 0 {
return nil
}
provider := workload.RegistrationProvider()
switch provider {
case structs.ServiceProviderNomad:
return h.nomadServiceProvider.RegisterWorkload(workload)
case structs.ServiceProviderConsul:
return h.consulServiceProvider.RegisterWorkload(workload)
default:
return fmt.Errorf("unknown service registration provider: %q", provider)
}
}
// RemoveWorkload wraps the serviceregistration.Handler RemoveWorkload
// function. It determines which backend provider to call and passes the
// workload unless the provider is unknown.
func (h *HandlerWrapper) RemoveWorkload(services *serviceregistration.WorkloadServices) {
var provider string
// It is possible the services field is empty depending on the exact
// situation which resulted in the call.
if len(services.Services) > 0 {
provider = services.RegistrationProvider()
}
// Call the correct provider, if we have managed to identify it. An empty
// string means you didn't find a provider, therefore default to consul.
//
// In certain situations this function is called with zero services,
// therefore meaning we make an assumption on the provider. When this
// happens, we need to ensure the allocation is removed from the Consul
// implementation. This tracking (allocRegistrations) is used by the
// allochealth tracker and so is critical to be removed. The test
// allocrunner.TestAllocRunner_Restore_RunningTerminal covers the case
// described here.
switch provider {
case structs.ServiceProviderNomad:
h.nomadServiceProvider.RemoveWorkload(services)
case structs.ServiceProviderConsul, "":
h.consulServiceProvider.RemoveWorkload(services)
default:
h.log.Error("unknown service registration provider", "provider", provider)
}
}
// UpdateWorkload identifies which provider to call for the new and old
// workloads provided. In the event both use the same provider, the
// UpdateWorkload function will be called, otherwise the register and remove
// functions will be called.
func (h *HandlerWrapper) UpdateWorkload(old, new *serviceregistration.WorkloadServices) error {
// Hot path to exit if there is nothing to do.
if len(old.Services) == 0 && len(new.Services) == 0 {
return nil
}
newProvider := new.RegistrationProvider()
oldProvider := old.RegistrationProvider()
// If the new and old services use the same provider, call the
// UpdateWorkload and leave it at that.
if newProvider == oldProvider {
switch newProvider {
case structs.ServiceProviderNomad:
return h.nomadServiceProvider.UpdateWorkload(old, new)
case structs.ServiceProviderConsul:
return h.consulServiceProvider.UpdateWorkload(old, new)
default:
return fmt.Errorf("unknown service registration provider for update: %q", newProvider)
}
}
// If we have new services, call the relevant provider. Registering can
// return an error. Do this before RemoveWorkload, so we can halt the
// process if needed, otherwise we may leave the task/group
// registration-less.
if len(new.Services) > 0 {
if err := h.RegisterWorkload(new); err != nil {
return err
}
}
if len(old.Services) > 0 {
h.RemoveWorkload(old)
}
return nil
}