Merge pull request #4242 from buildkite/backport-pr-4239

[Backport] Describe and retry undecodable API responses

Author: DrJosh9000

agent/agent_worker.go

@@ -458,7 +458,9 @@ func (a *AgentWorker) Heartbeat(ctx context.Context) error {
 	beat, err := roko.DoFunc(ctx, r, func(r *roko.Retrier) (*api.Heartbeat, error) {
 		b, resp, err := a.apiClient.Heartbeat(ctx)
 		if err != nil {
-			if resp != nil && !api.IsRetryableStatus(resp) {
+			// A response we couldn't decode is not the API telling us to stop,
+			// whatever status it arrived with, so check the error as well.
+			if resp != nil && !api.IsRetryableStatus(resp) && !api.IsRetryableError(err) {
 				r.Break()
 				return nil, &errUnrecoverable{action: "Heartbeat", response: resp, err: err}
 			}

agent/agent_worker_ping.go

@@ -223,7 +223,9 @@ func (a *AgentWorker) Ping(ctx context.Context) (jobID, action string, err error
 		// If the ping has a non-retryable status, we have to kill the agent, there's no way of recovering
 		// The reason we do this after the disconnect check is because the backend can (and does) send disconnect actions in
 		// responses with non-retryable statuses
-		if resp != nil && !api.IsRetryableStatus(resp) {
+		// A response we couldn't decode is not the API telling us to stop, whatever
+		// status it arrived with, so check the error as well
+		if resp != nil && !api.IsRetryableStatus(resp) && !api.IsRetryableError(pingErr) {
 			return "", action, &errUnrecoverable{action: "Ping", response: resp, err: pingErr}
 		}
 

api/client.go

@@ -24,6 +24,9 @@ import (
 const (
 	defaultEndpoint  = "https://agent-edge.buildkite.com/v3"
 	defaultUserAgent = "buildkite-agent/api"
+
+	// Maximum bytes of a response body to quote in an error message.
+	maxErrorBodySnippet = 512
 )
 
 // Config is configuration for the API Client
@@ -340,8 +343,11 @@ func (c *Client) doRequest(req *http.Request, v any) (*Response, error) {
 				return response, errors.New("msgpack not supported")
 			}
 
-			if err = json.NewDecoder(resp.Body).Decode(v); err != nil {
-				return response, fmt.Errorf("failed to decode JSON response: %w", err)
+			// Keep the start of the body as it is decoded, so that a decode
+			// failure can report what the response actually contained.
+			head := &headWriter{limit: maxErrorBodySnippet}
+			if err = json.NewDecoder(io.TeeReader(resp.Body, head)).Decode(v); err != nil {
+				return response, newUndecodableResponseError(resp, head.String(), err)
 			}
 		}
 	}
@@ -353,6 +359,11 @@ func (c *Client) doRequest(req *http.Request, v any) (*Response, error) {
 type ErrorResponse struct {
 	Response *http.Response // HTTP response that caused this error
 	Message  string         `json:"message"` // error message
+
+	// Details of a response body that wasn't an API error message. Taken from the
+	// response, never decoded from it, hence `json:"-"`.
+	ContentType string `json:"-"`
+	Snippet     string `json:"-"`
 }
 
 func (r *ErrorResponse) Error() string {
@@ -364,9 +375,92 @@ func (r *ErrorResponse) Error() string {
 		s = fmt.Sprintf("%s: %v", s, r.Message)
 	}
 
+	if r.Snippet != "" {
+		s = fmt.Sprintf("%s: non-JSON body (content-type %q) starting with %q", s, r.ContentType, r.Snippet)
+	}
+
 	return s
 }
 
+// UndecodableResponseError is returned when a response had a success status but
+// a body that could not be decoded. Usually this means something other than the
+// Buildkite Agent API answered the request - a proxy, load balancer, or CDN
+// serving its own error page - which makes it worth retrying. See
+// IsRetryableError.
+type UndecodableResponseError struct {
+	Method      string // request method
+	URL         string // request URL, after any redirects
+	Status      string // response status
+	ContentType string // response content type
+	Snippet     string // start of the response body; empty when it could hold credentials
+	Err         error  // the decoding error
+}
+
+func (e *UndecodableResponseError) Error() string {
+	s := fmt.Sprintf("%s %s: %s: could not decode response body", e.Method, e.URL, e.Status)
+
+	if e.ContentType != "" {
+		s = fmt.Sprintf("%s (content-type %q)", s, e.ContentType)
+	}
+
+	s = fmt.Sprintf("%s: %v", s, e.Err)
+
+	if e.Snippet != "" {
+		s = fmt.Sprintf("%s: body starts with %q", s, e.Snippet)
+	}
+
+	return s
+}
+
+func (e *UndecodableResponseError) Unwrap() error { return e.Err }
+
+func newUndecodableResponseError(r *http.Response, head string, err error) *UndecodableResponseError {
+	e := &UndecodableResponseError{
+		Status:      r.Status,
+		ContentType: r.Header.Get("Content-Type"),
+		Err:         err,
+	}
+
+	if r.Request != nil {
+		e.Method = r.Request.Method
+		e.URL = r.Request.URL.String()
+	}
+
+	// A body that didn't claim to be JSON came from something that isn't the
+	// Agent API, so quoting it leaks none of our credentials. A malformed JSON
+	// body did come from the API and could contain a token, so it stays out.
+	if !isJSONContent(e.ContentType) {
+		// head is capped by the caller; drop a trailing partial rune.
+		e.Snippet = strings.ToValidUTF8(head, "")
+	}
+
+	return e
+}
+
+// headWriter keeps the first limit bytes written to it, and discards the rest.
+type headWriter struct {
+	buf   bytes.Buffer
+	limit int
+}
+
+func (w *headWriter) Write(p []byte) (int, error) {
+	if rem := w.limit - w.buf.Len(); rem > 0 {
+		w.buf.Write(p[:min(rem, len(p))]) //nolint:errcheck // bytes.Buffer.Write never errors.
+	}
+	return len(p), nil
+}
+
+func (w *headWriter) String() string { return w.buf.String() }
+
+// truncate shortens s to at most limit bytes, marking it when it does.
+func truncate(s string, limit int) string {
+	if len(s) <= limit {
+		return s
+	}
+	// Cutting at a byte offset can split a multi-byte rune; drop the remnant.
+	return strings.ToValidUTF8(s[:limit], "") + "…"
+}
+
 func IsErrHavingStatus(err error, code int) bool {
 	var apierr *ErrorResponse
 	return errors.As(err, &apierr) && apierr.Response.StatusCode == code
@@ -377,15 +471,21 @@ func checkResponse(r *http.Response) error {
 		return nil
 	}
 
-	errorResponse := &ErrorResponse{Response: r}
+	contentType := r.Header.Get("Content-Type")
+	errorResponse := &ErrorResponse{Response: r, ContentType: contentType}
 	data, err := io.ReadAll(r.Body)
 	if err != nil {
 		return errorResponse
 	}
-	if data != nil {
+	if len(data) > 0 {
 		// Unmarshaling the error JSON is best-effort, but we could consider
 		// reporting unmarshaling problems.
-		json.Unmarshal(data, errorResponse) //nolint:errcheck // ^^
+		if err := json.Unmarshal(data, errorResponse); err != nil && !isJSONContent(contentType) {
+			// Not an API error message: most likely an error page from a proxy,
+			// load balancer, or CDN in front of the API. It's the only clue as to
+			// what really answered, and holds none of our credentials.
+			errorResponse.Snippet = truncate(string(data), maxErrorBodySnippet)
+		}
 	}
 
 	return errorResponse

api/retryable.go

@@ -38,9 +38,24 @@ func IsRetryableStatus(r *Response) bool {
 	}
 }
 
+// isUndecodableResponse reports whether err is (or wraps) an
+// UndecodableResponseError.
+func isUndecodableResponse(err error) bool {
+	var undecodable *UndecodableResponseError
+	return errors.As(err, &undecodable)
+}
+
 // Looks at a bunch of connection related errors, and returns true if the error
 // matches one of them.
 func IsRetryableError(err error) bool {
+	// A response we couldn't decode is usually an error page from a proxy, load
+	// balancer, or CDN in front of the API, and it can arrive with any status,
+	// including a 2xx. Either way the API's real answer is lost, so the only way
+	// to get it is to ask again.
+	if isUndecodableResponse(err) {
+		return true
+	}
+
 	var neterr net.Error
 	if errors.As(err, &neterr) {
 		if neterr.Timeout() {
@@ -70,10 +85,11 @@ func IsRetryableError(err error) bool {
 }
 
 // BreakOnNonRetryable calls r.Break() if the error from an API call is not
-// worth retrying. An error is retryable if the response has a retryable status
-// code (429, 5xx) or if there was no response and the error is a retryable
-// network-level error (connection reset, timeout, etc.). All other errors
-// — including all non-429 4xx status codes — cause a break.
+// worth retrying. An error is retryable if the response body could not be
+// decoded, if the response has a retryable status code (429, 5xx), or if there
+// was no response and the error is a retryable network-level error (connection
+// reset, timeout, etc.). All other errors — including all non-429 4xx status
+// codes — cause a break.
 //
 // This should be called inside roko retry callbacks after every API call.
 // If err is nil, this is a no-op.
@@ -84,6 +100,11 @@ func BreakOnNonRetryable(r *roko.Retrier, resp *Response, err error) (broke bool
 	if err == nil {
 		return false
 	}
+	// An undecodable response is retryable whatever status it arrived with, so
+	// classify the error before the response.
+	if isUndecodableResponse(err) {
+		return false
+	}
 	if resp != nil {
 		if !IsRetryableStatus(resp) {
 			r.Break()