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
49 InstanceSetID cloud.InstanceSetID
51 logger logrus.FieldLogger
52 reg *prometheus.Registry
53 instanceSet cloud.InstanceSet
55 queue scheduler.ContainerQueue
56 httpHandler http.Handler
64 // Start starts the dispatcher. Start can be called multiple times
65 // with no ill effect.
66 func (disp *dispatcher) Start() {
67 disp.setupOnce.Do(disp.setup)
70 // ServeHTTP implements service.Handler.
71 func (disp *dispatcher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
73 disp.httpHandler.ServeHTTP(w, r)
76 // CheckHealth implements service.Handler.
77 func (disp *dispatcher) CheckHealth() error {
82 // Stop dispatching containers and release resources. Typically used
84 func (disp *dispatcher) Close() {
87 case disp.stop <- struct{}{}:
93 // Make a worker.Executor for the given instance.
94 func (disp *dispatcher) newExecutor(inst cloud.Instance) worker.Executor {
95 exr := ssh_executor.New(inst)
96 exr.SetTargetPort(disp.Cluster.CloudVMs.SSHPort)
97 exr.SetSigners(disp.sshKey)
101 func (disp *dispatcher) typeChooser(ctr *arvados.Container) (arvados.InstanceType, error) {
102 return ChooseInstanceType(disp.Cluster, ctr)
105 func (disp *dispatcher) setup() {
110 func (disp *dispatcher) initialize() {
111 arvClient := arvados.NewClientFromEnv()
112 if disp.InstanceSetID == "" {
113 if strings.HasPrefix(arvClient.AuthToken, "v2/") {
114 disp.InstanceSetID = cloud.InstanceSetID(strings.Split(arvClient.AuthToken, "/")[1])
116 // Use some other string unique to this token
117 // that doesn't reveal the token itself.
118 disp.InstanceSetID = cloud.InstanceSetID(fmt.Sprintf("%x", md5.Sum([]byte(arvClient.AuthToken))))
121 disp.stop = make(chan struct{}, 1)
122 disp.stopped = make(chan struct{})
123 disp.logger = ctxlog.FromContext(disp.Context)
125 if key, err := ssh.ParsePrivateKey([]byte(disp.Cluster.Dispatch.PrivateKey)); err != nil {
126 disp.logger.Fatalf("error parsing configured Dispatch.PrivateKey: %s", err)
131 instanceSet, err := newInstanceSet(disp.Cluster, disp.InstanceSetID, disp.logger)
133 disp.logger.Fatalf("error initializing driver: %s", err)
135 disp.instanceSet = instanceSet
136 disp.reg = prometheus.NewRegistry()
137 disp.pool = worker.NewPool(disp.logger, arvClient, disp.reg, disp.instanceSet, disp.newExecutor, disp.sshKey.PublicKey(), disp.Cluster)
138 disp.queue = container.NewQueue(disp.logger, disp.reg, disp.typeChooser, arvClient)
140 if disp.Cluster.ManagementToken == "" {
141 disp.httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
142 http.Error(w, "Management API authentication is not configured", http.StatusForbidden)
145 mux := httprouter.New()
146 mux.HandlerFunc("GET", "/arvados/v1/dispatch/containers", disp.apiContainers)
147 mux.HandlerFunc("GET", "/arvados/v1/dispatch/instances", disp.apiInstances)
148 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/hold", disp.apiInstanceHold)
149 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/drain", disp.apiInstanceDrain)
150 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/run", disp.apiInstanceRun)
151 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/kill", disp.apiInstanceKill)
152 metricsH := promhttp.HandlerFor(disp.reg, promhttp.HandlerOpts{
153 ErrorLog: disp.logger,
155 mux.Handler("GET", "/metrics", metricsH)
156 mux.Handler("GET", "/metrics.json", metricsH)
157 disp.httpHandler = auth.RequireLiteralToken(disp.Cluster.ManagementToken, mux)
161 func (disp *dispatcher) run() {
162 defer close(disp.stopped)
163 defer disp.instanceSet.Stop()
164 defer disp.pool.Stop()
166 staleLockTimeout := time.Duration(disp.Cluster.Dispatch.StaleLockTimeout)
167 if staleLockTimeout == 0 {
168 staleLockTimeout = defaultStaleLockTimeout
170 pollInterval := time.Duration(disp.Cluster.Dispatch.PollInterval)
171 if pollInterval <= 0 {
172 pollInterval = defaultPollInterval
174 sched := scheduler.New(disp.Context, disp.queue, disp.pool, staleLockTimeout, pollInterval)
181 // Management API: all active and queued containers.
182 func (disp *dispatcher) apiContainers(w http.ResponseWriter, r *http.Request) {
184 Items []container.QueueEnt `json:"items"`
186 qEntries, _ := disp.queue.Entries()
187 for _, ent := range qEntries {
188 resp.Items = append(resp.Items, ent)
190 json.NewEncoder(w).Encode(resp)
193 // Management API: all active instances (cloud VMs).
194 func (disp *dispatcher) apiInstances(w http.ResponseWriter, r *http.Request) {
196 Items []worker.InstanceView `json:"items"`
198 resp.Items = disp.pool.Instances()
199 json.NewEncoder(w).Encode(resp)
202 // Management API: set idle behavior to "hold" for specified instance.
203 func (disp *dispatcher) apiInstanceHold(w http.ResponseWriter, r *http.Request) {
204 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorHold)
207 // Management API: set idle behavior to "drain" for specified instance.
208 func (disp *dispatcher) apiInstanceDrain(w http.ResponseWriter, r *http.Request) {
209 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorDrain)
212 // Management API: set idle behavior to "run" for specified instance.
213 func (disp *dispatcher) apiInstanceRun(w http.ResponseWriter, r *http.Request) {
214 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorRun)
217 // Management API: shutdown/destroy specified instance now.
218 func (disp *dispatcher) apiInstanceKill(w http.ResponseWriter, r *http.Request) {
219 id := cloud.InstanceID(r.FormValue("instance_id"))
221 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
224 err := disp.pool.KillInstance(id, "via management API: "+r.FormValue("reason"))
226 httpserver.Error(w, err.Error(), http.StatusNotFound)
231 func (disp *dispatcher) apiInstanceIdleBehavior(w http.ResponseWriter, r *http.Request, want worker.IdleBehavior) {
232 id := cloud.InstanceID(r.FormValue("instance_id"))
234 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
237 err := disp.pool.SetIdleBehavior(id, want)
239 httpserver.Error(w, err.Error(), http.StatusNotFound)