[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)
@@ -219,7 +221,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
@@ -233,7 +235,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,129 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+type clientRetryPolicy struct {
+	gem                    *globalEndpointManager
+	retryCount             int
+	sessionRetryCount      int
+	preferredLocationIndex int
+}
+
+const maxRetryCount = 10 // this number needs to be higher, keeping low for testing purposes right now
+const defaultBackoff = 1
+
+func (p *clientRetryPolicy) Do(req *policy.Request) (*http.Response, error) {
+	response, err := req.Next() // err can happen in weird scenarios (connectivity, etc) - need to test
+	o := pipelineRequestOptions{}
+	req.OperationValue(&o)
+	subStatus := response.Header.Get(cosmosHeaderSubstatus)
+	if p.shouldRetryStatus(response.StatusCode, subStatus) {
+		p.retryCount = 0
+		p.sessionRetryCount = 0
+		for {
+			resolvedEndpoint := p.gem.ResolveServiceEndpoint(p.retryCount, o.isWriteOperation)
+			req.Raw().Host = resolvedEndpoint.Host
+			req.Raw().URL.Host = resolvedEndpoint.Host
+			subStatus = response.Header.Get(cosmosHeaderSubstatus)
+			if p.shouldRetryStatus(response.StatusCode, subStatus) {
+				fmt.Println("Policy TIME")
+				if response.StatusCode == statusForbidden {
+					if !p.attemptRetryOnEndpointFailure(req, o.isWriteOperation) {
+						break
+					}
+				} else if response.StatusCode == statusNotFound {
+					if !p.attemptRetryOnSessionUnavailable(req, o.isWriteOperation) {
+						break
+					}
+				} else if response.StatusCode == statusServiceUnavailable {
+					if !p.attemptRetryOnServiceUnavailable(req, o.isWriteOperation) {
+						break
+					}
+				}
+				fmt.Println("bout to retry this")
+			} else {
+				fmt.Println("supposed to break this")
+				break
+			}
+			response, err = req.Next()
+			fmt.Println("should have retried")
+		}
+	}
+	return response, err
+}
+
+func (p *clientRetryPolicy) shouldRetryStatus(status int, subStatus string) (shouldRetry bool) {
+	if (status == statusForbidden && (subStatus == subStatusWriteForbidden || subStatus == subStatusDatabaseAccountNotFound)) ||
+		(status == statusNotFound && subStatus == subStatusReadSessionNotAvailable) ||
+		(status == statusServiceUnavailable) {
+		return true
+	}
+	return false
+}
+
+func (p *clientRetryPolicy) attemptRetryOnEndpointFailure(req *policy.Request, isWriteOperation bool) bool {
+	if (p.retryCount > maxRetryCount) || !p.gem.locationCache.enableCrossRegionRetries {
+		return false
+	}
+	if isWriteOperation {
+		p.gem.MarkEndpointUnavailableForWrite(*req.Raw().URL)
+	} else {
+		p.gem.MarkEndpointUnavailableForRead(*req.Raw().URL)
+	}
+	p.gem.Update(req.Raw().Context())
+
+	p.retryCount += 1
+	time.Sleep(defaultBackoff * time.Second)
+	return true
+}
+
+func (p *clientRetryPolicy) attemptRetryOnSessionUnavailable(req *policy.Request, isWriteOperation bool) bool {
+	if p.gem.CanUseMultipleWriteLocations() {
+		endpoints := []string{}
+		if isWriteOperation {
+			endpoints = p.gem.locationCache.locationInfo.availWriteLocations
+		} else {
+			endpoints = p.gem.locationCache.locationInfo.availReadLocations
+		}
+		if p.sessionRetryCount >= len(endpoints) {
+			return false
+		}
+	} else {
+		if p.sessionRetryCount > 0 {
+			return false
+		}
+	}
+	p.sessionRetryCount += 1
+	return true
+}
+
+func (p *clientRetryPolicy) attemptRetryOnServiceUnavailable(req *policy.Request, isWriteOperation bool) bool {
+	//On HTTP 503 response, if it's a read request and preferredRegions > 1,
+	//retry on the next preferredRegion. If it's a write request and account is multi master
+	//and preferredRegions > 1, retry on the next preferredRegion.
+	if !p.gem.locationCache.enableCrossRegionRetries || p.preferredLocationIndex >= len(p.gem.preferredLocations) {
+		return false
+	}
+	if isWriteOperation {
+		if p.gem.CanUseMultipleWriteLocations() {
+			locationalEndpoint := p.gem.GetPreferredLocationEndpoint(p.preferredLocationIndex, *req.Raw().URL)
+			req.Raw().URL = &locationalEndpoint
+		} else {
+			return false
+		}
+	} else {
+		locationalEndpoint := p.gem.GetPreferredLocationEndpoint(p.preferredLocationIndex, *req.Raw().URL)
+		req.Raw().URL = &locationalEndpoint
+	}
+	p.preferredLocationIndex += 1
+	return true
+}

