Hammer pushback support

hammer/client.go

@@ -17,6 +17,7 @@ package main
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -25,11 +26,14 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/transparency-dev/trillian-tessera/client"
 	"k8s.io/klog/v2"
 )
 
+var ErrRetry = errors.New("retry")
+
 // newLogClientsFromFlags returns a fetcher and a writer that will read
 // and write leaves to all logs in the `log_url` flag set.
 func newLogClientsFromFlags() (*roundRobinFetcher, *roundRobinLeafWriter) {
@@ -155,12 +159,20 @@ func (w httpLeafWriter) Write(ctx context.Context, newLeaf []byte) (uint64, erro
 	if err != nil {
 		return 0, fmt.Errorf("failed to read body: %v", err)
 	}
-	if resp.StatusCode != http.StatusOK {
+	switch resp.StatusCode {
+	case http.StatusOK:
+		if resp.Request.Method != http.MethodPost {
+			return 0, fmt.Errorf("write leaf was redirected to %s", resp.Request.URL)
+		}
+		// Continue below
+	case http.StatusServiceUnavailable, http.StatusBadGateway, http.StatusGatewayTimeout:
+		// These status codes may indicate a delay before retrying, so handle that here:
+		time.Sleep(retryDelay(resp.Header.Get("RetryAfter"), time.Second))
+
+		return 0, fmt.Errorf("log not available. Status code: %d. Body: %q %w", resp.StatusCode, body, ErrRetry)
+	default:
 		return 0, fmt.Errorf("write leaf was not OK. Status code: %d. Body: %q", resp.StatusCode, body)
 	}
-	if resp.Request.Method != http.MethodPost {
-		return 0, fmt.Errorf("write leaf was redirected to %s", resp.Request.URL)
-	}
 	parts := bytes.Split(body, []byte("\n"))
 	index, err := strconv.ParseUint(string(parts[0]), 10, 64)
 	if err != nil {
@@ -169,6 +181,21 @@ func (w httpLeafWriter) Write(ctx context.Context, newLeaf []byte) (uint64, erro
 	return index, nil
 }
 
+func retryDelay(retryAfter string, defaultDur time.Duration) time.Duration {
+	if retryAfter == "" {
+		return defaultDur
+	}
+	d, err := time.Parse(http.TimeFormat, retryAfter)
+	if err == nil {
+		return time.Until(d)
+	}
+	s, err := strconv.Atoi(retryAfter)
+	if err == nil {
+		return time.Duration(s) * time.Second
+	}
+	return defaultDur
+}
+
 // roundRobinLeafWriter ensures that write requests are sent to all configured
 // LeafWriters using a round-robin strategy.
 type roundRobinLeafWriter struct {

hammer/hammer.go

@@ -156,11 +156,22 @@ func (h *Hammer) Run(ctx context.Context) {
 
 	// Set up logging for any errors
 	go func() {
+		tick := time.NewTicker(time.Second)
+		pbCount := 0
 		for {
 			select {
 			case <-ctx.Done(): //context cancelled
 				return
+			case <-tick.C:
+				if pbCount > 0 {
+					klog.Warningf("%d requests received pushback from log", pbCount)
+					pbCount = 0
+				}
 			case err := <-h.errChan:
+				if errors.Is(err, ErrRetry) {
+					pbCount++
+					continue
+				}
 				klog.Warning(err)
 			}
 		}

hammer/workers.go

@@ -201,7 +201,7 @@ func (w *LogWriter) Run(ctx context.Context) {
 		lt := leafTime{queuedAt: time.Now()}
 		index, err := w.writer(ctx, newLeaf)
 		if err != nil {
-			w.errChan <- fmt.Errorf("failed to create request: %v", err)
+			w.errChan <- fmt.Errorf("failed to create request: %w", err)
 			continue
 		}
 		lt.idx, lt.assignedAt = index, time.Now()