1
package sqlite3
2
 
3
import (
4
	"context"
5
	"crypto/rand"
6
	"testing"
7
	"time"
8
 
9
	"go.miragespace.co/specter/spec/chord"
10
 
11
	"github.com/stretchr/testify/assert"
12
	"github.com/stretchr/testify/require"
13
)
14
 
15
func TestAcquireMutualExclusion(t *testing.T) {
16
	as := require.New(t)
17
 
18
	kv := testGetKV(t)
19
 
20
	key := make([]byte, 8)
21
	rand.Read(key)
22
 
23
	token, err := kv.Acquire(context.Background(), key, time.Second)
24
	as.NoError(err)
25
 
26
	_, err = kv.Acquire(context.Background(), key, time.Second)
27
	as.ErrorIs(err, chord.ErrKVLeaseConflict)
28
 
29
	as.NoError(kv.Release(context.Background(), key, token))
30
}
31
 
32
func TestAcquireExpired(t *testing.T) {
33
	as := require.New(t)
34
 
35
	kv := testGetKV(t)
36
 
37
	key := make([]byte, 8)
38
	rand.Read(key)
39
 
40
	// expired 2 seconds ago
41
	expiredToken := uint64(time.Now().Add(time.Duration(-2) * time.Second).UnixNano())
42
	_, err := kv.writer.Exec(
43
		"INSERT INTO `lease_entries` (`owner`, `token`) VALUES (?, ?)",
44
		key, bindUint64AsInt64(expiredToken),
45
	)
46
	as.NoError(err)
47
 
48
	token, err := kv.Acquire(context.Background(), key, time.Second)
49
	as.NoError(err)
50
 
51
	as.NoError(kv.Release(context.Background(), key, token))
52
}
53
 
54
func TestRenewValid(t *testing.T) {
55
	as := assert.New(t)
56
 
57
	kv := testGetKV(t)
58
 
59
	key := make([]byte, 8)
60
	rand.Read(key)
61
 
62
	token, err := kv.Acquire(context.Background(), key, time.Second)
63
	as.NoError(err)
64
 
65
	time.Sleep(time.Millisecond * 100)
66
 
67
	n, err := kv.Renew(context.Background(), key, time.Second*2, token)
68
	as.NoError(err)
69
	as.NotEqual(token, n)
70
 
71
	// no releasing with the wrong token
72
	as.ErrorIs(kv.Release(context.Background(), key, token), chord.ErrKVLeaseExpired)
73
	as.NoError(kv.Release(context.Background(), key, n))
74
}
75
 
76
func TestRenewExpired(t *testing.T) {
77
	as := assert.New(t)
78
 
79
	kv := testGetKV(t)
80
 
81
	key := make([]byte, 8)
82
	rand.Read(key)
83
 
84
	ttl := time.Second
85
 
86
	tk1, err := kv.Acquire(context.Background(), key, ttl)
87
	as.NoError(err)
88
 
89
	time.Sleep(ttl * 2)
90
 
91
	tk2, err := kv.Acquire(context.Background(), key, ttl)
92
	as.NoError(err)
93
 
94
	_, err = kv.Renew(context.Background(), key, ttl, tk1)
95
	as.ErrorIs(err, chord.ErrKVLeaseExpired)
96
 
97
	tk2, err = kv.Renew(context.Background(), key, ttl, tk2)
98
	as.NoError(err)
99
 
100
	// no releasing with the wrong token
101
	as.ErrorIs(kv.Release(context.Background(), key, tk1), chord.ErrKVLeaseExpired)
102
	as.NoError(kv.Release(context.Background(), key, tk2))
103
}
104
 
105
func TestTTLGuard(t *testing.T) {
106
	as := assert.New(t)
107
 
108
	kv := testGetKV(t)
109
 
110
	key := make([]byte, 8)
111
	rand.Read(key)
112
 
113
	ttl := time.Millisecond * 500
114
 
115
	_, err := kv.Acquire(context.Background(), key, ttl)
116
	as.ErrorIs(err, chord.ErrKVLeaseInvalidTTL)
117
 
118
	ttl = time.Second
119
	tk, err := kv.Acquire(context.Background(), key, ttl)
120
	as.NoError(err)
121
 
122
	ttl = time.Millisecond * 500
123
	_, err = kv.Renew(context.Background(), key, ttl, tk)
124
	as.ErrorIs(err, chord.ErrKVLeaseInvalidTTL)
125
 
126
	as.NoError(kv.Release(context.Background(), key, tk))
127
}
128
 
129
// The following tests were collaborated with GPT 4o-mini
130
 
131
func TestLeaseAcquisitionAndRenewal(t *testing.T) {
132
	as := require.New(t)
133
	kv := testGetKV(t)
134
 
135
	ctx := context.Background()
136
	leaseID := []byte("lease_test")
137
	ttl := 1 * time.Second
138
 
139
	// Acquire Lease
140
	token, err := kv.Acquire(ctx, leaseID, ttl)
141
	as.NoError(err)
142
 
143
	time.Sleep(50 * time.Millisecond)
144
 
145
	// Renew Lease before expiration
146
	newToken, err := kv.Renew(ctx, leaseID, ttl, token)
147
	as.NoError(err)
148
	as.Greater(newToken, token) // Ensure token is increasing
149
}
150
 