sdk/data/azcosmos/cosmos_global_endpoint_manager.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -26,7 +27,7 @@ type globalEndpointManager struct {
 	lastUpdateTime      time.Time
 }
 
-func newGlobalEndpointManager(clientEndpoint string, pipeline azruntime.Pipeline, preferredLocations []string, refreshTimeInterval time.Duration) (*globalEndpointManager, error) {
+func newGlobalEndpointManager(clientEndpoint string, pipeline azruntime.Pipeline, preferredLocations []string, refreshTimeInterval time.Duration, enableCrossRegionRetries bool) (*globalEndpointManager, error) {
 	endpoint, err := url.Parse(clientEndpoint)
 	if err != nil {
 		return &globalEndpointManager{}, err
@@ -40,7 +41,7 @@ func newGlobalEndpointManager(clientEndpoint string, pipeline azruntime.Pipeline
 		clientEndpoint:      clientEndpoint,
 		pipeline:            pipeline,
 		preferredLocations:  preferredLocations,
-		locationCache:       newLocationCache(preferredLocations, *endpoint),
+		locationCache:       newLocationCache(preferredLocations, *endpoint, enableCrossRegionRetries),
 		refreshTimeInterval: refreshTimeInterval,
 		lastUpdateTime:      time.Time{},
 	}
@@ -84,6 +85,29 @@ func (gem *globalEndpointManager) ShouldRefresh() bool {
 	return time.Since(gem.lastUpdateTime) > gem.refreshTimeInterval
 }
 
+func (gem *globalEndpointManager) ResolveServiceEndpoint(locationIndex int, isWriteOperation bool) url.URL {
+	return gem.locationCache.resolveServiceEndpoint(locationIndex, isWriteOperation)
+}
+
+func (gem *globalEndpointManager) GetPreferredLocationEndpoint(preferredLocationIndex int, currentUrl url.URL) url.URL {
+	endpointString := currentUrl.String()
+	location := gem.preferredLocations[preferredLocationIndex]
+	endpointParts := strings.Split(endpointString, ".")
+	if len(endpointParts) > 0 {
+		databaseAccountName := endpointParts[0]
+		locationalDatabaseAccountName := databaseAccountName + "-" + strings.ToLower(strings.ReplaceAll(location, " ", ""))
+		endpointParts[0] = locationalDatabaseAccountName
+		locationalString := strings.Join(endpointParts, ".")
+		locationalURL, err := url.Parse(locationalString)
+		if err != nil {
+			// error parsing the new url
+			return currentUrl
+		}
+		return *locationalURL
+	}
+	return currentUrl
+}
+
 func (gem *globalEndpointManager) Update(ctx context.Context) error {
 	gem.gemMutex.Lock()
 	defer gem.gemMutex.Unlock()

sdk/data/azcosmos/cosmos_global_endpoint_manager_test.go

@@ -34,7 +34,7 @@ func TestGlobalEndpointManagerGetWriteEndpoints(t *testing.T) {
 
 	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv})
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	writeEndpoints, err := gem.GetWriteEndpoints()
@@ -57,7 +57,7 @@ func TestGlobalEndpointManagerGetReadEndpoints(t *testing.T) {
 
 	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv})
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	readEndpoints, err := gem.GetReadEndpoints()
@@ -84,7 +84,7 @@ func TestGlobalEndpointManagerMarkEndpointUnavailableForRead(t *testing.T) {
 	endpoint, err := url.Parse(client.endpoint)
 	assert.NoError(t, err)
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	err = gem.MarkEndpointUnavailableForRead(*endpoint)
@@ -106,7 +106,7 @@ func TestGlobalEndpointManagerMarkEndpointUnavailableForWrite(t *testing.T) {
 	endpoint, err := url.Parse(client.endpoint)
 	assert.NoError(t, err)
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	err = gem.MarkEndpointUnavailableForWrite(*endpoint)
@@ -142,7 +142,7 @@ func TestGlobalEndpointManagerGetEndpointLocation(t *testing.T) {
 	serverEndpoint, err := url.Parse(srv.URL())
 	assert.NoError(t, err)
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	err = gem.Update(context.Background())
@@ -161,7 +161,7 @@ func TestGlobalEndpointManagerGetAccountProperties(t *testing.T) {
 
 	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv})
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	accountProps, err := gem.GetAccountProperties(context.Background())
@@ -189,9 +189,8 @@ func TestGlobalEndpointManagerCanUseMultipleWriteLocations(t *testing.T) {
 	serverEndpoint, err := url.Parse(srv.URL())
 	assert.NoError(t, err)
 
-	mockLc := newLocationCache(preferredRegions, *serverEndpoint)
+	mockLc := newLocationCache(preferredRegions, *serverEndpoint, true)
 	mockLc.enableMultipleWriteLocations = true
-	mockLc.useMultipleWriteLocations = true
 
 	mockGem := globalEndpointManager{
 		clientEndpoint:      client.endpoint,
@@ -200,7 +199,7 @@ func TestGlobalEndpointManagerCanUseMultipleWriteLocations(t *testing.T) {
 		refreshTimeInterval: 5 * time.Minute,
 	}
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Minute)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Minute, true)
 	assert.NoError(t, err)
 
 	// Multiple locations should be false for default GEM
@@ -237,7 +236,7 @@ func TestGlobalEndpointManagerConcurrentUpdate(t *testing.T) {
 
 	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{countPolicy}}, &policy.ClientOptions{Transport: srv})
 
-	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Second)
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{}, 5*time.Second, true)
 	assert.NoError(t, err)
 
 	// Call update concurrently and see how many times the policy gets called
@@ -271,5 +270,27 @@ func TestGlobalEndpointManagerConcurrentUpdate(t *testing.T) {
 	assert.NoError(t, err)
 	callCount = countPolicy.callCount
 	assert.Equal(t, callCount, 2)
+}
+
+func TestGlobalEndpointManagerGetPreferredLocationEndpoint(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(mock.WithStatusCode(http.StatusOK))
+
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv})
+
+	gem, err := newGlobalEndpointManager(srv.URL(), pl, []string{"West US", "Central US"}, 5*time.Minute, true)
+	assert.NoError(t, err)
+
+	testUrl, err := url.Parse("https://contoso.documents.azure.com:443/")
+	assert.NoError(t, err)
+
+	expectedWestLocationalEndpoint := "https://contoso-westus.documents.azure.com:443/"
+	expectedCentralLocationalEndpoint := "https://contoso-centralus.documents.azure.com:443/"
+
+	westLocationalEndpoint := gem.GetPreferredLocationEndpoint(0, *testUrl)
+	centralLocationalEndpoint := gem.GetPreferredLocationEndpoint(1, *testUrl)
 
