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
45 type dispatcher struct {
46 Cluster *arvados.Cluster
47 Context context.Context
48 InstanceSetID cloud.InstanceSetID
50 logger logrus.FieldLogger
51 reg *prometheus.Registry
52 instanceSet cloud.InstanceSet
54 queue scheduler.ContainerQueue
55 httpHandler http.Handler
63 // Start starts the dispatcher. Start can be called multiple times
64 // with no ill effect.
65 func (disp *dispatcher) Start() {
66 disp.setupOnce.Do(disp.setup)
69 // ServeHTTP implements service.Handler.
70 func (disp *dispatcher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
72 disp.httpHandler.ServeHTTP(w, r)
75 // CheckHealth implements service.Handler.
76 func (disp *dispatcher) CheckHealth() error {
81 // Stop dispatching containers and release resources. Typically used
83 func (disp *dispatcher) Close() {
86 case disp.stop <- struct{}{}:
92 // Make a worker.Executor for the given instance.
93 func (disp *dispatcher) newExecutor(inst cloud.Instance) worker.Executor {
94 exr := ssh_executor.New(inst)
95 exr.SetTargetPort(disp.Cluster.CloudVMs.SSHPort)
96 exr.SetSigners(disp.sshKey)
100 func (disp *dispatcher) typeChooser(ctr *arvados.Container) (arvados.InstanceType, error) {
101 return ChooseInstanceType(disp.Cluster, ctr)
104 func (disp *dispatcher) setup() {
109 func (disp *dispatcher) initialize() {
110 arvClient := arvados.NewClientFromEnv()
111 if disp.InstanceSetID == "" {
112 if strings.HasPrefix(arvClient.AuthToken, "v2/") {
113 disp.InstanceSetID = cloud.InstanceSetID(strings.Split(arvClient.AuthToken, "/")[1])
115 // Use some other string unique to this token
116 // that doesn't reveal the token itself.
117 disp.InstanceSetID = cloud.InstanceSetID(fmt.Sprintf("%x", md5.Sum([]byte(arvClient.AuthToken))))
120 disp.stop = make(chan struct{}, 1)
121 disp.stopped = make(chan struct{})
122 disp.logger = ctxlog.FromContext(disp.Context)
124 if key, err := ssh.ParsePrivateKey([]byte(disp.Cluster.Dispatch.PrivateKey)); err != nil {
125 disp.logger.Fatalf("error parsing configured Dispatch.PrivateKey: %s", err)
130 instanceSet, err := newInstanceSet(disp.Cluster, disp.InstanceSetID, disp.logger)
132 disp.logger.Fatalf("error initializing driver: %s", err)
134 disp.instanceSet = instanceSet
135 disp.reg = prometheus.NewRegistry()
136 disp.pool = worker.NewPool(disp.logger, arvClient, disp.reg, disp.instanceSet, disp.newExecutor, disp.sshKey.PublicKey(), disp.Cluster)
137 disp.queue = container.NewQueue(disp.logger, disp.reg, disp.typeChooser, arvClient)
139 if disp.Cluster.ManagementToken == "" {
140 disp.httpHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
141 http.Error(w, "Management API authentication is not configured", http.StatusForbidden)
144 mux := httprouter.New()
145 mux.HandlerFunc("GET", "/arvados/v1/dispatch/containers", disp.apiContainers)
146 mux.HandlerFunc("GET", "/arvados/v1/dispatch/instances", disp.apiInstances)
147 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/hold", disp.apiInstanceHold)
148 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/drain", disp.apiInstanceDrain)
149 mux.HandlerFunc("POST", "/arvados/v1/dispatch/instances/run", disp.apiInstanceRun)
150 metricsH := promhttp.HandlerFor(disp.reg, promhttp.HandlerOpts{
151 ErrorLog: disp.logger,
153 mux.Handler("GET", "/metrics", metricsH)
154 mux.Handler("GET", "/metrics.json", metricsH)
155 disp.httpHandler = auth.RequireLiteralToken(disp.Cluster.ManagementToken, mux)
159 func (disp *dispatcher) run() {
160 defer close(disp.stopped)
161 defer disp.instanceSet.Stop()
162 defer disp.pool.Stop()
164 staleLockTimeout := time.Duration(disp.Cluster.Dispatch.StaleLockTimeout)
165 if staleLockTimeout == 0 {
166 staleLockTimeout = defaultStaleLockTimeout
168 pollInterval := time.Duration(disp.Cluster.Dispatch.PollInterval)
169 if pollInterval <= 0 {
170 pollInterval = defaultPollInterval
172 sched := scheduler.New(disp.Context, disp.queue, disp.pool, staleLockTimeout, pollInterval)
179 // Management API: all active and queued containers.
180 func (disp *dispatcher) apiContainers(w http.ResponseWriter, r *http.Request) {
182 Items []container.QueueEnt `json:"items"`
184 qEntries, _ := disp.queue.Entries()
185 for _, ent := range qEntries {
186 resp.Items = append(resp.Items, ent)
188 json.NewEncoder(w).Encode(resp)
191 // Management API: all active instances (cloud VMs).
192 func (disp *dispatcher) apiInstances(w http.ResponseWriter, r *http.Request) {
194 Items []worker.InstanceView `json:"items"`
196 resp.Items = disp.pool.Instances()
197 json.NewEncoder(w).Encode(resp)
200 // Management API: set idle behavior to "hold" for specified instance.
201 func (disp *dispatcher) apiInstanceHold(w http.ResponseWriter, r *http.Request) {
202 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorHold)
205 // Management API: set idle behavior to "drain" for specified instance.
206 func (disp *dispatcher) apiInstanceDrain(w http.ResponseWriter, r *http.Request) {
207 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorDrain)
210 // Management API: set idle behavior to "run" for specified instance.
211 func (disp *dispatcher) apiInstanceRun(w http.ResponseWriter, r *http.Request) {
212 disp.apiInstanceIdleBehavior(w, r, worker.IdleBehaviorRun)
215 func (disp *dispatcher) apiInstanceIdleBehavior(w http.ResponseWriter, r *http.Request, want worker.IdleBehavior) {
216 id := cloud.InstanceID(r.FormValue("instance_id"))
218 httpserver.Error(w, "instance_id parameter not provided", http.StatusBadRequest)
221 err := disp.pool.SetIdleBehavior(id, want)
223 httpserver.Error(w, err.Error(), http.StatusNotFound)