backup: Track concurrency per storage in tests

Modifies the parallelism tests for the Pipeline so the number of
concurrent jobs can be tracked per storage. This allows us to ensure
that both the `parallel` and `parallelStorage` thresholds are being
respected, and will come in handy for a future optimisation that will
attempt to balance operations across storages.

1 files changed, 45 insertions, 21 deletions

diff --git a/internal/backup/pipeline_test.go b/internal/backup/pipeline_test.go
index a0a4f2205..f4164d26b 100644
--- a/internal/backup/pipeline_test.go
+++ b/internal/backup/pipeline_test.go
@@ -3,7 +3,7 @@ package backup
 import (
 	"context"
 	"fmt"
-	"sync/atomic"
+	"sync"
 	"testing"
 	"time"
 
@@ -28,39 +28,63 @@ func TestPipeline(t *testing.T) {
 	// Concurrent
 	t.Run("parallelism", func(t *testing.T) {
 		for _, tc := range []struct {
-			parallel            int
-			parallelStorage     int
-			expectedMaxParallel int64
+			parallel                   int
+			parallelStorage            int
+			expectedMaxParallel        int
+			expectedMaxStorageParallel int
 		}{
 			{
-				parallel:            2,
-				parallelStorage:     0,
-				expectedMaxParallel: 2,
+				parallel:                   2,
+				parallelStorage:            0,
+				expectedMaxParallel:        2,
+				expectedMaxStorageParallel: 2,
 			},
 			{
-				parallel:            2,
-				parallelStorage:     3,
-				expectedMaxParallel: 2,
+				parallel:                   2,
+				parallelStorage:            3,
+				expectedMaxParallel:        2,
+				expectedMaxStorageParallel: 2,
 			},
 			{
-				parallel:            0,
-				parallelStorage:     3,
-				expectedMaxParallel: 6, // 2 storages * 3 workers per storage
+				parallel:                   0,
+				parallelStorage:            3,
+				expectedMaxParallel:        6, // 2 storages * 3 workers per storage
+				expectedMaxStorageParallel: 3,
 			},
 			{
-				parallel:            3,
-				parallelStorage:     2,
-				expectedMaxParallel: 3, // `parallel` takes priority, which is why 2 storages * 2 workers is not the max
+				parallel:                   3,
+				parallelStorage:            2,
+				expectedMaxParallel:        3,
+				expectedMaxStorageParallel: 2,
 			},
 		} {
 			t.Run(fmt.Sprintf("parallel:%d,parallelStorage:%d", tc.parallel, tc.parallelStorage), func(t *testing.T) {
-				var calls int64
+				var mu sync.Mutex
+				// callsPerStorage tracks the number of concurrent jobs running for each storage.
+				callsPerStorage := map[string]int{
+					"storage1": 0,
+					"storage2": 0,
+				}
+
 				strategy := MockStrategy{
 					CreateFunc: func(ctx context.Context, req *CreateRequest) error {
-						currentCalls := atomic.AddInt64(&calls, 1)
-						defer atomic.AddInt64(&calls, -1)
-
-						assert.LessOrEqual(t, currentCalls, tc.expectedMaxParallel)
+						mu.Lock()
+						callsPerStorage[req.Repository.StorageName]++
+						allCalls := 0
+						for _, v := range callsPerStorage {
+							allCalls += v
+						}
+						// We ensure that the concurrency for each storage is not above the
+						// parallelStorage threshold, and also that the total number of concurrent
+						// jobs is not above the parallel threshold.
+						require.LessOrEqual(t, callsPerStorage[req.Repository.StorageName], tc.expectedMaxStorageParallel)
+						require.LessOrEqual(t, allCalls, tc.expectedMaxParallel)
+						mu.Unlock()
+						defer func() {
+							mu.Lock()
+							callsPerStorage[req.Repository.StorageName]--
+							mu.Unlock()
+						}()
 
 						time.Sleep(time.Millisecond)
 						return nil