add task ack

server/service.go

@@ -66,6 +66,7 @@ func (n *NodeManagerService) DispatchTask(ctx context.Context, request *omanager
 	dtask := &dispatchTask{
 		task:  request.TaskData,
 		errCh: make(chan error, 1),
+		ack:   make(chan interface{}, 1),
 	}
 
 	worker.taskCh <- dtask

server/workerManager.go

@@ -35,6 +35,7 @@ var (
 
 type dispatchTask struct {
 	task  *odysseus.TaskContent
+	ack   chan interface{}
 	errCh chan error
 }
 
@@ -367,10 +368,24 @@ func (wm *WorkerManager) manageWorker(worker *Worker) error {
 				TaskParam: task.TaskParam,
 			}
 			msg.Message = taskMsg
+			errCh := make(chan error, 1)
+
+			go func() {
+				tm := time.NewTicker(time.Second * 5)
+				defer tm.Stop()
+				select {
+				case <-tm.C:
+					dtask.errCh <- errors.New("ack timeout")
+				case err := <-errCh:
+					dtask.errCh <- err
+				}
+				return
+			}()
+
 			callback = func(err error) bool {
 				if err == nil {
 					// add task to cache.
-					worker.recentTask.Add(task.TaskId, task)
+					worker.recentTask.Add(task.TaskId, dtask)
 					if task.TaskKind == odysseus.TaskKind_ComputeTask {
 						if e := wm.setWorkerLastTaskTime(worker, time.Now().Unix()); e != nil {
 							log.WithField("worker", worker.uuid).WithError(e).Error("set worker last task time failed")
@@ -382,12 +397,7 @@ func (wm *WorkerManager) manageWorker(worker *Worker) error {
 					"worker":      worker.uuid,
 					"worker addr": worker.workerAddr,
 				}).Info("dispatch task to worker")
-
-				select {
-				case dtask.errCh <- err:
-				default:
-					// err ch is invalid
-				}
+				errCh <- err
 				return true
 			}
 		case result := <-worker.resultCh:
@@ -488,6 +498,15 @@ func (wm *WorkerManager) handleWorkerMsg(worker *Worker) {
 				worker.quit <- msg.GoodbyeMessage.Reason
 				close(worker.taskCh)
 				return
+			case *omanager.WorkerMessage_SubmitTaskAck:
+				if v, ok := worker.recentTask.Get(msg.SubmitTaskAck.TaskId); ok {
+					dtask := v.(dispatchTask)
+					select {
+					case dtask.ack <- true:
+						log.WithField("task", dtask.task.TaskId).Debug("got ack")
+					default:
+					}
+				}
 
 			case *omanager.WorkerMessage_SubmitTaskResult:
 				worker.resultCh <- msg.SubmitTaskResult