1 // Copyright (C) The Arvados Authors. All rights reserved.
3 // SPDX-License-Identifier: AGPL-3.0
17 "git.curoverse.com/arvados.git/lib/cloud"
18 "git.curoverse.com/arvados.git/lib/dispatchcloud/container"
19 "git.curoverse.com/arvados.git/lib/dispatchcloud/scheduler"
20 "git.curoverse.com/arvados.git/lib/dispatchcloud/ssh_executor"
21 "git.curoverse.com/arvados.git/lib/dispatchcloud/worker"
22 "git.curoverse.com/arvados.git/sdk/go/arvados"
23 "git.curoverse.com/arvados.git/sdk/go/auth"
24 "git.curoverse.com/arvados.git/sdk/go/ctxlog"
25 "git.curoverse.com/arvados.git/sdk/go/httpserver"
26 "github.com/julienschmidt/httprouter"
27 "github.com/prometheus/client_golang/prometheus"
28 "github.com/prometheus/client_golang/prometheus/promhttp"
29 "github.com/sirupsen/logrus"
30 "golang.org/x/crypto/ssh"
34 defaultPollInterval = time.Second
35 defaultStaleLockTimeout = time.Minute
40 Instances() []worker.InstanceView
41 SetIdleBehavior(cloud.InstanceID, worker.IdleBehavior) error
42 KillInstance(id cloud.InstanceID, reason string) error
46 type dispatcher struct {
47 Cluster *arvados.Cluster
48 Context context.Context
50 InstanceSetID cloud.InstanceSetID
52 logger logrus.FieldLogger
53 reg *prometheus.Registry
54 instanceSet cloud.InstanceSet
56 queue scheduler.ContainerQueue
57 httpHandler http.Handler
65 // Start starts the dispatcher. Start can be called multiple times
66 // with no ill effect.
67 func (disp *dispatcher) Start() {
68 disp.setupOnce.Do(disp.setup)
71 // ServeHTTP implements service.Handler.
72 func (disp *dispatcher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
74 disp.httpHandler.ServeHTTP(w, r)
77 // CheckHealth implements service.Handler.
78 func (disp *dispatcher) CheckHealth() error {
83 // Stop dispatching containers and release resources. Typically used
85 func (disp *dispatcher) Close() {
88 case disp.stop <- struct{}{}:
94 // Make a worker.Executor for the given instance.
95 func (disp *dispatcher) newExecutor(inst cloud.Instance) worker.Executor {
96 exr := ssh_executor.New(inst)
97 exr.SetTargetPort(disp.Cluster.CloudVMs.SSHPort)
98 exr.SetSigners(disp.sshKey)
102 func (disp *dispatcher) typeChooser(ctr *arvados.Container) (arvados.InstanceType, error) {
103 return ChooseInstanceType(disp.Cluster, ctr)
106 func (disp *dispatcher) setup() {
111 func (disp *dispatcher) initialize() {
112 disp.logger = ctxlog.FromContext(disp.Context)
114 arvClient, err := arvados.NewClientFromConfig(disp.Cluster)
116 disp.logger.WithError(err).Warn("error initializing client from cluster config, falling back to ARVADOS_API_HOST(_INSECURE) environment variables")
117 arvClient = arvados.NewClientFromEnv()
119 arvClient.AuthToken = disp.AuthToken
121 if disp.InstanceSetID == "" {
122 if strings.HasPrefix(arvClient.AuthToken, "v2/") {
123 disp.InstanceSetID = cloud.InstanceSetID(strings.Split(arvClient.AuthToken, "/")[1])
125 // Use some other string unique to this token
126 // that doesn't reveal the token itself.
127 disp.InstanceSetID = cloud.InstanceSetID(fmt.Sprintf("%x", md5.Sum([]byte(arvClient.AuthToken))))
130 disp.stop = make(chan struct{}, 1)
131 disp.stopped = make(chan struct{})
133 if key, err := ssh.ParsePrivateKey([]byte(disp.Cluster.Dispatch.PrivateKey)); err != nil {
134 disp.logger.Fatalf("error parsing configured Dispatch.PrivateKey: %s", err)
139 instanceSet, err := newInstanceSet(disp.Cluster, disp.InstanceSetID, disp.logger)
141 disp.logger.Fatalf("error initializing driver: %s", err)
143 disp.instanceSet = instanceSet
144 disp.reg = prometheus.NewRegistry()
145 disp.pool = worker.NewPool(disp.logger, arvClient, disp.reg, disp.instanceSet, disp.newExecutor, disp.sshKey.PublicKey(), disp.Cluster)
146 disp.queue = container.NewQueue(disp.logger, disp.reg, disp.typeChooser, arvClient)
148 if disp.Cluster.ManagementToken == "" {
149 disp.httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
150 http.Error(w, "Management API authentication is not configured", http.StatusForbidden)
153 mux := httprouter.New()
154 mux.HandlerFunc("GET", "/arvados/v1/dispatch/containers", disp.apiContainers)
155 mux.HandlerFunc("GET", "/arvados/v1/dispatch/instances", disp.apiInstances)
156 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/hold", disp.apiInstanceHold)
157 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/drain", disp.apiInstanceDrain)
158 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/run", disp.apiInstanceRun)
159 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/kill", disp.apiInstanceKill)
160 metricsH := promhttp.HandlerFor(disp.reg, promhttp.HandlerOpts{
161 ErrorLog: disp.logger,
163 mux.Handler("GET", "/metrics", metricsH)
164 mux.Handler("GET", "/metrics.json", metricsH)
165 disp.httpHandler = auth.RequireLiteralToken(disp.Cluster.ManagementToken, mux)
169 func (disp *dispatcher) run() {
170 defer close(disp.stopped)
171 defer disp.instanceSet.Stop()
172 defer disp.pool.Stop()
174 staleLockTimeout := time.Duration(disp.Cluster.Dispatch.StaleLockTimeout)
175 if staleLockTimeout == 0 {
176 staleLockTimeout = defaultStaleLockTimeout
178 pollInterval := time.Duration(disp.Cluster.Dispatch.PollInterval)
179 if pollInterval <= 0 {
180 pollInterval = defaultPollInterval
182 sched := scheduler.New(disp.Context, disp.queue, disp.pool, staleLockTimeout, pollInterval)
189 // Management API: all active and queued containers.
190 func (disp *dispatcher) apiContainers(w http.ResponseWriter, r *http.Request) {
192 Items []container.QueueEnt `json:"items"`
194 qEntries, _ := disp.queue.Entries()
195 for _, ent := range qEntries {
196 resp.Items = append(resp.Items, ent)
198 json.NewEncoder(w).Encode(resp)
201 // Management API: all active instances (cloud VMs).
202 func (disp *dispatcher) apiInstances(w http.ResponseWriter, r *http.Request) {
204 Items []worker.InstanceView `json:"items"`
206 resp.Items = disp.pool.Instances()
207 json.NewEncoder(w).Encode(resp)
210 // Management API: set idle behavior to "hold" for specified instance.
211 func (disp *dispatcher) apiInstanceHold(w http.ResponseWriter, r *http.Request) {
212 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorHold)
215 // Management API: set idle behavior to "drain" for specified instance.
216 func (disp *dispatcher) apiInstanceDrain(w http.ResponseWriter, r *http.Request) {
217 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorDrain)
220 // Management API: set idle behavior to "run" for specified instance.
221 func (disp *dispatcher) apiInstanceRun(w http.ResponseWriter, r *http.Request) {
222 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorRun)
225 // Management API: shutdown/destroy specified instance now.
226 func (disp *dispatcher) apiInstanceKill(w http.ResponseWriter, r *http.Request) {
227 id := cloud.InstanceID(r.FormValue("instance_id"))
229 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
232 err := disp.pool.KillInstance(id, "via management API: "+r.FormValue("reason"))
234 httpserver.Error(w, err.Error(), http.StatusNotFound)
239 func (disp *dispatcher) apiInstanceIdleBehavior(w http.ResponseWriter, r *http.Request, want worker.IdleBehavior) {
240 id := cloud.InstanceID(r.FormValue("instance_id"))
242 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
245 err := disp.pool.SetIdleBehavior(id, want)
247 httpserver.Error(w, err.Error(), http.StatusNotFound)