Merge pull request #317

* make `result.Wait()` configurable

* Merge remote-tracking branch 'origin/master' into configurable-wait

Author: noamcohen97

azkustoingest/internal/status/status_table_client.go

@@ -16,12 +16,18 @@ const (
 	fullMetadata          = aztables.MetadataFormatFull
 )
 
+type TableClientReader interface {
+	Read(ctx context.Context, ingestionSourceID string) (map[string]interface{}, error)
+}
+
 // TableClient allows reading and writing to azure tables.
 type TableClient struct {
 	tableURI resources.URI
 	client   *aztables.Client
 }
 
+var _ TableClientReader = (*TableClient)(nil)
+
 // NewTableClient Creates an azure table client.
 func NewTableClient(client policy.Transporter, uri resources.URI) (*TableClient, error) {
 	tableClient, err := aztables.NewClientWithNoCredential(uri.String(), &aztables.ClientOptions{

azkustoingest/result.go

@@ -14,7 +14,7 @@ import (
 // Result provides a way for users track the state of ingestion jobs.
 type Result struct {
 	record        statusRecord
-	tableClient   *status.TableClient
+	tableClient   status.TableClientReader
 	reportToTable bool
 }
 
@@ -78,9 +78,58 @@ func (r *Result) putQueued(ctx context.Context, i *Ingestion) {
 	r.tableClient = client
 }
 
+type waitConfig struct {
+	interval           time.Duration
+	immediateFirst     bool
+	retryBackoffDelay  []time.Duration
+	retryBackoffJitter time.Duration
+}
+
+type WaitOption func(o *waitConfig)
+
+func WithImmediateFirst() WaitOption {
+	return func(o *waitConfig) {
+		o.immediateFirst = true
+	}
+}
+
+func WithInterval(interval time.Duration) WaitOption {
+	return func(o *waitConfig) {
+		o.interval = interval
+	}
+}
+
+func WithRetryBackoffDelay(delay ...time.Duration) WaitOption {
+	return func(o *waitConfig) {
+		o.retryBackoffDelay = delay
+	}
+}
+
+func WithRetryBackoffJitter(jitter time.Duration) WaitOption {
+	return func(o *waitConfig) {
+		o.retryBackoffJitter = jitter
+	}
+}
+
+var (
+	DefaultWaitPollInterval           = 10 * time.Second
+	DefaultWaitPollRetryBackoffDelay  = []time.Duration{10 * time.Second, 60 * time.Second, 120 * time.Second}
+	DefaultWaitPollRetryBackoffJitter = 5 * time.Second
+)
+
 // Wait returns a channel that can be checked for ingestion results.
 // In order to check actual status please use the ReportResultToTable option when ingesting data.
-func (r *Result) Wait(ctx context.Context) <-chan error {
+func (r *Result) Wait(ctx context.Context, options ...WaitOption) <-chan error {
+	cfg := waitConfig{
+		interval:           DefaultWaitPollInterval,
+		retryBackoffDelay:  DefaultWaitPollRetryBackoffDelay,
+		retryBackoffJitter: DefaultWaitPollRetryBackoffJitter,
+	}
+
+	for _, o := range options {
+		o(&cfg)
+	}
+
 	ch := make(chan error, 1)
 
 	if r.record.Status.IsFinal() {
@@ -106,7 +155,7 @@ func (r *Result) Wait(ctx context.Context) <-chan error {
 	go func() {
 		defer close(ch)
 
-		r.poll(ctx)
+		r.poll(ctx, &cfg)
 		if !r.record.Status.IsSuccess() {
 			ch <- r.record
 		}
@@ -115,13 +164,13 @@ func (r *Result) Wait(ctx context.Context) <-chan error {
 	return ch
 }
 
-func (r *Result) poll(ctx context.Context) {
-	const pollInterval = 10 * time.Second
-	attempts := 3
-	delay := [3]int{120, 60, 10} // attempts are counted backwards
-
-	// Create a ticker to poll the table in 10 second intervals.
-	timer := time.NewTimer(pollInterval)
+func (r *Result) poll(ctx context.Context, cfg *waitConfig) {
+	initialInterval := cfg.interval
+	if cfg.immediateFirst {
+		initialInterval = 0
+	}
+	attempts := cfg.retryBackoffDelay[:]
+	timer := time.NewTimer(initialInterval)
 	defer timer.Stop()
 
 	for {
@@ -133,24 +182,28 @@ func (r *Result) poll(ctx context.Context) {
 
 		case <-timer.C:
 			smap, err := r.tableClient.Read(ctx, r.record.IngestionSourceID.String())
+			sleepTime := cfg.interval
 			if err != nil {
-				if attempts == 0 {
+				if len(attempts) == 0 {
 					r.record.Status = StatusRetrievalFailed
 					r.record.FailureStatus = Transient
 					r.record.Details = "Failed reading from Status Table: " + err.Error()
 					return
 				}
 
-				attempts = attempts - 1
-				time.Sleep(time.Duration(delay[attempts]+rand.Intn(5)) * time.Second)
+				sleepTime += attempts[0]
+				attempts = attempts[1:]
+				if cfg.retryBackoffJitter > 0 {
+					sleepTime += time.Duration(rand.Intn(int(cfg.retryBackoffJitter)))
+				}
 			} else {
 				r.record.FromMap(smap)
 				if r.record.Status.IsFinal() {
 					return
 				}
 			}
 
-			timer.Reset(pollInterval)
+			timer.Reset(sleepTime)
 		}
 	}
 }

azkustoingest/result_test.go

@@ -0,0 +1,160 @@
+package azkustoingest
+
+import (
+	"context"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/Azure/azure-kusto-go/azkustoingest/internal/status"
+	"github.com/stretchr/testify/assert"
+)
+
+type TableClientReaderFunc func(ctx context.Context, ingestionSourceID string) (map[string]interface{}, error)
+
+func (t TableClientReaderFunc) Read(ctx context.Context, ingestionSourceID string) (map[string]interface{}, error) {
+	return t(ctx, ingestionSourceID)
+}
+
+var _ status.TableClientReader = (*TableClientReaderFunc)(nil)
+
+func TestWait(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		res := &Result{
+			reportToTable: true,
+			tableClient: TableClientReaderFunc(func(ctx context.Context, ingestionSourceID string) (map[string]interface{}, error) {
+				assert.FailNow(t, "Expected table client to not be called")
+				return nil, nil
+			}),
+			record: statusRecord{
+				Status: Pending,
+			},
+		}
+
+		ch := res.Wait(t.Context())
+		synctest.Wait()
+
+		select {
+		case <-ch:
+			assert.FailNow(t, "Expected nothing to be sent on channel")
+		default:
+		}
+	})
+}
+
+func TestWait_ImmediateFirst(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		res := &Result{
+			reportToTable: true,
+			tableClient: TableClientReaderFunc(func(ctx context.Context, ingestionSourceID string) (map[string]any, error) {
+				return map[string]any{"Status": string(Succeeded)}, nil
+			}),
+			record: statusRecord{
+				Status: Pending,
+			},
+		}
+
+		ch := res.Wait(t.Context(), WithImmediateFirst())
+		synctest.Wait()
+
+		select {
+		case <-ch:
+		default:
+			assert.FailNow(t, "Expected something to be sent on channel")
+		}
+	})
+}
+
+func TestWait_WithInterval(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		startTime := time.Now()
+		var calledTimes []time.Duration
+
+		res := &Result{
+			reportToTable: true,
+			tableClient: TableClientReaderFunc(func(ctx context.Context, ingestionSourceID string) (map[string]any, error) {
+				calledTimes = append(calledTimes, time.Now().Sub(startTime))
+				ret := map[string]any{"Status": string(Pending)}
+				if len(calledTimes) >= 3 {
+					ret["Status"] = string(Failed)
+				}
+				return ret, nil
+			}),
+			record: statusRecord{
+				Status: Pending,
+			},
+		}
+
+		ch := res.Wait(t.Context(), WithInterval(5*time.Second))
+		synctest.Wait()
+		assert.Empty(t, calledTimes, "Expected no calls to table client")
+
+		time.Sleep(5 * time.Second)
+		synctest.Wait()
+		assert.Len(t, calledTimes, 1)
+		assert.Equal(t, 5*time.Second, calledTimes[0])
+
+		time.Sleep(5 * time.Second)
+		synctest.Wait()
+		assert.Len(t, calledTimes, 2)
+		assert.Equal(t, 10*time.Second, calledTimes[1])
+
+		time.Sleep(5 * time.Second)
+		synctest.Wait()
+		assert.Len(t, calledTimes, 3)
+		assert.Equal(t, 15*time.Second, calledTimes[2])
+
+		select {
+		case <-ch:
+		default:
+			assert.FailNow(t, "Expected something to be sent on channel")
+		}
+	})
+}
+
+func TestWait_WithRetryBackoffDelay(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		startTime := time.Now()
+		var calledTimes []time.Duration
+
+		res := &Result{
+			reportToTable: true,
+			tableClient: TableClientReaderFunc(func(ctx context.Context, ingestionSourceID string) (map[string]any, error) {
+				calledTimes = append(calledTimes, time.Now().Sub(startTime))
+				return nil, assert.AnError
+			}),
+			record: statusRecord{
+				Status: Pending,
+			},
+		}
+
+		ch := res.Wait(t.Context(), WithRetryBackoffDelay(1*time.Second, 3*time.Second), WithRetryBackoffJitter(0))
+		synctest.Wait()
+		assert.Empty(t, calledTimes, "Expected no calls to table client")
+
+		// First call after DefaultWaitPollInterval (10s)
+		time.Sleep(10 * time.Second)
+		synctest.Wait()
+		assert.Len(t, calledTimes, 1)
+		assert.Equal(t, 10*time.Second, calledTimes[0])
+
+		// Second call after first backoff delay (1s) + poll interval (10s)
+		time.Sleep(11 * time.Second)
+		synctest.Wait()
+		assert.Len(t, calledTimes, 2)
+		assert.Equal(t, 21*time.Second, calledTimes[1])
+
+		// Third call after second backoff delay (3s) + poll interval (10s)
+		time.Sleep(13 * time.Second)
+		synctest.Wait()
+		assert.Len(t, calledTimes, 3)
+		assert.Equal(t, 34*time.Second, calledTimes[2])
+
+		select {
+		case err := <-ch:
+			assert.NotNil(t, err)
+		default:
+			assert.FailNow(t, "Expected something to be sent on channel")
+		}
+	})
+}