15964: Remove qr1hi from a few more places. Delete unused includes.
[arvados.git] / lib / dispatchcloud / test / stub_driver.go
index 28fee6d9278db08482c9daa2f575e217bc6c6bb3..7a1f42301684a5a042f555c3e17ccda5d2f8b6c5 100644 (file)
@@ -17,8 +17,8 @@ import (
        "sync"
        "time"
 
-       "git.curoverse.com/arvados.git/lib/cloud"
-       "git.curoverse.com/arvados.git/sdk/go/arvados"
+       "git.arvados.org/arvados.git/lib/cloud"
+       "git.arvados.org/arvados.git/sdk/go/arvados"
        "github.com/sirupsen/logrus"
        "golang.org/x/crypto/ssh"
 )
@@ -56,7 +56,7 @@ type StubDriver struct {
 }
 
 // InstanceSet returns a new *StubInstanceSet.
-func (sd *StubDriver) InstanceSet(params json.RawMessage, id cloud.InstanceSetID, logger logrus.FieldLogger) (cloud.InstanceSet, error) {
+func (sd *StubDriver) InstanceSet(params json.RawMessage, id cloud.InstanceSetID, _ cloud.SharedResourceTags, logger logrus.FieldLogger) (cloud.InstanceSet, error) {
        if sd.holdCloudOps == nil {
                sd.holdCloudOps = make(chan bool)
        }
@@ -126,6 +126,8 @@ func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID,
                tags:         copyTags(tags),
                providerType: it.ProviderType,
                initCommand:  cmd,
+               running:      map[string]int64{},
+               killing:      map[string]bool{},
        }
        svm.SSHService = SSHService{
                HostKey:        sis.driver.HostKey,
@@ -179,6 +181,7 @@ func (e RateLimitError) EarliestRetry() time.Time { return e.Retry }
 type StubVM struct {
        Boot                  time.Time
        Broken                time.Time
+       ReportBroken          time.Time
        CrunchRunMissing      bool
        CrunchRunCrashRate    float64
        CrunchRunDetachDelay  time.Duration
@@ -192,6 +195,7 @@ type StubVM struct {
        providerType string
        SSHService   SSHService
        running      map[string]int64
+       killing      map[string]bool
        lastPID      int64
        sync.Mutex
 }
@@ -241,16 +245,13 @@ func (svm *StubVM) Exec(env map[string]string, command string, stdin io.Reader,
                }
                for _, name := range []string{"ARVADOS_API_HOST", "ARVADOS_API_TOKEN"} {
                        if stdinKV[name] == "" {
-                               fmt.Fprintf(stderr, "%s env var missing from stdin %q\n", name, stdin)
+                               fmt.Fprintf(stderr, "%s env var missing from stdin %q\n", name, stdinData)
                                return 1
                        }
                }
                svm.Lock()
                svm.lastPID++
                pid := svm.lastPID
-               if svm.running == nil {
-                       svm.running = map[string]int64{}
-               }
                svm.running[uuid] = pid
                svm.Unlock()
                time.Sleep(svm.CrunchRunDetachDelay)
@@ -314,26 +315,33 @@ func (svm *StubVM) Exec(env map[string]string, command string, stdin io.Reader,
                for uuid := range svm.running {
                        fmt.Fprintf(stdout, "%s\n", uuid)
                }
+               if !svm.ReportBroken.IsZero() && svm.ReportBroken.Before(time.Now()) {
+                       fmt.Fprintln(stdout, "broken")
+               }
                return 0
        }
        if strings.HasPrefix(command, "crunch-run --kill ") {
                svm.Lock()
                pid, running := svm.running[uuid]
-               svm.Unlock()
-               go func() {
-                       time.Sleep(time.Duration(math_rand.Float64()*20) * time.Millisecond)
+               if running && !svm.killing[uuid] {
+                       svm.killing[uuid] = true
+                       go func() {
+                               time.Sleep(time.Duration(math_rand.Float64()*30) * time.Millisecond)
+                               svm.Lock()
+                               defer svm.Unlock()
+                               if svm.running[uuid] == pid {
+                                       // Kill only if the running entry
+                                       // hasn't since been killed and
+                                       // replaced with a different one.
+                                       delete(svm.running, uuid)
+                               }
+                               delete(svm.killing, uuid)
+                       }()
+                       svm.Unlock()
+                       time.Sleep(time.Duration(math_rand.Float64()*2) * time.Millisecond)
                        svm.Lock()
-                       defer svm.Unlock()
-                       if svm.running[uuid] == pid {
-                               // Kill only if the running entry
-                               // hasn't since been killed and
-                               // replaced with a different one.
-                               delete(svm.running, uuid)
-                       }
-               }()
-               time.Sleep(time.Duration(math_rand.Float64()*2) * time.Millisecond)
-               svm.Lock()
-               _, running = svm.running[uuid]
+                       _, running = svm.running[uuid]
+               }
                svm.Unlock()
                if running {
                        fmt.Fprintf(stderr, "%s: container is running\n", uuid)