✨ parallisation Add support for more context when returning from RunActionWithParallelCheck

Author: joshjennings98

changes/20251205145557.feature

@@ -0,0 +1 @@
+:sparkles: `parallisation` Add support for more context when returning from RunActionWithParallelCheck

utils/parallelisation/parallelisation.go

@@ -9,6 +9,7 @@ package parallelisation
 import (
 	"context"
 	"reflect"
+	"sync"
 	"time"
 
 	"go.uber.org/atomic"
@@ -210,38 +211,72 @@ func RunActionWithTimeoutAndCancelStore(ctx context.Context, timeout time.Durati
 	}
 }
 
-// RunActionWithParallelCheck runs an action with a check in parallel
-// The function performing the check should return true if the check was favourable; false otherwise. If the check did not have the expected result and the whole function would be cancelled.
-func RunActionWithParallelCheck(ctx context.Context, action func(ctx context.Context) error, checkAction func(ctx context.Context) bool, checkPeriod time.Duration) error {
-	err := DetermineContextError(ctx)
-	if err != nil {
-		return err
-	}
+func runActionAndWait[T any](ctx context.Context, wg *sync.WaitGroup, action func(ctx context.Context) error, checkAction func(ctx context.Context) (res T, ok bool), checkPeriod time.Duration) (res T, ok bool, err error) {
 	cancelStore := NewCancelFunctionsStore()
 	defer cancelStore.Cancel()
+
 	cancellableCtx, cancelFunc := context.WithCancel(ctx)
 	cancelStore.RegisterCancelFunction(cancelFunc)
+
+	wg.Add(1)
 	go func(ctx context.Context, store *CancelFunctionStore) {
+		defer wg.Done()
 		for {
 			select {
 			case <-ctx.Done():
 				store.Cancel()
 				return
 			default:
-				if !checkAction(ctx) {
+				res, ok = checkAction(ctx)
+				if !ok {
 					store.Cancel()
 					return
 				}
+
 				SleepWithContext(ctx, checkPeriod)
 			}
 		}
 	}(cancellableCtx, cancelStore)
+
 	err = action(cancellableCtx)
-	err2 := DetermineContextError(cancellableCtx)
-	if err2 != nil {
-		return err2
+	if errCtx := DetermineContextError(cancellableCtx); errCtx != nil {
+		err = errCtx
+	}
+
+	return
+}
+
+// RunActionWithParallelCheck runs an action with a check in parallel
+// The function performing the check should return true if the check was favourable; false otherwise.
+// For more context, a result can be returned. If the check did not have the expected result and the
+// whole function would be cancelled.
+func RunActionWithParallelCheckAndResult[T any](ctx context.Context, action func(ctx context.Context) error, checkAction func(ctx context.Context) (res T, ok bool), checkPeriod time.Duration) (res T, ok bool, err error) {
+	err = DetermineContextError(ctx)
+	if err != nil {
+		return
 	}
-	return err
+
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	res, ok, err = runActionAndWait(ctx, &wg, action, checkAction, checkPeriod)
+	return
+}
+
+// RunActionWithParallelCheck runs an action with a check in parallel
+// The function performing the check should return true if the check was favourable; false otherwise. If the check did not have the expected result and the whole function would be cancelled.
+func RunActionWithParallelCheck(ctx context.Context, action func(ctx context.Context) error, checkAction func(ctx context.Context) bool, checkPeriod time.Duration) (err error) {
+	_, _, err = RunActionWithParallelCheckAndResult(
+		ctx,
+		action,
+		func(ctx context.Context) (_ struct{}, ok bool) {
+			ok = checkAction(ctx)
+			return
+		},
+		checkPeriod,
+	)
+
+	return
 }
 
 // WaitUntil waits for a condition evaluated by evalCondition to be verified

utils/parallelisation/parallelisation_test.go

@@ -14,6 +14,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-faker/faker/v4"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
@@ -418,6 +419,141 @@ func runActionWithParallelCheckFailAtRandom(t *testing.T, ctx context.Context) {
 	assert.GreaterOrEqual(t, counter.Load(), int32(1))
 }
 
+func TestRunActionWithParallelCheckAndResult(t *testing.T) {
+	type parallelisationCheckResult struct {
+		checks int32
+		status string
+	}
+
+	t.Run("Happy", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		counter := atomic.NewInt32(0)
+
+		res, ok, err := RunActionWithParallelCheckAndResult(
+			context.Background(),
+			func(ctx context.Context) (err error) {
+				time.Sleep(120 * time.Millisecond)
+				return
+			},
+			func(ctx context.Context) (res parallelisationCheckResult, ok bool) {
+				return parallelisationCheckResult{
+					checks: counter.Inc(),
+					status: "healthy",
+				}, true
+			},
+			10*time.Millisecond,
+		)
+
+		require.NoError(t, err)
+		require.True(t, ok)
+
+		assert.GreaterOrEqual(t, res.checks, int32(10))
+		assert.Equal(t, res.checks, counter.Load())
+		assert.Equal(t, "healthy", res.status)
+	})
+
+	t.Run("Check Fails With Reason", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		counter := atomic.NewInt32(0)
+		actionStarted := atomic.NewBool(false)
+
+		status := "adrien"
+
+		res, ok, err := RunActionWithParallelCheckAndResult(
+			context.Background(),
+			func(ctx context.Context) error {
+				actionStarted.Store(true)
+				<-ctx.Done()
+				return DetermineContextError(ctx)
+			},
+			func(ctx context.Context) (res parallelisationCheckResult, ok bool) {
+				if n := counter.Inc(); n >= 5 {
+					return parallelisationCheckResult{
+						checks: n,
+						status: status,
+					}, false
+				} else {
+					return parallelisationCheckResult{
+						checks: n,
+						status: "ok",
+					}, true
+				}
+			},
+			5*time.Millisecond,
+		)
+
+		require.True(t, actionStarted.Load())
+		require.Error(t, err)
+		errortest.AssertError(t, err, commonerrors.ErrCancelled)
+
+		require.False(t, ok)
+		assert.Equal(t, status, res.status)
+		assert.Equal(t, int32(5), res.checks)
+		assert.Equal(t, int32(5), counter.Load())
+	})
+	t.Run("Action Error (no context cancel)", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		counter := atomic.NewInt32(0)
+		status := "abdel"
+
+		res, ok, err := RunActionWithParallelCheckAndResult(
+			context.Background(),
+			func(ctx context.Context) error {
+				time.Sleep(30 * time.Millisecond)
+				return commonerrors.New(commonerrors.ErrForbidden, faker.Sentence())
+			},
+			func(ctx context.Context) (parallelisationCheckResult, bool) {
+				return parallelisationCheckResult{
+					checks: counter.Inc(),
+					status: status,
+				}, true
+			},
+			5*time.Millisecond,
+		)
+
+		require.Error(t, err)
+		errortest.AssertError(t, err, commonerrors.ErrForbidden)
+		require.True(t, ok)
+
+		assert.Equal(t, status, res.status)
+		assert.GreaterOrEqual(t, res.checks, int32(1))
+		assert.Equal(t, res.checks, counter.Load())
+	})
+
+	t.Run("Context cancel", func(t *testing.T) {
+		defer goleak.VerifyNone(t)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+		defer cancel()
+
+		counter := atomic.NewInt32(0)
+		status := "kem"
+
+		res, ok, err := RunActionWithParallelCheckAndResult(
+			ctx,
+			func(ctx context.Context) error {
+				<-ctx.Done()
+				return DetermineContextError(ctx)
+			},
+			func(ctx context.Context) (parallelisationCheckResult, bool) {
+				return parallelisationCheckResult{
+					checks: counter.Inc(),
+					status: status,
+				}, true
+			},
+			5*time.Millisecond,
+		)
+
+		require.Error(t, err)
+		errortest.AssertError(t, err, commonerrors.ErrTimeout)
+		assert.True(t, ok)
+		assert.GreaterOrEqual(t, res.checks, int32(1))
+	})
+}
+
 func TestWaitUntil(t *testing.T) {
 	defer goleak.VerifyNone(t)
 	verifiedCondition := func(ctx context.Context) (bool, error) {