[Cosmos] Implements Client Retry policy

sdk/data/azcosmos/ci.yml

@@ -25,6 +25,7 @@ stages:
   parameters:
     ServiceDirectory: 'data/azcosmos'
     UsePipelineProxy: false
+    ExcludeGoNMinus2: true
 - stage: Emulator
   displayName: 'Cosmos Emulator'
   variables:

sdk/data/azcosmos/cosmos_client.go

@@ -15,7 +15,6 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
 )
@@ -42,10 +41,11 @@ func (c *Client) Endpoint() string {
 // options - Optional Cosmos client options.  Pass nil to accept default values.
 func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) {
 	preferredRegions := []string{}
+	enableCrossRegionRetries := true
 	if o != nil {
 		preferredRegions = o.PreferredRegions
 	}
-	gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), preferredRegions, 0)
+	gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newSharedKeyCredPolicy(cred), o), preferredRegions, 0, enableCrossRegionRetries)
 	if err != nil {
 		return nil, err
 	}
@@ -62,10 +62,11 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) (
 		return nil, err
 	}
 	preferredRegions := []string{}
+	enableCrossRegionRetries := true
 	if o != nil {
 		preferredRegions = o.PreferredRegions
 	}
-	gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), preferredRegions, 0)
+	gem, err := newGlobalEndpointManager(endpoint, newInternalPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o), preferredRegions, 0, enableCrossRegionRetries)
 	if err != nil {
 		return nil, err
 	}
@@ -124,6 +125,7 @@ func newPipeline(authPolicy policy.Policy, gem *globalEndpointManager, options *
 			},
 			PerRetry: []policy.Policy{
 				authPolicy,
+				&clientRetryPolicy{gem: gem},
 			},
 		},
 		&options.ClientOptions)
