Merge branch 'main' into 19582-aws-s3v2-driver

[arvados.git] / services / crunch-dispatch-slurm / crunch-dispatch-slurm.go

diff --git a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
index c31d7997522fa1caa73507a009d680b3835a2f46..1c0f6ad28f5ba7b6d20bcc0cadbef0fb87fec634 100644 (file)
--- a/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
+++ b/services/crunch-dispatch-slurm/crunch-dispatch-slurm.go
@@ -7,6 +7,8 @@ package dispatchslurm
 
 import (
        "context"
+       "crypto/hmac"
+       "crypto/sha256"
        "fmt"
        "log"
        "math"
@@ -17,6 +19,8 @@ import (
        "time"
 
        "git.arvados.org/arvados.git/lib/cmd"
+       "git.arvados.org/arvados.git/lib/controller/dblock"
+       "git.arvados.org/arvados.git/lib/ctrlctx"
        "git.arvados.org/arvados.git/lib/dispatchcloud"
        "git.arvados.org/arvados.git/lib/service"
        "git.arvados.org/arvados.git/sdk/go/arvados"
@@ -53,10 +57,11 @@ const initialNiceValue int64 = 10000
 
 type Dispatcher struct {
        *dispatch.Dispatcher
-       logger  logrus.FieldLogger
-       cluster *arvados.Cluster
-       sqCheck *SqueueChecker
-       slurm   Slurm
+       logger      logrus.FieldLogger
+       cluster     *arvados.Cluster
+       sqCheck     *SqueueChecker
+       slurm       Slurm
+       dbConnector ctrlctx.DBConnector
 
        done chan struct{}
        err  error
@@ -88,6 +93,7 @@ func (disp *Dispatcher) configure() error {
        disp.Client.APIHost = disp.cluster.Services.Controller.ExternalURL.Host
        disp.Client.AuthToken = disp.cluster.SystemRootToken
        disp.Client.Insecure = disp.cluster.TLS.Insecure
+       disp.dbConnector = ctrlctx.DBConnector{PostgreSQL: disp.cluster.PostgreSQL}
 
        if disp.Client.APIHost != "" || disp.Client.AuthToken != "" {
                // Copy real configs into env vars so [a]
@@ -99,7 +105,6 @@ func (disp *Dispatcher) configure() error {
                if disp.Client.Insecure {
                        os.Setenv("ARVADOS_API_HOST_INSECURE", "1")
                }
-               os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
                for k, v := range disp.cluster.Containers.SLURM.SbatchEnvironmentVariables {
                        os.Setenv(k, v)
                }
@@ -136,6 +141,8 @@ func (disp *Dispatcher) setup() {
 }
 
 func (disp *Dispatcher) run() error {
+       dblock.Dispatch.Lock(context.Background(), disp.dbConnector.GetDB)
+       defer dblock.Dispatch.Unlock()
        defer disp.sqCheck.Stop()
 
        if disp.cluster != nil && len(disp.cluster.InstanceTypes) > 0 {
@@ -213,7 +220,12 @@ func (disp *Dispatcher) submit(container arvados.Container, crunchRunCommand []s
        crArgs := append([]string(nil), crunchRunCommand...)
        crArgs = append(crArgs, "--runtime-engine="+disp.cluster.Containers.RuntimeEngine)
        crArgs = append(crArgs, container.UUID)
-       crScript := strings.NewReader(execScript(crArgs))
+
+       h := hmac.New(sha256.New, []byte(disp.cluster.SystemRootToken))
+       fmt.Fprint(h, container.UUID)
+       authsecret := fmt.Sprintf("%x", h.Sum(nil))
+
+       crScript := strings.NewReader(execScript(crArgs, map[string]string{"GatewayAuthSecret": authsecret}))
 
        sbArgs, err := disp.sbatchArgs(container)
        if err != nil {