X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/12cd1df493cc0da3c6ed90b467470b336a36b42b..00ff9cf6b638f78df82fb5e15ca6788c8d095f40:/lib/dispatchcloud/test/stub_driver.go

diff --git a/lib/dispatchcloud/test/stub_driver.go b/lib/dispatchcloud/test/stub_driver.go
index 01af8e6d54..6e0b129487 100644
--- a/lib/dispatchcloud/test/stub_driver.go
+++ b/lib/dispatchcloud/test/stub_driver.go
@@ -20,6 +20,7 @@ import (
 	"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"
 )
@@ -45,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
@@ -53,6 +55,8 @@ type StubDriver struct {
 	MinTimeBetweenCreateCalls    time.Duration
 	MinTimeBetweenInstancesCalls time.Duration
 
+	QuotaMaxInstances int
+
 	// If true, Create and Destroy calls block until Release() is
 	// called.
 	HoldCloudOps bool
@@ -62,7 +66,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)
 	}
@@ -108,7 +112,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
 	}
@@ -120,6 +124,12 @@ 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)
+	}
+	if max := sis.driver.QuotaMaxInstances; max > 0 && len(sis.servers) >= max {
+		return nil, QuotaError{fmt.Errorf("StubInstanceSet: reached QuotaMaxInstances %d", max)}
+	}
 	sis.allowCreateCall = time.Now().Add(sis.driver.MinTimeBetweenCreateCalls)
 	ak := sis.driver.AuthorizedKeys
 	if authKey != nil {
@@ -127,11 +137,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{},
 	}
@@ -171,6 +181,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) }
@@ -196,10 +215,12 @@ type StubVM struct {
 	CrashRunningContainer func(arvados.Container)
 	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
@@ -474,3 +495,9 @@ func copyTags(src cloud.InstanceTags) cloud.InstanceTags {
 func (si stubInstance) PriceHistory(arvados.InstanceType) []cloud.InstancePrice {
 	return nil
 }
+
+type QuotaError struct {
+	error
+}
+
+func (QuotaError) IsQuotaError() bool { return true }