@@ -195,8 +197,10 @@ func (c *Client) CreateDatabase(
 	}
 
 	operationContext := pipelineRequestOptions{
-		resourceType:    resourceTypeDatabase,
-		resourceAddress: ""}
+		resourceType:     resourceTypeDatabase,
+		resourceAddress:  "",
+		isWriteOperation: true,
+	}
 
 	path, err := generatePathForNameBased(resourceTypeDatabase, "", true)
 	if err != nil {
@@ -220,7 +224,7 @@ func (c *Client) CreateDatabase(
 // NewQueryDatabasesPager executes query for databases.
 // query - The SQL query to execute.
 // o - Options for the operation.
-func (c *Client) NewQueryDatabasesPager(query string, o *QueryDatabasesOptions) *runtime.Pager[QueryDatabasesResponse] {
+func (c *Client) NewQueryDatabasesPager(query string, o *QueryDatabasesOptions) *azruntime.Pager[QueryDatabasesResponse] {
 	queryOptions := &QueryDatabasesOptions{}
 	if o != nil {
 		originalOptions := *o
@@ -234,7 +238,7 @@ func (c *Client) NewQueryDatabasesPager(query string, o *QueryDatabasesOptions)
 
 	path, _ := generatePathForNameBased(resourceTypeDatabase, operationContext.resourceAddress, true)
 
-	return runtime.NewPager(runtime.PagingHandler[QueryDatabasesResponse]{
+	return azruntime.NewPager(azruntime.PagingHandler[QueryDatabasesResponse]{
 		More: func(page QueryDatabasesResponse) bool {
 			return page.ContinuationToken != ""
 		},

sdk/data/azcosmos/cosmos_client_retry_policy.go

@@ -0,0 +1,188 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"net/http"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/errorinfo"
+)
+
+type clientRetryPolicy struct {
+	gem                    *globalEndpointManager
+	useWriteEndpoint       bool
+	retryCount             int
+	sessionRetryCount      int
+	preferredLocationIndex int
+}
+
+const maxRetryCount = 120
+const defaultBackoff = 1
+
+func (p *clientRetryPolicy) Do(req *policy.Request) (*http.Response, error) {
+	p.resetPolicyCounters()
+	o := pipelineRequestOptions{}
+	if !req.OperationValue(&o) {
+		return nil, fmt.Errorf("failed to obtain request options, please check request being sent: %s", req.Body())
+	}
+	for {
+		resolvedEndpoint := p.gem.ResolveServiceEndpoint(p.retryCount, o.isWriteOperation, p.useWriteEndpoint)
+		req.Raw().Host = resolvedEndpoint.Host
+		req.Raw().URL.Host = resolvedEndpoint.Host
+		response, err := req.Next() // err can happen in weird scenarios (connectivity, etc)
+		if err != nil {
+			if p.isNetworkConnectionError(err) {
+				shouldRetry, err := p.attemptRetryOnNetworkError(req)
+				if err != nil {
+					return nil, err
+				}
+				if !shouldRetry {
+					return nil, errorinfo.NonRetriableError(azruntime.NewResponseErrorWithErrorCode(response, response.Status))
+				}
+				err = req.RewindBody()
+				if err != nil {
+					return nil, err
+				}
+				p.retryCount += 1
+				continue
+			}
+
+			return nil, err
+		}
+		subStatus := response.Header.Get(cosmosHeaderSubstatus)
+		if p.shouldRetryStatus(response.StatusCode, subStatus) {
+			p.useWriteEndpoint = false
+			if response.StatusCode == http.StatusForbidden {
+				shouldRetry, err := p.attemptRetryOnEndpointFailure(req, o.isWriteOperation)
+				if err != nil {
+					return nil, err
+				}
+				if !shouldRetry {
+					return nil, errorinfo.NonRetriableError(azruntime.NewResponseErrorWithErrorCode(response, response.Status))
+				}
+			} else if response.StatusCode == http.StatusNotFound {
+				if !p.attemptRetryOnSessionUnavailable(req, o.isWriteOperation) {
+					return nil, errorinfo.NonRetriableError(azruntime.NewResponseErrorWithErrorCode(response, response.Status))
+				}
+			} else if response.StatusCode == http.StatusServiceUnavailable {
+				if !p.attemptRetryOnServiceUnavailable(req, o.isWriteOperation) {
+					return nil, errorinfo.NonRetriableError(azruntime.NewResponseErrorWithErrorCode(response, response.Status))
+				}
+			}
+			err = req.RewindBody()
+			if err != nil {
+				return response, err
+			}
+			p.retryCount += 1
+			continue
+		}
+
+		return response, err
+	}
+
+}
+
+func (p *clientRetryPolicy) shouldRetryStatus(status int, subStatus string) (shouldRetry bool) {
+	if (status == http.StatusForbidden && (subStatus == subStatusWriteForbidden || subStatus == subStatusDatabaseAccountNotFound)) ||
+		(status == http.StatusNotFound && subStatus == subStatusReadSessionNotAvailable) ||
+		(status == http.StatusServiceUnavailable) {
+		return true
+	}
+	return false
+}
+
+func (p *clientRetryPolicy) attemptRetryOnNetworkError(req *policy.Request) (bool, error) {
+	if (p.retryCount > maxRetryCount) || !p.gem.locationCache.enableCrossRegionRetries {
+		return false, nil
+	}
+
+	err := p.gem.MarkEndpointUnavailableForWrite(*req.Raw().URL)
+	if err != nil {
+		return false, err
+	}
+	err = p.gem.MarkEndpointUnavailableForRead(*req.Raw().URL)
+	if err != nil {
+		return false, err
+	}
+	err = p.gem.Update(req.Raw().Context(), false)
+	if err != nil {
+		return false, err
+	}
+
+	time.Sleep(defaultBackoff * time.Second)
+	return true, nil
+}
+
+func (p *clientRetryPolicy) attemptRetryOnEndpointFailure(req *policy.Request, isWriteOperation bool) (bool, error) {
+	if (p.retryCount > maxRetryCount) || !p.gem.locationCache.enableCrossRegionRetries {
+		return false, nil
+	}
+	if isWriteOperation {
+		err := p.gem.MarkEndpointUnavailableForWrite(*req.Raw().URL)
+		if err != nil {
+			return false, err
+		}
+	} else {
+		err := p.gem.MarkEndpointUnavailableForRead(*req.Raw().URL)
+		if err != nil {
+			return false, err
+		}
+	}
+
+	err := p.gem.Update(req.Raw().Context(), isWriteOperation)
+	if err != nil {
+		return false, err
+	}
+
+	time.Sleep(defaultBackoff * time.Second)
+	return true, nil
+}
+
+func (p *clientRetryPolicy) attemptRetryOnSessionUnavailable(req *policy.Request, isWriteOperation bool) bool {
+	if p.gem.CanUseMultipleWriteLocations() {
+		endpoints := p.gem.locationCache.locationInfo.availReadLocations
+		if isWriteOperation {
+			endpoints = p.gem.locationCache.locationInfo.availWriteLocations
+		}
+		if p.sessionRetryCount >= len(endpoints) {
+			return false
+		}
+	} else {
+		if p.sessionRetryCount > 0 {
+			return false
+		}
+		p.useWriteEndpoint = true
+	}
+	p.sessionRetryCount += 1
+	return true
+}
+
+func (p *clientRetryPolicy) attemptRetryOnServiceUnavailable(req *policy.Request, isWriteOperation bool) bool {
+	if !p.gem.locationCache.enableCrossRegionRetries || p.preferredLocationIndex >= len(p.gem.preferredLocations) {
+		return false
+	}
+	if isWriteOperation && !p.gem.CanUseMultipleWriteLocations() {
+		return false
+	}
+	p.preferredLocationIndex += 1
+	return true
+}
+
+func (p *clientRetryPolicy) resetPolicyCounters() {
+	p.retryCount = 0
+	p.sessionRetryCount = 0
+	p.preferredLocationIndex = 0
+}
+
+// isNetworkConnectionError checks if the error is related to failure to connect / resolve DNS
+func (p *clientRetryPolicy) isNetworkConnectionError(err error) bool {
+	var dnserror *net.DNSError
+	return errors.As(err, &dnserror)
+}