[WIP] Remove client timeout for ping streams

api/client.go

@@ -68,6 +68,10 @@ type Client struct {
 	// HTTP client used to communicate with the API.
 	client *http.Client
 
+	// HTTP client used for streaming responses. It shares the API client's
+	// transport but has no whole-response timeout.
+	streamingClient *http.Client
+
 	// The logger used
 	logger logger.Logger
 
@@ -87,9 +91,10 @@ func NewClient(l logger.Logger, conf Config) *Client {
 
 	if conf.HTTPClient != nil {
 		return &Client{
-			logger: l,
-			client: conf.HTTPClient,
-			conf:   conf,
+			logger:          l,
+			client:          conf.HTTPClient,
+			streamingClient: withoutTimeout(conf.HTTPClient),
+			conf:            conf,
 		}
 	}
 
@@ -103,14 +108,22 @@ func NewClient(l logger.Logger, conf Config) *Client {
 		clientOptions = append(clientOptions, agenthttp.WithTimeout(conf.Timeout))
 	}
 
+	client := agenthttp.NewClient(clientOptions...)
 	return &Client{
-		logger:         l,
-		client:         agenthttp.NewClient(clientOptions...),
-		conf:           conf,
-		requestHeaders: requestHeadersFromEnv(os.Environ()),
+		logger:          l,
+		client:          client,
+		streamingClient: withoutTimeout(client),
+		conf:            conf,
+		requestHeaders:  requestHeadersFromEnv(os.Environ()),
 	}
 }
 
+func withoutTimeout(client *http.Client) *http.Client {
+	streamingClient := *client
+	streamingClient.Timeout = 0
+	return &streamingClient
+}
+
 func requestHeadersFromEnv(environ []string) http.Header {
 	const prefix = "BUILDKITE_REQUEST_HEADER_"
 	headers := make(http.Header)

api/pings_streaming.go

@@ -23,7 +23,7 @@ func (c *Client) StreamPings(ctx context.Context, agentID string, opts ...connec
 	u.Path = "/"
 
 	cl := agentedgev1connect.NewAgentEdgeServiceClient(
-		c.client,
+		c.streamingClient,
 		u.String(),
 		connect.WithGRPC(),
 		connect.WithClientOptions(opts...),

api/pings_streaming_test.go

@@ -0,0 +1,80 @@
+package api_test
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"connectrpc.com/connect"
+	"github.com/buildkite/agent/v4/api"
+	agentedgev1 "github.com/buildkite/agent/v4/api/proto/gen"
+	"github.com/buildkite/agent/v4/api/proto/gen/agentedgev1connect"
+	"github.com/buildkite/agent/v4/logger"
+)
+
+type delayedPingStreamHandler struct{}
+
+func (delayedPingStreamHandler) StreamPings(_ context.Context, _ *connect.Request[agentedgev1.StreamPingsRequest], stream *connect.ServerStream[agentedgev1.StreamPingsResponse]) error {
+	if err := stream.Send(resumePing()); err != nil {
+		return err
+	}
+	time.Sleep(150 * time.Millisecond)
+	return stream.Send(resumePing())
+}
+
+func TestStreamPingsDoesNotUseAPIClientTimeout(t *testing.T) {
+	path, handler := agentedgev1connect.NewAgentEdgeServiceHandler(delayedPingStreamHandler{})
+	mux := http.NewServeMux()
+	mux.Handle(path, handler)
+	server := httptest.NewServer(mux)
+	t.Cleanup(server.Close)
+
+	client := api.NewClient(logger.Discard, api.Config{
+		Endpoint: server.URL + "/v3",
+		Timeout:  50 * time.Millisecond,
+	})
+	stream, err := client.StreamPings(t.Context(), "agent-id")
+	if err != nil {
+		t.Fatalf("client.StreamPings: %v", err)
+	}
+
+	var messages int
+	for _, err := range stream {
+		if err != nil {
+			t.Fatalf("stream error: %v", err)
+		}
+		messages++
+	}
+	if got, want := messages, 2; got != want {
+		t.Errorf("stream message count = %d, want %d", got, want)
+	}
+}
+
+func TestAPIRequestsStillUseClientTimeout(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		time.Sleep(150 * time.Millisecond)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{}`))
+	}))
+	t.Cleanup(server.Close)
+
+	client := api.NewClient(logger.Discard, api.Config{
+		Endpoint: server.URL + "/",
+		Timeout:  50 * time.Millisecond,
+	})
+	_, _, err := client.Ping(t.Context())
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Errorf("client.Ping error = %v, want context deadline exceeded", err)
+	}
+}
+
+func resumePing() *agentedgev1.StreamPingsResponse {
+	return &agentedgev1.StreamPingsResponse{
+		Action: &agentedgev1.StreamPingsResponse_Resume{
+			Resume: &agentedgev1.ResumeAction{},
+		},
+	}
+}