From 07f5297001c446b8855f356fad96077f1f119862 Mon Sep 17 00:00:00 2001
From: Tom Harvey <tombuildsstuff@users.noreply.github.com>
Date: Sat, 29 Sep 2018 06:46:21 +1000
Subject: [PATCH] [WIP] Using the Context from the timeout if provided (#315)

* Using the timeout from the context if available

- Makes PollingDuration optional

* Renaming the registration start time

* Making PollingDuration not a pointer

* fixing a broken reference
---
 autorest/adal/sender.go |  2 +-
 autorest/azure/async.go | 14 +++++++++++---
 autorest/azure/rp.go    |  6 +++---
 autorest/client.go      |  1 +
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/autorest/adal/sender.go b/autorest/adal/sender.go
index 0e5ad14d3..834401e00 100644
--- a/autorest/adal/sender.go
+++ b/autorest/adal/sender.go
@@ -38,7 +38,7 @@ func (sf SenderFunc) Do(r *http.Request) (*http.Response, error) {
 	return sf(r)
 }
 
-// SendDecorator takes and possibily decorates, by wrapping, a Sender. Decorators may affect the
+// SendDecorator takes and possibly decorates, by wrapping, a Sender. Decorators may affect the
 // http.Request and pass it along or, first, pass the http.Request along then react to the
 // http.Response result.
 type SendDecorator func(Sender) Sender
diff --git a/autorest/azure/async.go b/autorest/azure/async.go
index cda1e180a..d3dc73919 100644
--- a/autorest/azure/async.go
+++ b/autorest/azure/async.go
@@ -164,9 +164,17 @@ func (f Future) WaitForCompletion(ctx context.Context, client autorest.Client) e
 // running operation has completed, the provided context is cancelled, or the client's
 // polling duration has been exceeded.  It will retry failed polling attempts based on
 // the retry value defined in the client up to the maximum retry attempts.
-func (f *Future) WaitForCompletionRef(ctx context.Context, client autorest.Client) error {
-	ctx, cancel := context.WithTimeout(ctx, client.PollingDuration)
-	defer cancel()
+func (f *Future) WaitForCompletionRef(inputCtx context.Context, client autorest.Client) error {
+
+	var ctx context.Context
+	if d := client.PollingDuration; d != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(inputCtx, d)
+		defer cancel()
+	} else {
+		ctx = inputCtx
+	}
+
 	done, err := f.Done(client)
 	for attempts := 0; !done; done, err = f.Done(client) {
 		if attempts >= client.RetryAttempts {
diff --git a/autorest/azure/rp.go b/autorest/azure/rp.go
index bd34f0ed5..86ce9f2b5 100644
--- a/autorest/azure/rp.go
+++ b/autorest/azure/rp.go
@@ -140,8 +140,8 @@ func register(client autorest.Client, originalReq *http.Request, re RequestError
 	}
 
 	// poll for registered provisioning state
-	now := time.Now()
-	for err == nil && time.Since(now) < client.PollingDuration {
+	registrationStartTime := time.Now()
+	for err == nil && (client.PollingDuration == 0 || (client.PollingDuration != 0 && time.Since(registrationStartTime) < client.PollingDuration)) {
 		// taken from the resources SDK
 		// https://github.com/Azure/azure-sdk-for-go/blob/9f366792afa3e0ddaecdc860e793ba9d75e76c27/arm/resources/resources/providers.go#L45
 		preparer := autorest.CreatePreparer(
@@ -183,7 +183,7 @@ func register(client autorest.Client, originalReq *http.Request, re RequestError
 			return originalReq.Context().Err()
 		}
 	}
-	if !(time.Since(now) < client.PollingDuration) {
+	if client.PollingDuration != 0 && !(time.Since(registrationStartTime) < client.PollingDuration) {
 		return errors.New("polling for resource provider registration has exceeded the polling duration")
 	}
 	return err
diff --git a/autorest/client.go b/autorest/client.go
index 4e92dcad0..95420a9da 100644
--- a/autorest/client.go
+++ b/autorest/client.go
@@ -153,6 +153,7 @@ type Client struct {
 	PollingDelay time.Duration
 
 	// PollingDuration sets the maximum polling time after which an error is returned.
+	// if zero, the timeout from the Context is used
 	PollingDuration time.Duration
 
 	// RetryAttempts sets the default number of retry attempts for client.