151
func TestLeaseExpiration(t *testing.T) {
152
	as := require.New(t)
153
	kv := testGetKV(t)
154
 
155
	ctx := context.Background()
156
	leaseID := []byte("lease_test")
157
	ttl := 1 * time.Second
158
 
159
	// Acquire Lease
160
	token, err := kv.Acquire(ctx, leaseID, ttl)
161
	as.NoError(err)
162
 
163
	// Wait for lease to expire
164
	time.Sleep(ttl + 200*time.Millisecond)
165
 
166
	// Attempt to renew expired lease
167
	_, err = kv.Renew(ctx, leaseID, ttl, token)
168
	as.ErrorIs(err, chord.ErrKVLeaseExpired)
169
}
170
 
171
func TestLeaseRenewalWithParallelLoad(t *testing.T) {
172
	as := require.New(t)
173
	kv := testGetKV(t)
174
 
175
	ctx := context.Background()
176
	leaseID := []byte("lease_test")
177
	ttl := 1 * time.Second
178
 
179
	// Acquire Lease
180
	token, err := kv.Acquire(ctx, leaseID, ttl)
181
	as.NoError(err)
182
 
183
	// Start multiple goroutines that try to renew the lease concurrently
184
	const numWorkers = 50
185
	type result struct {
186
		token uint64
187
		err   error
188
	}
189
	resultChan := make(chan result, numWorkers)
190
	for range numWorkers {
191
		go func(tokenSnapshot uint64) {
192
			token, err := kv.Renew(ctx, leaseID, ttl, tokenSnapshot)
193
			resultChan <- result{
194
				token: token,
195
				err:   err,
196
			}
197
		}(token)
198
	}
199
 
200
	// Collect results
201
	var successCount int
202
	for range numWorkers {
203
		result := <-resultChan
204
		if result.err == nil && result.token > 0 {
205
			t.Log("success")
206
			successCount++
207
		} else {
208
			t.Logf("token: %d; err: %+v\n", result.token, result.err)
209
			as.ErrorIs(result.err, chord.ErrKVLeaseExpired)
210
		}
211
	}
212
	as.Equal(1, successCount, "Only one renewal should succeed")
213
}
214
 
215
func TestLeaseExpirationUnderLoad(t *testing.T) {
216
	as := require.New(t)
217
	kv := testGetKV(t)
218
 
219
	ctx := context.Background()
220
	leaseID := []byte("lease_test")
221
	ttl := 1 * time.Second
222
 
223
	// Acquire Lease
224
	token, err := kv.Acquire(ctx, leaseID, ttl)
225
	as.NoError(err)
226
 
227
	// Start concurrent access while waiting for lease expiration
228
	const numWorkers = 50
229
	errChan := make(chan error, numWorkers)
230
	for range numWorkers {
231
		go func(tokenSnapshot uint64) {
232
			time.Sleep(ttl + 200*time.Millisecond) // Wait for expiration
233
			_, err := kv.Renew(ctx, leaseID, ttl, tokenSnapshot)
234
			errChan <- err
235
		}(token)
236
	}
237
 
238
	// Expect all workers to fail renewal due to expiration
239
	for range numWorkers {
240
		err := <-errChan
241
		if err == nil {
242
			t.Errorf("Renewal should have failed, but it succeeded")
243
		} else {
244
			as.ErrorIs(err, chord.ErrKVLeaseExpired, "Expired leases should not be renewable")
245
		}
246
	}
247
}
248
 
249
func TestConcurrentLeaseAcquisition(t *testing.T) {
250
	as := require.New(t)
251
	kv := testGetKV(t)
252
 
253
	ctx := context.Background()
254
	leaseID := []byte("lease_test")
255
	ttl := 1 * time.Second
256
 
257
	const numWorkers = 50
258
	type result struct {
259
		token uint64
260
		err   error
261
	}
262
	resultChan := make(chan result, numWorkers)
263
 
264
	for range numWorkers {
265
		go func() {
266
			token, err := kv.Acquire(ctx, leaseID, ttl)
267
			resultChan <- result{
268
				token: token,
269
				err:   err,
270
			}
271
		}()
272
	}
273
 
274
	var successCount int
275
	for range numWorkers {
276
		result := <-resultChan
277
		if result.err == nil && result.token > 0 {
278
			t.Log("success")
279
			successCount++
280
		} else {
281
			t.Logf("token: %d; err: %+v\n", result.token, result.err)
282
			as.ErrorIs(result.err, chord.ErrKVLeaseConflict) // Ensure others fail
283
		}
284
	}
285
 
286
	as.Equal(1, successCount, "Only one acquisition should succeed")
287
}