+	assert.Equal(t, expectedWestLocationalEndpoint, westLocationalEndpoint.String())
+	assert.Equal(t, expectedCentralLocationalEndpoint, centralLocationalEndpoint.String())
 }

sdk/data/azcosmos/cosmos_headers.go → sdk/data/azcosmos/cosmos_http_constants.go

@@ -7,6 +7,7 @@ const (
 	cosmosHeaderRequestCharge                      string = "x-ms-request-charge"
 	cosmosHeaderActivityId                         string = "x-ms-activity-id"
 	cosmosHeaderEtag                               string = "etag"
+	cosmosHeaderSubstatus                          string = "x-ms-substatus"
 	cosmosHeaderPopulateQuotaInfo                  string = "x-ms-documentdb-populatequotainfo"
 	cosmosHeaderPreTriggerInclude                  string = "x-ms-documentdb-pre-trigger-include"
 	cosmosHeaderPostTriggerInclude                 string = "x-ms-documentdb-post-trigger-include"
@@ -46,3 +47,15 @@ const (
 	cosmosHeaderValuesPreferMinimal string = "return=minimal"
 	cosmosHeaderValuesQuery         string = "application/query+json"
 )
+
+const (
+	statusForbidden          int = 403
+	statusNotFound           int = 404
+	statusServiceUnavailable int = 503
+)
+
+const (
+	subStatusWriteForbidden          string = "3"
+	subStatusDatabaseAccountNotFound string = "1008"
+	subStatusReadSessionNotAvailable string = "1002"
+)