X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/b205525d0b7c7b9042513fe77d2e8061534208ae..a27348943c6e35eee983d6a333eb6977394c3f13:/lib/dispatchcloud/dispatcher.go

diff --git a/lib/dispatchcloud/dispatcher.go b/lib/dispatchcloud/dispatcher.go
index 1a47e8490b..2415094ac0 100644
--- a/lib/dispatchcloud/dispatcher.go
+++ b/lib/dispatchcloud/dispatcher.go
@@ -21,9 +21,9 @@ import (
 	"git.curoverse.com/arvados.git/sdk/go/arvados"
 	"git.curoverse.com/arvados.git/sdk/go/auth"
 	"git.curoverse.com/arvados.git/sdk/go/httpserver"
-	"github.com/Sirupsen/logrus"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -35,6 +35,7 @@ const (
 type pool interface {
 	scheduler.WorkerPool
 	Instances() []worker.InstanceView
+	Stop()
 }
 
 type dispatcher struct {
@@ -51,6 +52,7 @@ type dispatcher struct {
 
 	setupOnce sync.Once
 	stop      chan struct{}
+	stopped   chan struct{}
 }
 
 // Start starts the dispatcher. Start can be called multiple times
@@ -79,6 +81,7 @@ func (disp *dispatcher) Close() {
 	case disp.stop <- struct{}{}:
 	default:
 	}
+	<-disp.stopped
 }
 
 // Make a worker.Executor for the given instance.
@@ -109,6 +112,7 @@ func (disp *dispatcher) initialize() {
 		}
 	}
 	disp.stop = make(chan struct{}, 1)
+	disp.stopped = make(chan struct{})
 	disp.logger = logrus.StandardLogger()
 
 	if key, err := ssh.ParsePrivateKey(disp.Cluster.Dispatch.PrivateKey); err != nil {
@@ -117,7 +121,7 @@ func (disp *dispatcher) initialize() {
 		disp.sshKey = key
 	}
 
-	instanceSet, err := newInstanceSet(disp.Cluster, disp.InstanceSetID)
+	instanceSet, err := newInstanceSet(disp.Cluster, disp.InstanceSetID, disp.logger)
 	if err != nil {
 		disp.logger.Fatalf("error initializing driver: %s", err)
 	}
@@ -144,7 +148,9 @@ func (disp *dispatcher) initialize() {
 }
 
 func (disp *dispatcher) run() {
+	defer close(disp.stopped)
 	defer disp.instanceSet.Stop()
+	defer disp.pool.Stop()
 
 	staleLockTimeout := time.Duration(disp.Cluster.Dispatch.StaleLockTimeout)
 	if staleLockTimeout == 0 {