argo-workflows/workflow/sync/database_mutex_test.go at 05bea5f77e806d35f04b175d1b00ed286b75515d · codefresh-io/argo-workflows · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package sync

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/argoproj/argo-workflows/v3/util/sqldb"
)

// createTestDatabaseMutex creates a database-backed mutex for testing
func createTestDatabaseMutex(t *testing.T, name, namespace string, nextWorkflow NextWorkflow, dbType sqldb.DBType) (*databaseSemaphore, *transaction, func()) {
	t.Helper()
	info, deferfunc, _, err := createTestDBSession(t, dbType)
	require.NoError(t, err)

	dbKey := namespace + "/" + name

	// Create a mutex (which is a semaphore with limit=1)
	mutex := newDatabaseMutex(name, dbKey, nextWorkflow, info)
	require.NotNil(t, mutex)
	tx := &transaction{db: &info.session}
	return mutex, tx, deferfunc
}

// TestDatabaseMutexAcquireRelease tests basic acquire and release functionality
func TestDatabaseMutexAcquireRelease(t *testing.T) {
	for _, dbType := range testDBTypes {
		t.Run(string(dbType), func(t *testing.T) {
			nextWorkflow := func(key string) {}

			mutex, tx, deferfunc := createTestDatabaseMutex(t, "test-mutex", "default", nextWorkflow, dbType)
			defer deferfunc()

			now := time.Now()
			require.NoError(t, mutex.addToQueue("default/workflow1", 0, now))
			require.NoError(t, mutex.addToQueue("default/workflow2", 0, now.Add(time.Second)))

			// First acquisition should succeed
			acquired, _, err := mutex.tryAcquire("default/workflow1", tx)
			require.NoError(t, err)
			assert.True(t, acquired, "First acquisition should succeed")

			// Second acquisition should fail
			acquired, _, err = mutex.tryAcquire("default/workflow2", tx)
			require.NoError(t, err)
			assert.False(t, acquired, "Second acquisition should fail")

			// Release the mutex
			mutex.release("default/workflow1")

			// Now acquisition should succeed again
			acquired, _, err = mutex.tryAcquire("default/workflow2", tx)
			require.NoError(t, err)
			assert.True(t, acquired, "Acquisition after release should succeed")
		})
	}
}

// TestDatabaseMutexQueueOrder tests that workflows are processed in order
func TestDatabaseMutexQueueOrder(t *testing.T) {
	for _, dbType := range testDBTypes {
		t.Run(string(dbType), func(t *testing.T) {
			notified := make(map[string]bool)
			nextWorkflow := func(key string) {
				notified[key] = true
			}

			mutex, tx, deferfunc := createTestDatabaseMutex(t, "test-mutex", "default", nextWorkflow, dbType)
			defer deferfunc()

			// Add items to the queue
			now := time.Now()
			require.NoError(t, mutex.addToQueue("default/workflow1", 0, now))
			require.NoError(t, mutex.addToQueue("default/workflow2", 0, now.Add(time.Second)))

			acquired, _, err := mutex.tryAcquire("default/workflow2", tx)
			require.NoError(t, err)
			assert.False(t, acquired, "Second workflow should not acquire the mutex")

			// Acquire the first one
			acquired, _, err = mutex.tryAcquire("default/workflow1", tx)
			require.NoError(t, err)
			assert.True(t, acquired, "First workflow should acquire the mutex")

			// Release it - this should notify the next one
			mutex.release("default/workflow1")

			// Check that workflow2 was notified
			assert.True(t, notified["default/workflow2"], "workflow2 should be notified")
		})
	}
}