Retry once when failed (#189)

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### Changed
+- Modified `once.go` to reset `sync.Once` instance when an error occurs
+
 ## [0.13.0] - 2023-05-09
 
 ### Added

kusto/utils/once.go

@@ -1,6 +1,9 @@
 package utils
 
-import "sync"
+import (
+	"sync"
+	"sync/atomic"
+)
 
 type Once[Out any] interface {
 	Do(f func() (Out, error)) (Out, error)
@@ -14,10 +17,10 @@ type OnceWithInit[Out any] interface {
 }
 
 type once[Out any] struct {
-	inner  sync.Once
+	mutex  sync.Mutex
 	result Out
 	err    error
-	done   bool
+	done   uint32
 }
 
 type onceWithInit[Out any] struct {
@@ -28,10 +31,10 @@ type onceWithInit[Out any] struct {
 func NewOnce[Out any]() Once[Out] {
 	var empty Out
 	return &once[Out]{
-		inner:  sync.Once{},
+		mutex:  sync.Mutex{},
 		result: empty,
 		err:    nil,
-		done:   false,
+		done:   0,
 	}
 }
 
@@ -58,17 +61,27 @@ func (o *onceWithInit[Out]) Result() (bool, Out, error) {
 }
 
 func (o *once[Out]) Do(f func() (Out, error)) (Out, error) {
-	o.inner.Do(func() {
+	if atomic.LoadUint32(&o.done) != 0 {
+		return o.result, o.err
+	}
+
+	o.mutex.Lock()
+	defer o.mutex.Unlock()
+	if o.done == 0 {
 		o.result, o.err = f()
-		o.done = true
-	})
+
+		if o.err == nil {
+			defer atomic.StoreUint32(&o.done, 1)
+		}
+	}
+
 	return o.result, o.err
 }
 
 func (o *once[Out]) Done() bool {
-	return o.done
+	return o.done != 0
 }
 
 func (o *once[Out]) Result() (bool, Out, error) {
-	return o.done, o.result, o.err
+	return o.Done(), o.result, o.err
 }

kusto/utils/once_test.go

@@ -0,0 +1,110 @@
+package utils
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name     string
+		withInit bool
+		failOnce bool
+		err      error
+	}{
+		{
+			name:     "Test Success",
+			withInit: false,
+		},
+		{
+			name:     "Test Success With Init",
+			withInit: true,
+		},
+		{
+			name:     "Test Failure",
+			withInit: false,
+			err:      errors.New("test"),
+		},
+		{
+			name:     "Test Failure With Init",
+			withInit: true,
+			err:      errors.New("test"),
+		},
+		{
+			name:     "Test Twice",
+			withInit: false,
+			failOnce: true,
+			err:      errors.New("test"),
+		},
+		{
+			name:     "Test Twice With Init",
+			withInit: true,
+			failOnce: true,
+			err:      errors.New("test"),
+		},
+	}
+
+	for _, test := range tests {
+		test := test // Capture
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+			counter := 0
+			funErr := test.err
+			funFailOnce := test.failOnce
+
+			f := func() (int, error) {
+				counter++
+				if (funFailOnce && counter == 1) || (!funFailOnce && funErr != nil) {
+					return 0, errors.New("test")
+				} else {
+					return 1, nil
+				}
+			}
+
+			var result int
+			var err error
+			var once Once[int]
+			if test.withInit {
+				once = NewOnce[int]()
+			} else {
+				once = NewOnceWithInit[int](f)
+			}
+
+			expectedOnSuccess := 1
+
+			for i := 0; i < 10; i++ {
+				if withInit, ok := once.(OnceWithInit[int]); ok {
+					result, err = withInit.DoWithInit()
+				} else {
+					result, err = once.Do(f)
+				}
+
+				isDone, onceResult, onceErr := once.Result()
+				if test.err != nil {
+					assert.Equal(t, i+1, counter)
+					assert.False(t, isDone)
+					assert.Equal(t, 0, onceResult)
+					assert.Equal(t, test.err, onceErr)
+					assert.Equal(t, test.err, err)
+				} else {
+					assert.Equal(t, counter, expectedOnSuccess)
+					assert.True(t, once.Done())
+					assert.True(t, isDone)
+					assert.Equal(t, 1, onceResult)
+					require.NoError(t, err)
+					assert.Equal(t, 1, result)
+				}
+
+				if test.failOnce {
+					test.err = nil
+					expectedOnSuccess = 2
+				}
+			}
+
+		})
+	}
+}