Fix some race conditions when running very fast tasks (#108)

Author: eagleonhill

pkg/agent/report_status_plugin.go

@@ -4,7 +4,7 @@
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
-//     http://www.apache.org/licenses/LICENSE-2.0
+//	http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
@@ -39,22 +39,28 @@ func NewReportStatusPlugin(sshdPlugin any) *ReportStatusPlugin {
 
 func (p *ReportStatusPlugin) Run(ctx context.Context) error {
 	sshd := ctx.Value(p.sshdPlugin).(*SSHD)
+	send := func() {
+		addr := sshd.listener.Addr()
+		port := addr.(*net.TCPAddr).Port
+		sandboxConnection := agentd.SandboxConnection{
+			Port:    port,
+			HostKey: sshd.HostPublicKey.Marshal(),
+		}
+		err := gob.NewEncoder(os.Stdout).Encode(sandboxConnection)
+		log.Printf("Sent sandbox connection info")
+		if err != nil {
+			log.Fatal("Failed to encode sandbox connection:", err)
+		}
+	}
+	// Handle future CHECK_POINT recovers.
 	go func() {
 		for {
-			addr := sshd.listener.Addr()
-			port := addr.(*net.TCPAddr).Port
-			sandboxConnection := agentd.SandboxConnection{
-				Port:    port,
-				HostKey: sshd.HostPublicKey.Marshal(),
-			}
-			err := gob.NewEncoder(os.Stdout).Encode(sandboxConnection)
-			if err != nil {
-				log.Fatal("Failed to encode sandbox connection:", err)
-			}
 			sig := make(chan os.Signal, 1)
 			signal.Notify(sig, syscall.SIGCONT)
 			<-sig
+			send()
 		}
 	}()
+	send()
 	return p.RunNext(ctx)
 }

pkg/agent/sshd.go

@@ -144,7 +144,6 @@ func (s *SSHD) Start() (net.Addr, error) {
 	}
 
 	s.listener = listener
-	go s.waiter.Run()
 	go func() {
 		for {
 			conn, err := listener.Accept()

pkg/agent/wait.go

@@ -4,7 +4,7 @@
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
-//     http://www.apache.org/licenses/LICENSE-2.0
+//	http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,

pkg/broker/agent.go

@@ -283,6 +283,9 @@ func (a *Agent) Run(
 				response.SessionRequest = req.Request
 				if err := stream.Send(response); err != nil {
 					log.Printf("Failed to send session request %s to agent %s", req.Request.SessionId, a.id)
+					a.delegate.RevokeAccessToAgent(stream.Context(), a, req)
+					delete(a.mySessions, req.Key())
+					// If this error happens, the stream will likely be broken. Not registering back.
 					req.agentResponse <- AgentSessionResponse{
 						Error: err,
 					}
@@ -314,7 +317,7 @@ func (a *Agent) RequestKill(ctx context.Context, instanceId int64, sessionId str
 func (a *Agent) handleSessionUpdate(ctx context.Context, resp *proto.AgentUpdateRequest) error {
 	// TODO: Handle multi-error.
 	if resp.SessionInitResponse != nil {
-		if err := a.handleSessionInitResponse(resp.SessionInitResponse); err != nil {
+		if err := a.handleSessionInitResponse(ctx, resp.SessionInitResponse); err != nil {
 			return err
 		}
 	}
@@ -326,7 +329,7 @@ func (a *Agent) handleSessionUpdate(ctx context.Context, resp *proto.AgentUpdate
 	return nil
 }
 
-func (a *Agent) handleSessionInitResponse(resp *proto.SessionInitResponse) error {
+func (a *Agent) handleSessionInitResponse(ctx context.Context, resp *proto.SessionInitResponse) error {
 	sessionKey := SessionKey{
 		InstanceId: resp.InstanceId,
 		SessionId:  resp.SessionId,
@@ -342,6 +345,7 @@ func (a *Agent) handleSessionInitResponse(resp *proto.SessionInitResponse) error
 			Error: fmt.Errorf("%w: %s", ErrorFailedToStartSession, a.id),
 		}
 		session.Complete(proto.SessionExecutionFinalState_SESSION_EXECUTION_FINAL_STATE_STARTUP_FAILURE)
+		a.delegate.RevokeAccessToAgent(ctx, a, session)
 		a.scheduler.AddAgent(a)
 		a.scheduler.PoolManager.MarkIdle(a.id, a.slots-len(a.mySessions))
 		a.receiving = true

pkg/broker/broker.go

@@ -93,7 +93,7 @@ func (s *server) AgentUpdate(stream proto.BrokerService_AgentUpdateServer) error
 		}
 		update := func(sess *Session) error {
 			agent.mySessions[key] = sess
-			agent.handleSessionInitResponse(session.Response)
+			agent.handleSessionInitResponse(stream.Context(), session.Response)
 			sess.Reconnect()
 			return nil
 		}

pkg/tasks/task_service.go

@@ -205,7 +205,7 @@ func (s *TaskServiceServer) WatchTask(in *proto.GetTaskRequest, stream proto.Tas
 	if err != nil {
 		return err
 	}
-	if isTaskStatusFinal(task.Status) {
+	if IsTaskStatusFinal(task.Status) {
 		return nil
 	}
 
@@ -224,15 +224,15 @@ func (s *TaskServiceServer) WatchTask(in *proto.GetTaskRequest, stream proto.Tas
 			// Client cancel
 			return ctx.Err()
 		case status := <-statusChan:
-			if (isTaskStatusFinal(status) || status == proto.TaskStatus_TASK_STATUS_RUNNING_SUBTASKS) && sendErr != nil && leased {
+			if (IsTaskStatusFinal(status) || status == proto.TaskStatus_TASK_STATUS_RUNNING_SUBTASKS) && sendErr != nil && leased {
 				// Task execution finished, stop watching live status, ensure updates are processed.
 				err := <-sendErr
 				if err != nil && !errors.Is(err, io.EOF) {
 					return err
 				}
 				sendErr = nil
 			}
-			if isTaskStatusFinal(status) {
+			if IsTaskStatusFinal(status) {
 				// Get full task from db
 				task, err := s.db.GetTask(ctx, in.TaskId)
 				if err != nil {
@@ -348,6 +348,6 @@ func (s *TaskServiceServer) patchTaskStatus(ctx context.Context, tasks ...*proto
 	return nil
 }
 
-func isTaskStatusFinal(status proto.TaskStatus) bool {
+func IsTaskStatusFinal(status proto.TaskStatus) bool {
 	return status == proto.TaskStatus_TASK_STATUS_SUCCESS || status == proto.TaskStatus_TASK_STATUS_FAILURE || status == proto.TaskStatus_TASK_STATUS_CANCELLED || status == proto.TaskStatus_TASK_STATUS_FAILED_UPSTREAM
 }