Bugfix: Return notebook worker to the pool when query is cancelled. (#…

…3252)

This potentially fixes #3201

services/notebook/worker.go

@@ -410,14 +410,16 @@ func (self *NotebookWorker) RegisterWorker(
 
 		case job, ok := <-job_chan:
 			if !ok {
-				return errors.New("Cancellation")
+				job.Done("", errors.New("Cancellation"))
+				return nil
 			}
 
 			request := &NotebookRequest{}
 			err := json.Unmarshal([]byte(job.Job), request)
 			if err != nil {
 				logger.Error("NotebookManager: Invalid job request in worker: %v: %v",
 					err, job.Job)
+				job.Done("", err)
 				continue
 			}
 
@@ -570,6 +572,8 @@ func (self *NotebookWorker) startNanny(
 			}
 			scope.Log("ERROR:NotebookManager: Detected cell %v is cancelled. Stopping.", cell_id)
 			notifier.NotifyDirectListener(cell_id)
+
+			return
 		}
 	}
 }

services/scheduler/scheduler.go

@@ -167,20 +167,32 @@ func (self *Scheduler) Schedule(ctx context.Context,
 		})
 
 		for _, w := range available_workers {
+			// The worker can get back to the pool immediately while
+			// we wait for our consumer.
 			job.Done = func(result string, err error) {
-				result_chan <- services.JobResponse{
+				w.SetBusy(false)
+				w.SetRequest(vfilter.Null{})
+				defer close(result_chan)
+
+				select {
+				case <-ctx.Done():
+					return
+
+				case result_chan <- services.JobResponse{
 					Job: result,
 					Err: err,
+				}:
 				}
-
-				w.SetRequest(vfilter.Null{})
-				w.SetBusy(false)
-				close(result_chan)
 			}
 
 			select {
+
+			// If the caller is cancelled we can return the worker to
+			// the pool.
 			case <-ctx.Done():
 				self.mu.Unlock()
+				w.SetBusy(false)
+				w.SetRequest(vfilter.Null{})
 				close(result_chan)
 				return result_chan, errors.New("Cancelled")