X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/55da806a494b13e88ee1b0c1e6cad211786628f3..abd21f165dbfd84eaea9ca1d7f20ce195c52d679:/lib/dispatchcloud/worker/pool_test.go

diff --git a/lib/dispatchcloud/worker/pool_test.go b/lib/dispatchcloud/worker/pool_test.go
index 3867e2c63e..526bc552cd 100644
--- a/lib/dispatchcloud/worker/pool_test.go
+++ b/lib/dispatchcloud/worker/pool_test.go
@@ -5,13 +5,14 @@
 package worker
 
 import (
-	"io"
+	"sort"
+	"strings"
 	"time"
 
 	"git.curoverse.com/arvados.git/lib/cloud"
 	"git.curoverse.com/arvados.git/lib/dispatchcloud/test"
 	"git.curoverse.com/arvados.git/sdk/go/arvados"
-	"github.com/Sirupsen/logrus"
+	"github.com/prometheus/client_golang/prometheus"
 	check "gopkg.in/check.v1"
 )
 
@@ -31,29 +32,114 @@ var less = &lessChecker{&check.CheckerInfo{Name: "less", Params: []string{"obtai
 
 type PoolSuite struct{}
 
-func (suite *PoolSuite) SetUpSuite(c *check.C) {
-	logrus.StandardLogger().SetLevel(logrus.DebugLevel)
-}
+func (suite *PoolSuite) TestResumeAfterRestart(c *check.C) {
+	type1 := test.InstanceType(1)
+	type2 := test.InstanceType(2)
+	type3 := test.InstanceType(3)
+	waitForIdle := func(pool *Pool, notify <-chan struct{}) {
+		timeout := time.NewTimer(time.Second)
+		for {
+			instances := pool.Instances()
+			sort.Slice(instances, func(i, j int) bool {
+				return strings.Compare(instances[i].ArvadosInstanceType, instances[j].ArvadosInstanceType) < 0
+			})
+			if len(instances) == 3 &&
+				instances[0].ArvadosInstanceType == type1.Name &&
+				instances[0].WorkerState == StateIdle.String() &&
+				instances[1].ArvadosInstanceType == type1.Name &&
+				instances[1].WorkerState == StateIdle.String() &&
+				instances[2].ArvadosInstanceType == type2.Name &&
+				instances[2].WorkerState == StateIdle.String() {
+				return
+			}
+			select {
+			case <-timeout.C:
+				c.Logf("pool.Instances() == %#v", instances)
+				c.Error("timed out")
+				return
+			case <-notify:
+			}
+		}
+	}
 
-func (suite *PoolSuite) TestStartContainer(c *check.C) {
-	// TODO: use an instanceSet stub with an SSH server
-}
+	logger := test.Logger()
+	driver := &test.StubDriver{}
+	is, err := driver.InstanceSet(nil, "", logger)
+	c.Assert(err, check.IsNil)
+
+	newExecutor := func(cloud.Instance) Executor {
+		return stubExecutor{
+			"crunch-run --list": stubResp{},
+			"true":              stubResp{},
+		}
+	}
+
+	cluster := &arvados.Cluster{
+		Dispatch: arvados.Dispatch{
+			MaxProbesPerSecond: 1000,
+			ProbeInterval:      arvados.Duration(time.Millisecond * 10),
+		},
+		CloudVMs: arvados.CloudVMs{
+			BootProbeCommand: "true",
+			SyncInterval:     arvados.Duration(time.Millisecond * 10),
+		},
+		InstanceTypes: arvados.InstanceTypeMap{
+			type1.Name: type1,
+			type2.Name: type2,
+			type3.Name: type3,
+		},
+	}
 
-func (suite *PoolSuite) TestVerifyHostKey(c *check.C) {
-	// TODO: use an instanceSet stub with an SSH server
+	pool := NewPool(logger, arvados.NewClientFromEnv(), prometheus.NewRegistry(), is, newExecutor, nil, cluster)
+	notify := pool.Subscribe()
+	defer pool.Unsubscribe(notify)
+	pool.Create(type1)
+	pool.Create(type1)
+	pool.Create(type2)
+	waitForIdle(pool, notify)
+	var heldInstanceID cloud.InstanceID
+	for _, inst := range pool.Instances() {
+		if inst.ArvadosInstanceType == type2.Name {
+			heldInstanceID = cloud.InstanceID(inst.Instance)
+			pool.SetIdleBehavior(heldInstanceID, IdleBehaviorHold)
+		}
+	}
+	pool.Stop()
+
+	c.Log("------- starting new pool, waiting to recover state")
+
+	pool2 := NewPool(logger, arvados.NewClientFromEnv(), prometheus.NewRegistry(), is, newExecutor, nil, cluster)
+	notify2 := pool2.Subscribe()
+	defer pool2.Unsubscribe(notify2)
+	waitForIdle(pool2, notify2)
+	for _, inst := range pool2.Instances() {
+		if inst.ArvadosInstanceType == type2.Name {
+			c.Check(inst.Instance, check.Equals, heldInstanceID)
+			c.Check(inst.IdleBehavior, check.Equals, IdleBehaviorHold)
+		} else {
+			c.Check(inst.IdleBehavior, check.Equals, IdleBehaviorRun)
+		}
+	}
+	pool2.Stop()
 }
 
 func (suite *PoolSuite) TestCreateUnallocShutdown(c *check.C) {
-	lameInstanceSet := &test.LameInstanceSet{Hold: make(chan bool)}
+	logger := test.Logger()
+	driver := test.StubDriver{HoldCloudOps: true}
+	instanceSet, err := driver.InstanceSet(nil, "", logger)
+	c.Assert(err, check.IsNil)
+
 	type1 := arvados.InstanceType{Name: "a1s", ProviderType: "a1.small", VCPUs: 1, RAM: 1 * GiB, Price: .01}
 	type2 := arvados.InstanceType{Name: "a2m", ProviderType: "a2.medium", VCPUs: 2, RAM: 2 * GiB, Price: .02}
+	type3 := arvados.InstanceType{Name: "a2l", ProviderType: "a2.large", VCPUs: 4, RAM: 4 * GiB, Price: .04}
 	pool := &Pool{
-		logger:      logrus.StandardLogger(),
-		newExecutor: func(cloud.Instance) Executor { return &stubExecutor{} },
-		instanceSet: lameInstanceSet,
+		logger:      logger,
+		newExecutor: func(cloud.Instance) Executor { return stubExecutor{} },
+		instanceSet: &throttledInstanceSet{InstanceSet: instanceSet},
 		instanceTypes: arvados.InstanceTypeMap{
 			type1.Name: type1,
 			type2.Name: type2,
+			type3.Name: type3,
 		},
 	}
 	notify := pool.Subscribe()
@@ -63,23 +149,42 @@ func (suite *PoolSuite) TestCreateUnallocShutdown(c *check.C) {
 
 	c.Check(pool.Unallocated()[type1], check.Equals, 0)
 	c.Check(pool.Unallocated()[type2], check.Equals, 0)
+	c.Check(pool.Unallocated()[type3], check.Equals, 0)
 	pool.Create(type2)
 	pool.Create(type1)
 	pool.Create(type2)
+	pool.Create(type3)
 	c.Check(pool.Unallocated()[type1], check.Equals, 1)
 	c.Check(pool.Unallocated()[type2], check.Equals, 2)
+	c.Check(pool.Unallocated()[type3], check.Equals, 1)
 
 	// Unblock the pending Create calls.
-	go lameInstanceSet.Release(3)
+	go driver.ReleaseCloudOps(4)
 
 	// Wait for each instance to either return from its Create
 	// call, or show up in a poll.
 	suite.wait(c, pool, notify, func() bool {
 		pool.mtx.RLock()
 		defer pool.mtx.RUnlock()
-		return len(pool.workers) == 3
+		return len(pool.workers) == 4
 	})
 
+	// Place type3 node on admin-hold
+	ivs := suite.instancesByType(pool, type3)
+	c.Assert(ivs, check.HasLen, 1)
+	type3instanceID := ivs[0].Instance
+	err = pool.SetIdleBehavior(type3instanceID, IdleBehaviorHold)
+	c.Check(err, check.IsNil)
+
+	// Check admin-hold behavior: refuse to shutdown, and don't
+	// report as Unallocated ("available now or soon").
+	c.Check(pool.Shutdown(type3), check.Equals, false)
+	suite.wait(c, pool, notify, func() bool {
+		return pool.Unallocated()[type3] == 0
+	})
+	c.Check(suite.instancesByType(pool, type3), check.HasLen, 1)
+
+	// Shutdown both type2 nodes
 	c.Check(pool.Shutdown(type2), check.Equals, true)
 	suite.wait(c, pool, notify, func() bool {
 		return pool.Unallocated()[type1] == 1 && pool.Unallocated()[type2] == 1
@@ -99,16 +204,58 @@ func (suite *PoolSuite) TestCreateUnallocShutdown(c *check.C) {
 		}
 		break
 	}
+
+	// Shutdown type1 node
 	c.Check(pool.Shutdown(type1), check.Equals, true)
 	suite.wait(c, pool, notify, func() bool {
-		return pool.Unallocated()[type1] == 0 && pool.Unallocated()[type2] == 0
+		return pool.Unallocated()[type1] == 0 && pool.Unallocated()[type2] == 0 && pool.Unallocated()[type3] == 0
 	})
 	select {
 	case <-notify2:
 	case <-time.After(time.Second):
 		c.Error("notify did not receive")
 	}
-	go lameInstanceSet.Release(3) // unblock Destroy calls
+
+	// Put type3 node back in service.
+	err = pool.SetIdleBehavior(type3instanceID, IdleBehaviorRun)
+	c.Check(err, check.IsNil)
+	suite.wait(c, pool, notify, func() bool {
+		return pool.Unallocated()[type3] == 1
+	})
+
+	// Check admin-drain behavior: shut down right away, and don't
+	// report as Unallocated.
+	err = pool.SetIdleBehavior(type3instanceID, IdleBehaviorDrain)
+	c.Check(err, check.IsNil)
+	suite.wait(c, pool, notify, func() bool {
+		return pool.Unallocated()[type3] == 0
+	})
+	suite.wait(c, pool, notify, func() bool {
+		ivs := suite.instancesByType(pool, type3)
+		return len(ivs) == 1 && ivs[0].WorkerState == StateShutdown.String()
+	})
+
+	// Unblock all pending Destroy calls. Pool calls Destroy again
+	// if a node still appears in the provider list after a
+	// previous attempt, so there might be more than 4 Destroy
+	// calls to unblock.
+	go driver.ReleaseCloudOps(4444)
+
+	// Sync until all instances disappear from the provider list.
+	suite.wait(c, pool, notify, func() bool {
+		pool.getInstancesAndSync()
+		return len(pool.Instances()) == 0
+	})
+}
+
+func (suite *PoolSuite) instancesByType(pool *Pool, it arvados.InstanceType) []InstanceView {
+	var ivs []InstanceView
+	for _, iv := range pool.Instances() {
+		if iv.ArvadosInstanceType == it.Name {
+			ivs = append(ivs, iv)
+		}
+	}
+	return ivs
 }
 
 func (suite *PoolSuite) wait(c *check.C, pool *Pool, notify <-chan struct{}, ready func() bool) {
@@ -123,13 +270,3 @@ func (suite *PoolSuite) wait(c *check.C, pool *Pool, notify <-chan struct{}, rea
 	}
 	c.Check(ready(), check.Equals, true)
 }
-
-type stubExecutor struct{}
-
-func (*stubExecutor) SetTarget(cloud.ExecutorTarget) {}
-
-func (*stubExecutor) Execute(cmd string, stdin io.Reader) ([]byte, []byte, error) {
-	return nil, nil, nil
-}
-
-func (*stubExecutor) Close() {}