Merge pull request #2160 from buildkite/api-client-header-buildkite-t…

…imeout-milliseconds

`Buildkite-Timeout-Milliseconds` request header

api/api_internal_test.go

@@ -0,0 +1,43 @@
+package api
+
+import (
+	"context"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/buildkite/agent/v3/logger"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewRequestBuildkiteTimeoutMilliseconds(t *testing.T) {
+	c := NewClient(logger.NewBuffer(), Config{})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	r, err := c.newRequest(ctx, "GET", "/foo", nil)
+	assert.NoError(t, err)
+
+	value := r.Header.Get("Buildkite-Timeout-Milliseconds")
+	ms, err := strconv.ParseInt(value, 10, 64)
+	assert.NoError(t, err)
+
+	// Allow a generous 1000ms between setting the timeout and the header being set.
+	if ms <= 9_000 || ms > 10_000 {
+		t.Errorf("Expected Buildkite-Timeout-Milliseconds to reflect 10 second timeout, got %q (%d ms)", value, ms)
+	}
+}
+
+func TestNewRequestWithoutBuildkiteTimeoutMilliseconds(t *testing.T) {
+	c := NewClient(logger.NewBuffer(), Config{})
+
+	ctx := context.Background() // no timeout/deadline
+
+	r, err := c.newRequest(ctx, "GET", "/foo", nil)
+	assert.NoError(t, err)
+
+	if value, ok := r.Header["Buildkite-Timeout-Milliseconds"]; ok {
+		t.Errorf("Expected no Buildkite-Timeout-Milliseconds header, got %q", value)
+	}
+}

api/client.go

@@ -15,6 +15,7 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"reflect"
+	"strconv"
 	"strings"
 	"time"
 
@@ -172,6 +173,15 @@ func (c *Client) newRequest(
 
 	req.Header.Add("User-Agent", c.conf.UserAgent)
 
+	// If our context has a timeout/deadline, tell the server how long is remaining.
+	// This may allow the server to configure its own timeouts accordingly.
+	if deadline, ok := ctx.Deadline(); ok {
+		ms := time.Until(deadline).Milliseconds()
+		if ms > 0 {
+			req.Header.Add("Buildkite-Timeout-Milliseconds", strconv.FormatInt(ms, 10))
+		}
+	}
+
 	for _, header := range headers {
 		req.Header.Add(header.Name, header.Value)
 	}