20755: Merge branch 'main' into 20755-ec2-multiple-subnets

[arvados.git] / lib / dispatchcloud / test / stub_driver.go

diff --git a/lib/dispatchcloud/test/stub_driver.go b/lib/dispatchcloud/test/stub_driver.go
index 1b31a71a264fabf865f981f5f94eab1649847ac4..9bc1e4c8f773ed338fcf455abc25be79808f12f5 100644 (file)
--- a/lib/dispatchcloud/test/stub_driver.go
+++ b/lib/dispatchcloud/test/stub_driver.go
@@ -18,7 +18,9 @@ import (
        "time"
 
        "git.arvados.org/arvados.git/lib/cloud"
+       "git.arvados.org/arvados.git/lib/crunchrun"
        "git.arvados.org/arvados.git/sdk/go/arvados"
+       "github.com/prometheus/client_golang/prometheus"
        "github.com/sirupsen/logrus"
        "golang.org/x/crypto/ssh"
 )
@@ -44,7 +46,8 @@ type StubDriver struct {
        Queue *Queue
 
        // Frequency of artificially introduced errors on calls to
-       // Destroy. 0=always succeed, 1=always fail.
+       // Create and Destroy. 0=always succeed, 1=always fail.
+       ErrorRateCreate  float64
        ErrorRateDestroy float64
 
        // If Create() or Instances() is called too frequently, return
@@ -61,7 +64,7 @@ type StubDriver struct {
 }
 
 // InstanceSet returns a new *StubInstanceSet.
-func (sd *StubDriver) InstanceSet(params json.RawMessage, id cloud.InstanceSetID, _ cloud.SharedResourceTags, logger logrus.FieldLogger) (cloud.InstanceSet, error) {
+func (sd *StubDriver) InstanceSet(params json.RawMessage, id cloud.InstanceSetID, _ cloud.SharedResourceTags, logger logrus.FieldLogger, reg *prometheus.Registry) (cloud.InstanceSet, error) {
        if sd.holdCloudOps == nil {
                sd.holdCloudOps = make(chan bool)
        }
@@ -107,7 +110,7 @@ type StubInstanceSet struct {
        lastInstanceID     int
 }
 
-func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID, tags cloud.InstanceTags, cmd cloud.InitCommand, authKey ssh.PublicKey) (cloud.Instance, error) {
+func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID, tags cloud.InstanceTags, initCommand cloud.InitCommand, authKey ssh.PublicKey) (cloud.Instance, error) {
        if sis.driver.HoldCloudOps {
                sis.driver.holdCloudOps <- true
        }
@@ -119,6 +122,9 @@ func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID,
        if sis.allowCreateCall.After(time.Now()) {
                return nil, RateLimitError{sis.allowCreateCall}
        }
+       if math_rand.Float64() < sis.driver.ErrorRateCreate {
+               return nil, fmt.Errorf("StubInstanceSet: rand < ErrorRateCreate %f", sis.driver.ErrorRateCreate)
+       }
        sis.allowCreateCall = time.Now().Add(sis.driver.MinTimeBetweenCreateCalls)
        ak := sis.driver.AuthorizedKeys
        if authKey != nil {
@@ -126,11 +132,11 @@ func (sis *StubInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID,
        }
        sis.lastInstanceID++
        svm := &StubVM{
+               InitCommand:  initCommand,
                sis:          sis,
                id:           cloud.InstanceID(fmt.Sprintf("inst%d,%s", sis.lastInstanceID, it.ProviderType)),
                tags:         copyTags(tags),
                providerType: it.ProviderType,
-               initCommand:  cmd,
                running:      map[string]stubProcess{},
                killing:      map[string]bool{},
        }
@@ -170,6 +176,15 @@ func (sis *StubInstanceSet) Stop() {
        sis.stopped = true
 }
 
+func (sis *StubInstanceSet) StubVMs() (svms []*StubVM) {
+       sis.mtx.Lock()
+       defer sis.mtx.Unlock()
+       for _, vm := range sis.servers {
+               svms = append(svms, vm)
+       }
+       return
+}
+
 type RateLimitError struct{ Retry time.Time }
 
 func (e RateLimitError) Error() string            { return fmt.Sprintf("rate limited until %s", e.Retry) }
@@ -193,12 +208,14 @@ type StubVM struct {
        ArvMountDeadlockRate  float64
        ExecuteContainer      func(arvados.Container) int
        CrashRunningContainer func(arvados.Container)
-       ExtraCrunchRunArgs    string // extra args expected after "crunch-run --detach --stdin-env "
+       ExtraCrunchRunArgs    string // extra args expected after "crunch-run --detach --stdin-config "
+
+       // Populated by (*StubInstanceSet)Create()
+       InitCommand cloud.InitCommand
 
        sis          *StubInstanceSet
        id           cloud.InstanceID
        tags         cloud.InstanceTags
-       initCommand  cloud.InitCommand
        providerType string
        SSHService   SSHService
        running      map[string]stubProcess
@@ -252,15 +269,15 @@ func (svm *StubVM) Exec(env map[string]string, command string, stdin io.Reader,
                fmt.Fprint(stderr, "crunch-run: command not found\n")
                return 1
        }
-       if strings.HasPrefix(command, "crunch-run --detach --stdin-env "+svm.ExtraCrunchRunArgs) {
-               var stdinKV map[string]string
-               err := json.Unmarshal(stdinData, &stdinKV)
+       if strings.HasPrefix(command, "crunch-run --detach --stdin-config "+svm.ExtraCrunchRunArgs) {
+               var configData crunchrun.ConfigData
+               err := json.Unmarshal(stdinData, &configData)
                if err != nil {
                        fmt.Fprintf(stderr, "unmarshal stdin: %s (stdin was: %q)\n", err, stdinData)
                        return 1
                }
                for _, name := range []string{"ARVADOS_API_HOST", "ARVADOS_API_TOKEN"} {
-                       if stdinKV[name] == "" {
+                       if configData.Env[name] == "" {
                                fmt.Fprintf(stderr, "%s env var missing from stdin %q\n", name, stdinData)
                                return 1
                        }
@@ -469,3 +486,7 @@ func copyTags(src cloud.InstanceTags) cloud.InstanceTags {
        }
        return dst
 }
+
+func (si stubInstance) PriceHistory(arvados.InstanceType) []cloud.InstancePrice {
+       return nil
+}