From cf584e536a8cd1a22d9a2131543cff6143b3d891 Mon Sep 17 00:00:00 2001
From: Marsel Mavletkulov <mmavletkulov@maxmind.com>
Date: Thu, 25 Jan 2024 12:59:14 -0500
Subject: [PATCH] Retry HTTP2 INTERNAL_ERRORs

---
 go.mod                           |  2 +
 go.sum                           |  4 ++
 pkg/geoipupdate/geoip_updater.go | 70 +++++++++++++++++++++++++++-----
 3 files changed, 66 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index b587bb85..9b521727 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/gofrs/flock v0.8.1
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.8.4
+	golang.org/x/net v0.20.0
 	golang.org/x/sync v0.6.0
 )
 
@@ -16,6 +17,7 @@ require (
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/sys v0.16.0 // indirect
+	golang.org/x/text v0.14.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/go.sum b/go.sum
index 5648f23b..5dc82478 100644
--- a/go.sum
+++ b/go.sum
@@ -17,10 +17,14 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo=
+golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
 golang.org/x/sync v0.6.0 h1:5BMeUDZ7vkXGfEr1x9B4bRcTH4lpkTkpdh0T/J+qjbQ=
 golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
 golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
+golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b h1:QRR6H1YWRnHb4Y/HeNFCTJLFVxaq6wH4YuVdsUOr75U=
 gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
diff --git a/pkg/geoipupdate/geoip_updater.go b/pkg/geoipupdate/geoip_updater.go
index 24847709..76b7f8da 100644
--- a/pkg/geoipupdate/geoip_updater.go
+++ b/pkg/geoipupdate/geoip_updater.go
@@ -5,12 +5,16 @@ package geoipupdate
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log"
 	"os"
 	"sync"
 	"time"
 
+	"github.com/cenkalti/backoff/v4"
+	"golang.org/x/net/http2"
+
 	"github.com/maxmind/geoipupdate/v6/pkg/geoipupdate/database"
 	"github.com/maxmind/geoipupdate/v6/pkg/geoipupdate/internal"
 )
@@ -85,20 +89,11 @@ func (c *Client) Run(ctx context.Context) error {
 	for _, editionID := range c.config.EditionIDs {
 		editionID := editionID
 		processFunc := func(ctx context.Context) error {
-			editionHash, err := writer.GetHash(editionID)
+			edition, err := c.downloadEdition(ctx, editionID, reader, writer)
 			if err != nil {
 				return err
 			}
 
-			edition, err := reader.Read(ctx, editionID, editionHash)
-			if err != nil {
-				return err
-			}
-
-			if err := writer.Write(edition); err != nil {
-				return err
-			}
-
 			edition.CheckedAt = time.Now().In(time.UTC)
 
 			mu.Lock()
@@ -126,3 +121,58 @@ func (c *Client) Run(ctx context.Context) error {
 
 	return nil
 }
+
+// downloadEdition downloads the file with retries on HTTP2 INTERNAL_ERRORs.
+func (c *Client) downloadEdition(
+	ctx context.Context,
+	editionID string,
+	r database.Reader,
+	w database.Writer,
+) (*database.ReadResult, error) {
+	editionHash, err := w.GetHash(editionID)
+	if err != nil {
+		return nil, err
+	}
+
+	// RetryFor value of 0 means that no retries should be performed.
+	// Max zero retries has to be set to achieve that
+	// because the backoff never stops if MaxElapsedTime is zero.
+	exp := backoff.NewExponentialBackOff()
+	exp.MaxElapsedTime = c.config.RetryFor
+	b := backoff.BackOff(exp)
+	if exp.MaxElapsedTime == 0 {
+		b = backoff.WithMaxRetries(exp, 0)
+	}
+
+	var edition *database.ReadResult
+	err = backoff.RetryNotify(
+		func() error {
+			edition, err = r.Read(ctx, editionID, editionHash)
+			if err != nil {
+				return backoff.Permanent(err)
+			}
+
+			if err = w.Write(edition); err != nil {
+				streamErr := http2.StreamError{}
+				if errors.As(err, &streamErr) && streamErr.Code.String() == "INTERNAL_ERROR" {
+					return err
+				}
+
+				return backoff.Permanent(err)
+			}
+
+			return nil
+		},
+		b,
+		func(err error, d time.Duration) {
+			if c.config.Verbose {
+				log.Printf("Couldn't download %s, retrying in %v: %v", editionID, d, err)
+			}
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return edition, nil
+}