add onCheckResult

Author: joshjennings98

utils/parallelisation/parallelisation.go

@@ -211,7 +211,7 @@ func RunActionWithTimeoutAndCancelStore(ctx context.Context, timeout time.Durati
 	}
 }
 
-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) {
+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), onCheckResult func(res T), checkPeriod time.Duration) (res T, ok bool, err error) {
 	cancelStore := NewCancelFunctionsStore()
 	defer cancelStore.Cancel()
 
@@ -228,6 +228,9 @@ func runActionAndWait[T any](ctx context.Context, wg *sync.WaitGroup, action fun
 				return
 			default:
 				res, ok = checkAction(ctx)
+
+				onCheckResult(res)
+
 				if !ok {
 					store.Cancel()
 					return
@@ -250,7 +253,7 @@ func runActionAndWait[T any](ctx context.Context, wg *sync.WaitGroup, action fun
 // 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) {
+func RunActionWithParallelCheckAndResult[T any](ctx context.Context, action func(ctx context.Context) error, checkAction func(ctx context.Context) (res T, ok bool), onCheckResult func(res T), checkPeriod time.Duration) (res T, ok bool, err error) {
 	err = DetermineContextError(ctx)
 	if err != nil {
 		return
@@ -259,7 +262,7 @@ func RunActionWithParallelCheckAndResult[T any](ctx context.Context, action func
 	var wg sync.WaitGroup
 	defer wg.Wait()
 
-	res, ok, err = runActionAndWait(ctx, &wg, action, checkAction, checkPeriod)
+	res, ok, err = runActionAndWait(ctx, &wg, action, checkAction, onCheckResult, checkPeriod)
 	return
 }
 
@@ -273,6 +276,7 @@ func RunActionWithParallelCheck(ctx context.Context, action func(ctx context.Con
 			ok = checkAction(ctx)
 			return
 		},
+		func(_ struct{}) {},
 		checkPeriod,
 	)
 

utils/parallelisation/parallelisation_test.go

@@ -428,7 +428,8 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 	t.Run("Happy", func(t *testing.T) {
 		defer goleak.VerifyNone(t)
 
-		counter := atomic.NewInt32(0)
+		checkCounter := atomic.NewInt32(0)
+		checkResultCounter := atomic.NewInt32(0)
 
 		res, ok, err := RunActionWithParallelCheckAndResult(
 			context.Background(),
@@ -438,25 +439,30 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 			},
 			func(ctx context.Context) (res parallelisationCheckResult, ok bool) {
 				return parallelisationCheckResult{
-					checks: counter.Inc(),
+					checks: checkCounter.Inc(),
 					status: "healthy",
 				}, true
 			},
+			func(_ parallelisationCheckResult) {
+				checkResultCounter.Inc()
+			},
 			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, res.checks, checkCounter.Load())
 		assert.Equal(t, "healthy", res.status)
+		assert.Equal(t, checkCounter.Load(), checkResultCounter.Load())
 	})
 
 	t.Run("Check Fails With Reason", func(t *testing.T) {
 		defer goleak.VerifyNone(t)
 
-		counter := atomic.NewInt32(0)
+		checkCounter := atomic.NewInt32(0)
+		checkResultCounter := atomic.NewInt32(0)
 		actionStarted := atomic.NewBool(false)
 
 		status := "adrien"
@@ -469,7 +475,7 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 				return DetermineContextError(ctx)
 			},
 			func(ctx context.Context) (res parallelisationCheckResult, ok bool) {
-				if n := counter.Inc(); n >= 5 {
+				if n := checkCounter.Inc(); n >= 5 {
 					return parallelisationCheckResult{
 						checks: n,
 						status: status,
@@ -481,6 +487,9 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 					}, true
 				}
 			},
+			func(_ parallelisationCheckResult) {
+				checkResultCounter.Inc()
+			},
 			5*time.Millisecond,
 		)
 
@@ -491,12 +500,14 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 		require.False(t, ok)
 		assert.Equal(t, status, res.status)
 		assert.Equal(t, int32(5), res.checks)
-		assert.Equal(t, int32(5), counter.Load())
+		assert.Equal(t, int32(5), checkCounter.Load())
+		assert.Equal(t, checkCounter.Load(), checkResultCounter.Load())
 	})
 	t.Run("Action Error (no context cancel)", func(t *testing.T) {
 		defer goleak.VerifyNone(t)
 
-		counter := atomic.NewInt32(0)
+		checkCounter := atomic.NewInt32(0)
+		checkResultCounter := atomic.NewInt32(0)
 		status := "abdel"
 
 		res, ok, err := RunActionWithParallelCheckAndResult(
@@ -507,10 +518,13 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 			},
 			func(ctx context.Context) (parallelisationCheckResult, bool) {
 				return parallelisationCheckResult{
-					checks: counter.Inc(),
+					checks: checkCounter.Inc(),
 					status: status,
 				}, true
 			},
+			func(_ parallelisationCheckResult) {
+				checkResultCounter.Inc()
+			},
 			5*time.Millisecond,
 		)
 
@@ -520,7 +534,8 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 
 		assert.Equal(t, status, res.status)
 		assert.GreaterOrEqual(t, res.checks, int32(1))
-		assert.Equal(t, res.checks, counter.Load())
+		assert.Equal(t, res.checks, checkCounter.Load())
+		assert.Equal(t, checkCounter.Load(), checkResultCounter.Load())
 	})
 
 	t.Run("Context cancel", func(t *testing.T) {
@@ -529,7 +544,8 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 		defer cancel()
 
-		counter := atomic.NewInt32(0)
+		checkCounter := atomic.NewInt32(0)
+		checkResultCounter := atomic.NewInt32(0)
 		status := "kem"
 
 		res, ok, err := RunActionWithParallelCheckAndResult(
@@ -540,17 +556,22 @@ func TestRunActionWithParallelCheckAndResult(t *testing.T) {
 			},
 			func(ctx context.Context) (parallelisationCheckResult, bool) {
 				return parallelisationCheckResult{
-					checks: counter.Inc(),
+					checks: checkCounter.Inc(),
 					status: status,
 				}, true
 			},
+			func(_ parallelisationCheckResult) {
+				checkResultCounter.Inc()
+			},
 			5*time.Millisecond,
 		)
 
 		require.Error(t, err)
 		errortest.AssertError(t, err, commonerrors.ErrTimeout)
 		assert.True(t, ok)
 		assert.GreaterOrEqual(t, res.checks, int32(1))
+		assert.Equal(t, res.checks, checkCounter.Load())
+		assert.Equal(t, checkCounter.Load(), checkResultCounter.Load())
 	})
 }