From be4112ee68263b2523bfb1670ee7c4c34ca0a336 Mon Sep 17 00:00:00 2001
From: Howard Yeh <hayeah@gmail.com>
Date: Tue, 23 Jun 2020 17:42:47 +0800
Subject: [PATCH] add RPC timeout to maybeSchedRequest

---
 sched.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/sched.go b/sched.go
index 715a823e1..5b99f48ff 100644
--- a/sched.go
+++ b/sched.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"sort"
 	"sync"
+	"time"
 
 	"github.com/hashicorp/go-multierror"
 	"golang.org/x/xerrors"
@@ -214,7 +215,10 @@ func (sh *scheduler) maybeSchedRequest(req *workerRequest) (bool, error) {
 	needRes := ResourceTable[req.taskType][sh.spt]
 
 	for wid, worker := range sh.workers {
-		ok, err := req.sel.Ok(req.ctx, req.taskType, sh.spt, worker)
+		rpcCtx, cancel := context.WithTimeout(req.ctx, 5*time.Second)
+		ok, err := req.sel.Ok(rpcCtx, req.taskType, sh.spt, worker)
+		cancel()
+
 		if err != nil {
 			return false, err
 		}