From fd3bc4f1cff1fd1dc55ebb5a647ac5b081da4937 Mon Sep 17 00:00:00 2001
From: Neha Gupta <nehagupta161995@gmail.com>
Date: Fri, 3 Feb 2023 16:27:07 +0530
Subject: [PATCH] retry exponential in case of get page failure

---
 src/go.mod                               |  1 +
 src/go.sum                               |  2 +
 src/internal/connector/onedrive/drive.go | 58 ++++++++++++------------
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/src/go.mod b/src/go.mod
index b9cb5ec53..05b801f9f 100644
--- a/src/go.mod
+++ b/src/go.mod
@@ -36,6 +36,7 @@ require (
 	github.com/VividCortex/ewma v1.2.0 // indirect
 	github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d // indirect
 	github.com/andybalholm/brotli v1.0.4 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.0 // indirect
 	github.com/dnaeon/go-vcr v1.2.0 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
diff --git a/src/go.sum b/src/go.sum
index b8f926e79..7813e532d 100644
--- a/src/go.sum
+++ b/src/go.sum
@@ -73,6 +73,8 @@ github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4YnC6+E63dPcxHo2sUxDIu8g3QgEJdRY=
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
+github.com/cenkalti/backoff/v4 v4.2.0 h1:HN5dHm3WBOgndBH6E8V0q2jIYIR3s9yglV8k/+MN3u4=
+github.com/cenkalti/backoff/v4 v4.2.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE=
diff --git a/src/internal/connector/onedrive/drive.go b/src/internal/connector/onedrive/drive.go
index 6270ec08a..9680812b8 100644
--- a/src/internal/connector/onedrive/drive.go
+++ b/src/internal/connector/onedrive/drive.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"strings"
-	"time"
 
 	msdrive "github.com/microsoftgraph/msgraph-sdk-go/drive"
 	msdrives "github.com/microsoftgraph/msgraph-sdk-go/drives"
@@ -18,6 +17,7 @@ import (
 	"github.com/alcionai/corso/src/internal/connector/onedrive/api"
 	"github.com/alcionai/corso/src/internal/connector/support"
 	"github.com/alcionai/corso/src/pkg/logger"
+	backoff "github.com/cenkalti/backoff/v4"
 )
 
 var errFolderNotFound = errors.New("folder not found")
@@ -70,41 +70,41 @@ func drives(
 		drives          = []models.Driveable{}
 	)
 
-	if !retry {
-		numberOfRetries = 0
+	getPageOperation := func() error {
+		page, err = pager.GetPage(ctx)
+		if err != nil {
+			// Various error handling. May return an error or perform a retry.
+			detailedError := support.ConnectorStackErrorTrace(err)
+			if strings.Contains(detailedError, userMysiteURLNotFound) ||
+				strings.Contains(detailedError, userMysiteNotFound) {
+				logger.Ctx(ctx).Infof("resource owner does not have a drive")
+				drives = make([]models.Driveable, 0)
+				return nil // no license or drives.
+			}
+			return errors.Wrapf(
+				err,
+				"failed to retrieve drives. details: %s",
+				detailedError,
+			)
+		}
+		return err
 	}
 
 	// Loop through all pages returned by Graph API.
 	for {
-		// Retry Loop for Drive retrieval. Request can timeout
-		for i := 0; i <= numberOfRetries; i++ {
-			page, err = pager.GetPage(ctx)
+
+		if !retry {
+			err = getPageOperation()
 			if err != nil {
-				// Various error handling. May return an error or perform a retry.
-				detailedError := support.ConnectorStackErrorTrace(err)
-				if strings.Contains(detailedError, userMysiteURLNotFound) ||
-					strings.Contains(detailedError, userMysiteNotFound) {
-					logger.Ctx(ctx).Infof("resource owner does not have a drive")
-					return make([]models.Driveable, 0), nil // no license or drives.
-				}
-
-				if strings.Contains(detailedError, contextDeadlineExceeded) && i < numberOfRetries {
-					time.Sleep(time.Duration(3*(i+1)) * time.Second)
-					continue
-				}
-
-				return nil, errors.Wrapf(
-					err,
-					"failed to retrieve drives. details: %s",
-					detailedError,
-				)
+				return nil, errors.Wrap(err, "extracting drives from response")
+			}
+		} else {
+			exponentialBackoff := backoff.WithMaxRetries(backoff.NewExponentialBackOff(), uint64(numberOfRetries))
+			err = backoff.Retry(getPageOperation, exponentialBackoff)
+			if err != nil {
+				return nil, errors.Wrap(err, "extracting drives from response")
 			}
-
-			// No error encountered, break the retry loop so we can extract results
-			// and see if there's another page to fetch.
-			break
 		}
-
 		tmp, err := pager.ValuesIn(page)
 		if err != nil {
 			return nil, errors.Wrap(err, "extracting drives from response")