No public description

internal/cli/custom/envruntime/client.go

@@ -1134,6 +1134,7 @@ type HeartbeatStream struct {
 	hbCancel  context.CancelFunc
 	sendMu    sync.Mutex
 	ackSeq    atomic.Int64
+	closing   atomic.Bool
 	stoppedWg sync.WaitGroup
 }
 
@@ -1160,6 +1161,10 @@ func NewHeartbeatStream(ctx context.Context, inner InteractionStream, interval t
 				return
 			case <-ticker.C:
 				hs.sendMu.Lock()
+				if hs.closing.Load() {
+					hs.sendMu.Unlock()
+					return
+				}
 				err := SendHeartbeat(hbCtx, hs.inner, hs.ackSeq.Load())
 				hs.sendMu.Unlock()
 				if err != nil {
@@ -1197,15 +1202,18 @@ func (s *HeartbeatStream) Recv(ctx context.Context) (*ServerEventBatch, error) {
 
 // CloseSend stops the background heartbeat ticker and half-closes the underlying stream.
 func (s *HeartbeatStream) CloseSend() error {
+	s.closing.Store(true)
+	s.sendMu.Lock()
 	s.hbCancel()
+	err := s.inner.CloseSend()
+	s.sendMu.Unlock()
 	s.stoppedWg.Wait()
-	s.sendMu.Lock()
-	defer s.sendMu.Unlock()
-	return s.inner.CloseSend()
+	return err
 }
 
 // Close stops the background heartbeat ticker and closes the underlying stream.
 func (s *HeartbeatStream) Close() error {
+	s.closing.Store(true)
 	s.hbCancel()
 	err := s.inner.Close()
 	s.stoppedWg.Wait()

internal/cli/custom/envruntime/client_test.go

@@ -666,6 +666,27 @@ func TestHTTPClientLeaseHeartbeatPreviewAndStream(t *testing.T) {
 			t.Errorf("HeartbeatStream (useClose=%v) continued sending after stop: before=%d after=%d", useClose, afterStopCount, sendCalls.Load())
 		}
 	}
+	reqPr, reqPw := io.Pipe()
+	respPr, respPw := io.Pipe()
+	hbHalf := NewHeartbeatStream(t.Context(), &grpcInteractionStream{resp: &http.Response{Body: respPr}, pw: reqPw}, 5*time.Millisecond, 7)
+	closeSendErr := make(chan error, 1)
+	go func() {
+		_, _ = reqPr.Read(make([]byte, 1))
+		time.Sleep(12 * time.Millisecond)
+		closeSendErr <- hbHalf.CloseSend()
+	}()
+	for !hbHalf.closing.Load() {
+		time.Sleep(time.Millisecond)
+	}
+	time.Sleep(2 * time.Millisecond)
+	if _, err := io.ReadFull(reqPr, make([]byte, 6)); err != nil {
+		t.Fatalf("ReadFull(in-flight heartbeat remainder) = %v", err)
+	}
+	if n, err := reqPr.Read(make([]byte, 1)); n != 0 || !errors.Is(err, io.EOF) || <-closeSendErr != nil {
+		t.Errorf("reqPr.Read after CloseSend() = (%d, %v), want (0, io.EOF)", n, err)
+	}
+	_ = respPw.Close()
+	_ = hbHalf.Close()
 
 	var noDebugBuf bytes.Buffer
 	noDebugClient := NewHTTPClient(HTTPClientConfig{