Only submit result after upstream responsed.

hoveychen · Apr 29, 2024 · 95305a9 · 95305a9
1 parent 46d017d
commit 95305a9
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/pkg/agent/server.go b/pkg/agent/server.go
@@ -275,8 +275,15 @@ func (as *AgentServer) runWorker(ctx context.Context, agentID int, workerNum int
 			pr, pw := io.Pipe()
 
 			grp, ctx := errgroup.WithContext(ctx)
+			recvHeader := make(chan struct{})
 			grp.Go(func() error {
 				defer pr.Close()
+
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				case <-recvHeader:
+				}
 				submitReq := as.newHubAPIRequest(ctx, agentID, hub.PathSubmit, pr)
 				submitReq.Header.Set("slime-connection-id", connectionID)
 				submitResp, err := http.DefaultClient.Do(submitReq)
@@ -311,6 +318,12 @@ func (as *AgentServer) runWorker(ctx context.Context, agentID int, workerNum int
 					"content_length": upResp.ContentLength,
 				}).Info("Upstream responsed")
 
+				select {
+				case recvHeader <- struct{}{}:
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+
 				if err := upResp.Write(pw); err != nil {
 					log.WithError(err).Error("Write upstream response")
 					return err