20182: API server sets "supervisor" flag now
[arvados.git] / lib / dispatchcloud / dispatcher.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package dispatchcloud
6
7 import (
8         "context"
9         "crypto/md5"
10         "encoding/json"
11         "fmt"
12         "net/http"
13         "strings"
14         "sync"
15         "time"
16
17         "git.arvados.org/arvados.git/lib/cloud"
18         "git.arvados.org/arvados.git/lib/controller/dblock"
19         "git.arvados.org/arvados.git/lib/ctrlctx"
20         "git.arvados.org/arvados.git/lib/dispatchcloud/container"
21         "git.arvados.org/arvados.git/lib/dispatchcloud/scheduler"
22         "git.arvados.org/arvados.git/lib/dispatchcloud/sshexecutor"
23         "git.arvados.org/arvados.git/lib/dispatchcloud/worker"
24         "git.arvados.org/arvados.git/sdk/go/arvados"
25         "git.arvados.org/arvados.git/sdk/go/auth"
26         "git.arvados.org/arvados.git/sdk/go/ctxlog"
27         "git.arvados.org/arvados.git/sdk/go/health"
28         "git.arvados.org/arvados.git/sdk/go/httpserver"
29         "github.com/julienschmidt/httprouter"
30         "github.com/prometheus/client_golang/prometheus"
31         "github.com/prometheus/client_golang/prometheus/promhttp"
32         "github.com/sirupsen/logrus"
33         "golang.org/x/crypto/ssh"
34 )
35
36 const (
37         defaultPollInterval     = time.Second
38         defaultStaleLockTimeout = time.Minute
39 )
40
41 type pool interface {
42         scheduler.WorkerPool
43         CheckHealth() error
44         Instances() []worker.InstanceView
45         SetIdleBehavior(cloud.InstanceID, worker.IdleBehavior) error
46         KillInstance(id cloud.InstanceID, reason string) error
47         Stop()
48 }
49
50 type dispatcher struct {
51         Cluster       *arvados.Cluster
52         Context       context.Context
53         ArvClient     *arvados.Client
54         AuthToken     string
55         Registry      *prometheus.Registry
56         InstanceSetID cloud.InstanceSetID
57
58         dbConnector ctrlctx.DBConnector
59         logger      logrus.FieldLogger
60         instanceSet cloud.InstanceSet
61         pool        pool
62         queue       scheduler.ContainerQueue
63         httpHandler http.Handler
64         sshKey      ssh.Signer
65
66         setupOnce sync.Once
67         stop      chan struct{}
68         stopped   chan struct{}
69 }
70
71 // Start starts the dispatcher. Start can be called multiple times
72 // with no ill effect.
73 func (disp *dispatcher) Start() {
74         disp.setupOnce.Do(disp.setup)
75 }
76
77 // ServeHTTP implements service.Handler.
78 func (disp *dispatcher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
79         disp.Start()
80         disp.httpHandler.ServeHTTP(w, r)
81 }
82
83 // CheckHealth implements service.Handler.
84 func (disp *dispatcher) CheckHealth() error {
85         disp.Start()
86         return disp.pool.CheckHealth()
87 }
88
89 // Done implements service.Handler.
90 func (disp *dispatcher) Done() <-chan struct{} {
91         return disp.stopped
92 }
93
94 // Stop dispatching containers and release resources. Typically used
95 // in tests.
96 func (disp *dispatcher) Close() {
97         disp.Start()
98         select {
99         case disp.stop <- struct{}{}:
100         default:
101         }
102         <-disp.stopped
103 }
104
105 // Make a worker.Executor for the given instance.
106 func (disp *dispatcher) newExecutor(inst cloud.Instance) worker.Executor {
107         exr := sshexecutor.New(inst)
108         exr.SetTargetPort(disp.Cluster.Containers.CloudVMs.SSHPort)
109         exr.SetSigners(disp.sshKey)
110         return exr
111 }
112
113 func (disp *dispatcher) typeChooser(ctr *arvados.Container) (arvados.InstanceType, error) {
114         return ChooseInstanceType(disp.Cluster, ctr)
115 }
116
117 func (disp *dispatcher) setup() {
118         disp.initialize()
119         go disp.run()
120 }
121
122 func (disp *dispatcher) initialize() {
123         disp.logger = ctxlog.FromContext(disp.Context)
124         disp.dbConnector = ctrlctx.DBConnector{PostgreSQL: disp.Cluster.PostgreSQL}
125
126         disp.ArvClient.AuthToken = disp.AuthToken
127
128         if disp.InstanceSetID == "" {
129                 if strings.HasPrefix(disp.AuthToken, "v2/") {
130                         disp.InstanceSetID = cloud.InstanceSetID(strings.Split(disp.AuthToken, "/")[1])
131                 } else {
132                         // Use some other string unique to this token
133                         // that doesn't reveal the token itself.
134                         disp.InstanceSetID = cloud.InstanceSetID(fmt.Sprintf("%x", md5.Sum([]byte(disp.AuthToken))))
135                 }
136         }
137         disp.stop = make(chan struct{}, 1)
138         disp.stopped = make(chan struct{})
139
140         if key, err := ssh.ParsePrivateKey([]byte(disp.Cluster.Containers.DispatchPrivateKey)); err != nil {
141                 disp.logger.Fatalf("error parsing configured Containers.DispatchPrivateKey: %s", err)
142         } else {
143                 disp.sshKey = key
144         }
145
146         instanceSet, err := newInstanceSet(disp.Cluster, disp.InstanceSetID, disp.logger, disp.Registry)
147         if err != nil {
148                 disp.logger.Fatalf("error initializing driver: %s", err)
149         }
150         dblock.Dispatch.Lock(disp.Context, disp.dbConnector.GetDB)
151         disp.instanceSet = instanceSet
152         disp.pool = worker.NewPool(disp.logger, disp.ArvClient, disp.Registry, disp.InstanceSetID, disp.instanceSet, disp.newExecutor, disp.sshKey.PublicKey(), disp.Cluster)
153         disp.queue = container.NewQueue(disp.logger, disp.Registry, disp.typeChooser, disp.ArvClient)
154
155         if disp.Cluster.ManagementToken == "" {
156                 disp.httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
157                         http.Error(w, "Management API authentication is not configured", http.StatusForbidden)
158                 })
159         } else {
160                 mux := httprouter.New()
161                 mux.HandlerFunc("GET", "/arvados/v1/dispatch/containers", disp.apiContainers)
162                 mux.HandlerFunc("POST", "/arvados/v1/dispatch/containers/kill", disp.apiContainerKill)
163                 mux.HandlerFunc("GET", "/arvados/v1/dispatch/instances", disp.apiInstances)
164                 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/hold", disp.apiInstanceHold)
165                 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/drain", disp.apiInstanceDrain)
166                 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/run", disp.apiInstanceRun)
167                 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/kill", disp.apiInstanceKill)
168                 metricsH := promhttp.HandlerFor(disp.Registry, promhttp.HandlerOpts{
169                         ErrorLog: disp.logger,
170                 })
171                 mux.Handler("GET", "/metrics", metricsH)
172                 mux.Handler("GET", "/metrics.json", metricsH)
173                 mux.Handler("GET", "/_health/:check", &health.Handler{
174                         Token:  disp.Cluster.ManagementToken,
175                         Prefix: "/_health/",
176                         Routes: health.Routes{"ping": disp.CheckHealth},
177                 })
178                 disp.httpHandler = auth.RequireLiteralToken(disp.Cluster.ManagementToken, mux)
179         }
180 }
181
182 func (disp *dispatcher) run() {
183         defer dblock.Dispatch.Unlock()
184         defer close(disp.stopped)
185         defer disp.instanceSet.Stop()
186         defer disp.pool.Stop()
187
188         staleLockTimeout := time.Duration(disp.Cluster.Containers.StaleLockTimeout)
189         if staleLockTimeout == 0 {
190                 staleLockTimeout = defaultStaleLockTimeout
191         }
192         pollInterval := time.Duration(disp.Cluster.Containers.CloudVMs.PollInterval)
193         if pollInterval <= 0 {
194                 pollInterval = defaultPollInterval
195         }
196         maxSupervisors := int(float64(disp.Cluster.Containers.CloudVMs.MaxInstances) * disp.Cluster.Containers.CloudVMs.SupervisorFraction)
197         if maxSupervisors == 0 && disp.Cluster.Containers.CloudVMs.SupervisorFraction > 0 {
198                 maxSupervisors = 1
199         }
200         sched := scheduler.New(disp.Context, disp.ArvClient, disp.queue, disp.pool, disp.Registry, staleLockTimeout, pollInterval, maxSupervisors)
201         sched.Start()
202         defer sched.Stop()
203
204         <-disp.stop
205 }
206
207 // Management API: all active and queued containers.
208 func (disp *dispatcher) apiContainers(w http.ResponseWriter, r *http.Request) {
209         var resp struct {
210                 Items []container.QueueEnt `json:"items"`
211         }
212         qEntries, _ := disp.queue.Entries()
213         for _, ent := range qEntries {
214                 resp.Items = append(resp.Items, ent)
215         }
216         json.NewEncoder(w).Encode(resp)
217 }
218
219 // Management API: all active instances (cloud VMs).
220 func (disp *dispatcher) apiInstances(w http.ResponseWriter, r *http.Request) {
221         var resp struct {
222                 Items []worker.InstanceView `json:"items"`
223         }
224         resp.Items = disp.pool.Instances()
225         json.NewEncoder(w).Encode(resp)
226 }
227
228 // Management API: set idle behavior to "hold" for specified instance.
229 func (disp *dispatcher) apiInstanceHold(w http.ResponseWriter, r *http.Request) {
230         disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorHold)
231 }
232
233 // Management API: set idle behavior to "drain" for specified instance.
234 func (disp *dispatcher) apiInstanceDrain(w http.ResponseWriter, r *http.Request) {
235         disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorDrain)
236 }
237
238 // Management API: set idle behavior to "run" for specified instance.
239 func (disp *dispatcher) apiInstanceRun(w http.ResponseWriter, r *http.Request) {
240         disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorRun)
241 }
242
243 // Management API: shutdown/destroy specified instance now.
244 func (disp *dispatcher) apiInstanceKill(w http.ResponseWriter, r *http.Request) {
245         id := cloud.InstanceID(r.FormValue("instance_id"))
246         if id == "" {
247                 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
248                 return
249         }
250         err := disp.pool.KillInstance(id, "via management API: "+r.FormValue("reason"))
251         if err != nil {
252                 httpserver.Error(w, err.Error(), http.StatusNotFound)
253                 return
254         }
255 }
256
257 // Management API: send SIGTERM to specified container's crunch-run
258 // process now.
259 func (disp *dispatcher) apiContainerKill(w http.ResponseWriter, r *http.Request) {
260         uuid := r.FormValue("container_uuid")
261         if uuid == "" {
262                 httpserver.Error(w, "container_uuid parameter not provided", http.StatusBadRequest)
263                 return
264         }
265         if !disp.pool.KillContainer(uuid, "via management API: "+r.FormValue("reason")) {
266                 httpserver.Error(w, "container not found", http.StatusNotFound)
267                 return
268         }
269 }
270
271 func (disp *dispatcher) apiInstanceIdleBehavior(w http.ResponseWriter, r *http.Request, want worker.IdleBehavior) {
272         id := cloud.InstanceID(r.FormValue("instance_id"))
273         if id == "" {
274                 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
275                 return
276         }
277         err := disp.pool.SetIdleBehavior(id, want)
278         if err != nil {
279                 httpserver.Error(w, err.Error(), http.StatusNotFound)
280                 return
281         }
282 }