chore: shutdown provisioner should stop waiting on client

Emyrk · Emyrk · commit 4142fb3576ef · 2024-05-01T13:00:12.000-05:00
diff --git a/enterprise/cli/provisionerdaemons.go b/enterprise/cli/provisionerdaemons.go
@@ -239,6 +239,12 @@ func (r *RootCmd) provisionerDaemonStart() *serpent.Command {
 				return xerrors.Errorf("shutdown: %w", err)
 			}
 
+			// Shutdown does not call close. Must call it manually.
+			err = srv.Close()
+			if err != nil {
+				return xerrors.Errorf("close server: %w", err)
+			}
+
 			cancel()
 			if xerrors.Is(exitErr, context.Canceled) {
 				return nil
diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go
@@ -236,6 +236,9 @@ func (p *Server) client() (proto.DRPCProvisionerDaemonClient, bool) {
 	select {
 	case <-p.closeContext.Done():
 		return nil, false
+	case <-p.shuttingDownCh:
+		// Shutting down should return a nil client and unblock
+		return nil, false
 	case client := <-p.clientCh:
 		return client, true
 	}
