monitor: Retry get-sth and proof fetching

Should ensure we don't get into a position where we always fail to get
3x queries that succeed in a row when trying to persist chunks.

1 files changed, 80 insertions, 3 deletions

diff --git a/internal/monitor/backoff.go b/internal/monitor/backoff.go
index 63c5f55..3bfff7e 100644
--- a/internal/monitor/backoff.go
+++ b/internal/monitor/backoff.go
@@ -2,6 +2,7 @@ package monitor
 
 import (
 	"context"
+	"time"
 
 	ct "github.com/google/certificate-transparency-go"
 	"github.com/google/certificate-transparency-go/client"
@@ -15,6 +16,10 @@ import (
 //
 // For reference on this issue, see:
 // https://github.com/google/certificate-transparency-go/issues/898
+//
+// Update: retries was added for get-sth and proof fetching.
+// Only because we need 3x queries that succeed in a row, and some logs seem to
+// rate limit globally for all endpoints rather than per endpoint.
 type backoffClient struct {
 	cli *client.LogClient
 }
@@ -24,15 +29,75 @@ func (bc *backoffClient) BaseURI() string {
 }
 
 func (bc *backoffClient) GetSTH(ctx context.Context) (*ct.SignedTreeHead, error) {
-	return bc.cli.GetSTH(ctx)
+	backoff := 1
+	for {
+		rsp, err := bc.cli.GetSTH(ctx)
+		if err == nil {
+			return rsp, nil
+		}
+
+		jcErr, ok := err.(jsonclient.RspError)
+		if !ok {
+			return rsp, err
+		}
+		if jcErr.StatusCode != 429 {
+			return rsp, err
+		}
+		if err := sleep(ctx, time.Duration(backoff)*time.Second); err != nil {
+			return nil, err
+		}
+		if backoff < 32 {
+			backoff = 2 * backoff
+		}
+	}
 }
 
 func (bc *backoffClient) GetSTHConsistency(ctx context.Context, first, second uint64) ([][]byte, error) {
-	return bc.cli.GetSTHConsistency(ctx, first, second)
+	backoff := 1
+	for {
+		rsp, err := bc.cli.GetSTHConsistency(ctx, first, second)
+		if err == nil {
+			return rsp, nil
+		}
+
+		jcErr, ok := err.(jsonclient.RspError)
+		if !ok {
+			return rsp, err
+		}
+		if jcErr.StatusCode != 429 {
+			return rsp, err
+		}
+		if err := sleep(ctx, time.Duration(backoff)*time.Second); err != nil {
+			return nil, err
+		}
+		if backoff < 32 {
+			backoff = 2 * backoff
+		}
+	}
 }
 
 func (bc *backoffClient) GetProofByHash(ctx context.Context, hash []byte, treeSize uint64) (*ct.GetProofByHashResponse, error) {
-	return bc.cli.GetProofByHash(ctx, hash, treeSize)
+	backoff := 1
+	for {
+		rsp, err := bc.cli.GetProofByHash(ctx, hash, treeSize)
+		if err == nil {
+			return rsp, nil
+		}
+
+		jcErr, ok := err.(jsonclient.RspError)
+		if !ok {
+			return rsp, err
+		}
+		if jcErr.StatusCode != 429 {
+			return rsp, err
+		}
+		if err := sleep(ctx, time.Duration(backoff)*time.Second); err != nil {
+			return nil, err
+		}
+		if backoff < 32 {
+			backoff = 2 * backoff
+		}
+	}
 }
 
 func (bc *backoffClient) GetRawEntries(ctx context.Context, start, end int64) (*ct.GetEntriesResponse, error) {
@@ -54,3 +119,15 @@ func (bc *backoffClient) GetRawEntries(ctx context.Context, start, end int64) (*
 	}
 	return rsp, err
 }
+
+func sleep(ctx context.Context, d time.Duration) error {
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}