Exporter: Add retries for Search, ReadContext and Import operations when importing the resource (#3202)

* Exporter: Add retries for Search/ReadContext and Import operations when importing the resource

Under a very high load, Databricks backend may not answer on time, or return specific
errors, so it makes sense to retry operation few times.

This PR uses "naive" implementation, I need to play a bit more with `retries` package
before adopting it

* Review comments

* Code refactoring, reworking retries for directory listings as well

* Adjust logging levels

Author: alexott

exporter/context.go

@@ -758,7 +758,7 @@ func (ic *importContext) generateAndWriteResources(sh *os.File) {
 	for i, r := range resources {
 		ic.waitGroup.Add(1)
 		resourcesChan <- r
-		if i%50 == 0 {
+		if i%500 == 0 {
 			log.Printf("[INFO] Submitted %d of %d resources", i+1, scopeSize)
 		}
 	}

exporter/context_test.go

@@ -201,6 +201,50 @@ func TestEmitNoSearchNoId(t *testing.T) {
 	close(ch)
 }
 
+func TestEmitNoSearchNoIdWithRetry(t *testing.T) {
+	ch := make(resourceChannel, 10)
+	state := newStateApproximation([]string{"a"})
+	i := 0
+	ic := &importContext{
+		importing: map[string]bool{},
+		Resources: map[string]*schema.Resource{
+			"a": {},
+		},
+		Importables: map[string]importable{
+			"a": {
+				Service: "e",
+				Search: func(ic *importContext, r *resource) error {
+					if i > 0 {
+						return nil
+					}
+					i = i + 1
+					return fmt.Errorf("context deadline exceeded (Client.Timeout exceeded while awaiting headers)")
+				},
+			},
+		},
+		waitGroup: &sync.WaitGroup{},
+		channels: map[string]resourceChannel{
+			"a": ch,
+		},
+		ignoredResources: map[string]struct{}{},
+		State:            state,
+	}
+	ic.enableServices("e")
+	go func() {
+		for r := range ch {
+			r.ImportResource(ic)
+		}
+	}()
+	ic.Emit(&resource{
+		Resource:  "a",
+		Attribute: "b",
+		Value:     "d",
+		Name:      "c",
+	})
+	ic.waitGroup.Wait()
+	close(ch)
+}
+
 func TestEmitNoSearchSucceedsImportFails(t *testing.T) {
 	ch := make(resourceChannel, 10)
 	state := newStateApproximation([]string{"a"})

exporter/model.go

@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/terraform-provider-databricks/common"
 
 	"github.com/hashicorp/hcl/v2/hclwrite"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 )
@@ -265,8 +266,12 @@ func (r *resource) ImportResource(ic *importContext) {
 			log.Printf("[ERROR] Searching %s is not available", r)
 			return
 		}
-		if err := ir.Search(ic, r); err != nil {
-			log.Printf("[ERROR] Cannot search for a resource %s: %v", err, r)
+		err := runWithRetries(func() error {
+			return ir.Search(ic, r)
+		},
+			fmt.Sprintf("searching of %v", r))
+		if err != nil {
+			log.Printf("[ERROR] Error searching %s#%s: %v", r.Resource, r.ID, err)
 			return
 		}
 		if r.ID == "" {
@@ -288,7 +293,11 @@ func (r *resource) ImportResource(ic *importContext) {
 		if apiVersion != "" {
 			ctx = context.WithValue(ctx, common.Api, apiVersion)
 		}
-		if dia := pr.ReadContext(ctx, r.Data, ic.Client); dia != nil {
+		dia := runWithRetries(func() diag.Diagnostics {
+			return pr.ReadContext(ctx, r.Data, ic.Client)
+		},
+			fmt.Sprintf("reading %s#%s", r.Resource, r.ID))
+		if dia != nil {
 			log.Printf("[ERROR] Error reading %s#%s: %v", r.Resource, r.ID, dia)
 			return
 		}
@@ -298,7 +307,11 @@ func (r *resource) ImportResource(ic *importContext) {
 	}
 	r.Name = ic.ResourceName(r)
 	if ir.Import != nil {
-		if err := ir.Import(ic, r); err != nil {
+		err := runWithRetries(func() error {
+			return ir.Import(ic, r)
+		},
+			fmt.Sprintf("importing of %s#%s", r.Resource, r.ID))
+		if err != nil {
 			log.Printf("[ERROR] Failed custom import of %s: %s", r, err)
 			return
 		}

exporter/util.go

@@ -1088,7 +1088,6 @@ type directoryInfo struct {
 
 // constants related to the parallel listing
 const (
-	directoryListingMaxAttempts = 3
 	envVarListParallelism       = "EXPORTER_WS_LIST_PARALLELISM"
 	envVarDirectoryChannelSize  = "EXPORTER_DIRECTORIES_CHANNEL_SIZE"
 	defaultWorkersPoolSize      = 10
@@ -1101,8 +1100,11 @@ func recursiveAddPathsParallel(a workspace.NotebooksAPI, directory directoryInfo
 	notebookInfoList, err := a.ListInternalImpl(directory.Path)
 	if err != nil {
 		log.Printf("[WARN] error listing '%s': %v", directory.Path, err)
-		if directory.Attempts < directoryListingMaxAttempts {
+		if isRetryableError(err.Error(), directory.Attempts) {
 			wg.Add(1)
+			log.Printf("[INFO] attempt %d of retrying listing of '%s' after error: %v",
+				directory.Attempts+1, directory.Path, err)
+			time.Sleep(time.Duration(retryDelaySeconds) * time.Second)
 			dirChannel <- directoryInfo{Path: directory.Path, Attempts: directory.Attempts + 1}
 		}
 	}
@@ -1166,3 +1168,40 @@ func ListParallel(a workspace.NotebooksAPI, path string, shouldIncludeDir func(w
 	defer answer.MU.Unlock()
 	return answer.data, nil
 }
+
+var (
+	maxRetries        = 5
+	retryDelaySeconds = 2
+	retriableErrors   = []string{"context deadline exceeded", "Error handling request", "Timed out after "}
+)
+
+func isRetryableError(err string, i int) bool {
+	if i < (maxRetries - 1) {
+		for _, msg := range retriableErrors {
+			if strings.Contains(err, msg) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func runWithRetries[ERR any](runFunc func() ERR, msg string) ERR {
+	var err ERR
+	delay := 1
+	for i := 0; i < maxRetries; i++ {
+		err = runFunc()
+		valOf := reflect.ValueOf(&err).Elem()
+		if valOf.IsNil() || valOf.IsZero() {
+			break
+		}
+		if !isRetryableError(fmt.Sprintf("%v", err), i) {
+			log.Printf("[ERROR] Error %s after %d retries: %v", msg, i, err)
+			return err
+		}
+		delay = delay * retryDelaySeconds
+		log.Printf("[INFO] next retry (%d) for %s after %d seconds", (i + 1), msg, delay)
+		time.Sleep(time.Duration(delay) * time.Second)
+	}
+	return err
+}