From d57877af31b8ba88f63390c456eac0dba8b15a47 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Magiera?= <magik6k@gmail.com>
Date: Mon, 27 Apr 2020 22:59:17 +0200
Subject: [PATCH] sched: Fix error handling in assignWorker

---
 sched.go | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/sched.go b/sched.go
index 78f2a6664..2d988d082 100644
--- a/sched.go
+++ b/sched.go
@@ -247,8 +247,28 @@ func (sh *scheduler) assignWorker(wid WorkerID, w *workerHandle, req *workerRequ
 
 	go func() {
 		err := req.prepare(req.ctx, w.w)
-
 		sh.workersLk.Lock()
+
+		if err != nil {
+			w.preparing.free(w.info.Resources, needRes)
+			sh.workersLk.Unlock()
+
+			select {
+			case sh.workerFree <- wid:
+			case <-sh.closing:
+				log.Warnf("scheduler closed while sending response (prepare error: %+v)", err)
+			}
+
+			select {
+			case req.ret <- workerResponse{err: err}:
+			case <-req.ctx.Done():
+				log.Warnf("request got cancelled before we could respond (prepare error: %+v)", err)
+			case <-sh.closing:
+				log.Warnf("scheduler closed while sending response (prepare error: %+v)", err)
+			}
+			return
+		}
+
 		err = w.active.withResources(sh.spt, wid, w.info.Resources, needRes, &sh.workersLk, func() error {
 			w.preparing.free(w.info.Resources, needRes)
 			sh.workersLk.Unlock()