Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package resource
18

19
import (
20
	"errors"
21
	"reflect"
22
	"testing"
23
	"time"
24

25
	"github.com/stretchr/testify/assert"
26
	gomock "go.uber.org/mock/gomock"
27

28
	commonv2 "d7y.io/api/v2/pkg/apis/common/v2"
29

30
	"d7y.io/dragonfly/v2/pkg/container/set"
31
	"d7y.io/dragonfly/v2/pkg/gc"
32
	"d7y.io/dragonfly/v2/scheduler/config"
33
)
34

35
var (
36
	mockHostGCConfig = &config.GCConfig{
37
		HostGCInterval: 1 * time.Second,
38
	}
39
)
40

41
func TestHostManager_newHostManager(t *testing.T) {
42
	tests := []struct {
43
		name   string
44
		mock   func(m *gc.MockGCMockRecorder)
45
		expect func(t *testing.T, hostManager HostManager, err error)
46
	}{
47
		{
48
			name: "new host manager",
49
			mock: func(m *gc.MockGCMockRecorder) {
50
				m.Add(gomock.Any()).Return(nil).Times(1)
51
			},
52
			expect: func(t *testing.T, hostManager HostManager, err error) {
53
				assert := assert.New(t)
54
				assert.Equal(reflect.TypeOf(hostManager).Elem().Name(), "hostManager")
55
			},
56
		},
57
		{
58
			name: "new host manager failed because of gc error",
59
			mock: func(m *gc.MockGCMockRecorder) {
60
				m.Add(gomock.Any()).Return(errors.New("foo")).Times(1)
61
			},
62
			expect: func(t *testing.T, hostManager HostManager, err error) {
63
				assert := assert.New(t)
64
				assert.EqualError(err, "foo")
65
			},
66
		},
67
	}
68

69
	for _, tc := range tests {
70
		t.Run(tc.name, func(t *testing.T) {
71
			ctl := gomock.NewController(t)
72
			defer ctl.Finish()
73
			gc := gc.NewMockGC(ctl)
74
			tc.mock(gc.EXPECT())
75
			hostManager, err := newHostManager(mockHostGCConfig, gc)
76

77
			tc.expect(t, hostManager, err)
78
		})
79
	}
80
}
81

82
func TestHostManager_Load(t *testing.T) {
83
	tests := []struct {
84
		name   string
85
		mock   func(m *gc.MockGCMockRecorder)
86
		expect func(t *testing.T, hostManager HostManager, mockHost *Host)
87
	}{
88
		{
89
			name: "load host",
90
			mock: func(m *gc.MockGCMockRecorder) {
91
				m.Add(gomock.Any()).Return(nil).Times(1)
92
			},
93
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
94
				assert := assert.New(t)
95
				hostManager.Store(mockHost)
96
				host, loaded := hostManager.Load(mockHost.ID)
97
				assert.Equal(loaded, true)
98
				assert.Equal(host.ID, mockHost.ID)
99
			},
100
		},
101
		{
102
			name: "host does not exist",
103
			mock: func(m *gc.MockGCMockRecorder) {
104
				m.Add(gomock.Any()).Return(nil).Times(1)
105
			},
106
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
107
				assert := assert.New(t)
108
				_, loaded := hostManager.Load(mockHost.ID)
109
				assert.Equal(loaded, false)
110
			},
111
		},
112
		{
113
			name: "load key is empty",
114
			mock: func(m *gc.MockGCMockRecorder) {
115
				m.Add(gomock.Any()).Return(nil).Times(1)
116
			},
117
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
118
				assert := assert.New(t)
119
				mockHost.ID = ""
120
				hostManager.Store(mockHost)
121
				host, loaded := hostManager.Load(mockHost.ID)
122
				assert.Equal(loaded, true)
123
				assert.Equal(host.ID, mockHost.ID)
124
			},
125
		},
126
	}
127

128
	for _, tc := range tests {
129
		t.Run(tc.name, func(t *testing.T) {
130
			ctl := gomock.NewController(t)
131
			defer ctl.Finish()
132
			gc := gc.NewMockGC(ctl)
133
			tc.mock(gc.EXPECT())
134

135
			mockHost := NewHost(
136
				mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
137
				mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type)
138
			hostManager, err := newHostManager(mockHostGCConfig, gc)
139
			if err != nil {
140
				t.Fatal(err)
141
			}
142

143
			tc.expect(t, hostManager, mockHost)
144
		})
145
	}
146
}
147

148
func TestHostManager_Store(t *testing.T) {
149
	tests := []struct {
150
		name   string
151
		mock   func(m *gc.MockGCMockRecorder)
152
		expect func(t *testing.T, hostManager HostManager, mockHost *Host)
153
	}{
154
		{
155
			name: "store host",
156
			mock: func(m *gc.MockGCMockRecorder) {
157
				m.Add(gomock.Any()).Return(nil).Times(1)
158
			},
159
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
160
				assert := assert.New(t)
161
				hostManager.Store(mockHost)
162
				host, loaded := hostManager.Load(mockHost.ID)
163
				assert.Equal(loaded, true)
164
				assert.Equal(host.ID, mockHost.ID)
165
			},
166
		},
167
		{
168
			name: "store key is empty",
169
			mock: func(m *gc.MockGCMockRecorder) {
170
				m.Add(gomock.Any()).Return(nil).Times(1)
171
			},
172
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
173
				assert := assert.New(t)
174
				mockHost.ID = ""
175
				hostManager.Store(mockHost)
176
				host, loaded := hostManager.Load(mockHost.ID)
177
				assert.Equal(loaded, true)
178
				assert.Equal(host.ID, mockHost.ID)
179
			},
180
		},
181
	}
182

183
	for _, tc := range tests {
184
		t.Run(tc.name, func(t *testing.T) {
185
			ctl := gomock.NewController(t)
186
			defer ctl.Finish()
187
			gc := gc.NewMockGC(ctl)
188
			tc.mock(gc.EXPECT())
189

190
			mockHost := NewHost(
191
				mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
192
				mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type)
193
			hostManager, err := newHostManager(mockHostGCConfig, gc)
194
			if err != nil {
195
				t.Fatal(err)
196
			}
197

198
			tc.expect(t, hostManager, mockHost)
199
		})
200
	}
201
}
202

203
func TestHostManager_LoadOrStore(t *testing.T) {
204
	tests := []struct {
205
		name   string
206
		mock   func(m *gc.MockGCMockRecorder)
207
		expect func(t *testing.T, hostManager HostManager, mockHost *Host)
208
	}{
209
		{
210
			name: "load host exist",
211
			mock: func(m *gc.MockGCMockRecorder) {
212
				m.Add(gomock.Any()).Return(nil).Times(1)
213
			},
214
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
215
				assert := assert.New(t)
216
				hostManager.Store(mockHost)
217
				host, loaded := hostManager.LoadOrStore(mockHost)
218
				assert.Equal(loaded, true)
219
				assert.Equal(host.ID, mockHost.ID)
220
			},
221
		},
222
		{
223
			name: "load host does not exist",
224
			mock: func(m *gc.MockGCMockRecorder) {
225
				m.Add(gomock.Any()).Return(nil).Times(1)
226
			},
227
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
228
				assert := assert.New(t)
229
				host, loaded := hostManager.LoadOrStore(mockHost)
230
				assert.Equal(loaded, false)
231
				assert.Equal(host.ID, mockHost.ID)
232
			},
233
		},
234
	}
235

236
	for _, tc := range tests {
237
		t.Run(tc.name, func(t *testing.T) {
238
			ctl := gomock.NewController(t)
239
			defer ctl.Finish()
240
			gc := gc.NewMockGC(ctl)
241
			tc.mock(gc.EXPECT())
242

243
			mockHost := NewHost(
244
				mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
245
				mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type)
246
			hostManager, err := newHostManager(mockHostGCConfig, gc)
247
			if err != nil {
248
				t.Fatal(err)
249
			}
250

251
			tc.expect(t, hostManager, mockHost)
252
		})
253
	}
254
}
255

256
func TestHostManager_Delete(t *testing.T) {
257
	tests := []struct {
258
		name   string
259
		mock   func(m *gc.MockGCMockRecorder)
260
		expect func(t *testing.T, hostManager HostManager, mockHost *Host)
261
	}{
262
		{
263
			name: "delete host",
264
			mock: func(m *gc.MockGCMockRecorder) {
265
				m.Add(gomock.Any()).Return(nil).Times(1)
266
			},
267
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
268
				assert := assert.New(t)
269
				hostManager.Store(mockHost)
270
				hostManager.Delete(mockHost.ID)
271
				_, loaded := hostManager.Load(mockHost.ID)
272
				assert.Equal(loaded, false)
273
			},
274
		},
275
		{
276
			name: "delete key does not exist",
277
			mock: func(m *gc.MockGCMockRecorder) {
278
				m.Add(gomock.Any()).Return(nil).Times(1)
279
			},
280
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host) {
281
				assert := assert.New(t)
282
				mockHost.ID = ""
283
				hostManager.Store(mockHost)
284
				hostManager.Delete(mockHost.ID)
285
				_, loaded := hostManager.Load(mockHost.ID)
286
				assert.Equal(loaded, false)
287
			},
288
		},
289
	}
290

291
	for _, tc := range tests {
292
		t.Run(tc.name, func(t *testing.T) {
293
			ctl := gomock.NewController(t)
294
			defer ctl.Finish()
295
			gc := gc.NewMockGC(ctl)
296
			tc.mock(gc.EXPECT())
297

298
			mockHost := NewHost(
299
				mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
300
				mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type)
301
			hostManager, err := newHostManager(mockHostGCConfig, gc)
302
			if err != nil {
303
				t.Fatal(err)
304
			}
305

306
			tc.expect(t, hostManager, mockHost)
307
		})
308
	}
309
}
310

311
func TestHostManager_LoadRandomHosts(t *testing.T) {
312
	tests := []struct {
313
		name   string
314
		hosts  []*Host
315
		mock   func(m *gc.MockGCMockRecorder)
316
		expect func(t *testing.T, hm HostManager, hosts []*Host)
317
	}{
318
		{
319
			name: "load random hosts",
320
			hosts: []*Host{
321
				NewHost(
322
					mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
323
					mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type),
324
				NewHost(
325
					mockRawSeedHost.ID, mockRawSeedHost.IP, mockRawSeedHost.Hostname,
326
					mockRawSeedHost.Port, mockRawSeedHost.DownloadPort, mockRawSeedHost.Type),
327
			},
328
			mock: func(m *gc.MockGCMockRecorder) {
329
				m.Add(gomock.Any()).Return(nil).Times(1)
330
			},
331
			expect: func(t *testing.T, hm HostManager, hosts []*Host) {
332
				assert := assert.New(t)
333
				for _, host := range hosts {
334
					hm.Store(host)
335
				}
336

337
				blocklist := set.NewSafeSet[string]()
338
				blocklist.Add(mockRawSeedHost.ID)
339
				h := hm.LoadRandomHosts(2, blocklist)
340
				assert.Equal(len(h), 1)
341
			},
342
		},
343
		{
344
			name: "load random hosts when the load number is 0",
345
			hosts: []*Host{
346
				NewHost(
347
					mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
348
					mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type),
349
				NewHost(
350
					mockRawSeedHost.ID, mockRawSeedHost.IP, mockRawSeedHost.Hostname,
351
					mockRawSeedHost.Port, mockRawSeedHost.DownloadPort, mockRawSeedHost.Type),
352
			},
353
			mock: func(m *gc.MockGCMockRecorder) {
354
				m.Add(gomock.Any()).Return(nil).Times(1)
355
			},
356
			expect: func(t *testing.T, hm HostManager, hosts []*Host) {
357
				assert := assert.New(t)
358
				for _, host := range hosts {
359
					hm.Store(host)
360
				}
361

362
				blocklist := set.NewSafeSet[string]()
363
				blocklist.Add(mockRawSeedHost.ID)
364
				h := hm.LoadRandomHosts(0, blocklist)
365
				assert.Equal(len(h), 0)
366
			},
367
		},
368
		{
369
			name:  "map is empty",
370
			hosts: []*Host{},
371
			mock: func(m *gc.MockGCMockRecorder) {
372
				m.Add(gomock.Any()).Return(nil).Times(1)
373
			},
374
			expect: func(t *testing.T, hm HostManager, hosts []*Host) {
375
				assert := assert.New(t)
376
				for _, host := range hosts {
377
					hm.Store(host)
378
				}
379

380
				blocklist := set.NewSafeSet[string]()
381
				blocklist.Add(mockRawSeedHost.ID)
382
				h := hm.LoadRandomHosts(1, blocklist)
383
				assert.Equal(len(h), 0)
384
			},
385
		},
386
		{
387
			name: "the number of hosts in the map is insufficient",
388
			hosts: []*Host{
389
				NewHost(
390
					mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
391
					mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type),
392
				NewHost(
393
					mockRawSeedHost.ID, mockRawSeedHost.IP, mockRawSeedHost.Hostname,
394
					mockRawSeedHost.Port, mockRawSeedHost.DownloadPort, mockRawSeedHost.Type),
395
			},
396
			mock: func(m *gc.MockGCMockRecorder) {
397
				m.Add(gomock.Any()).Return(nil).Times(1)
398
			},
399
			expect: func(t *testing.T, hm HostManager, hosts []*Host) {
400
				assert := assert.New(t)
401
				for _, host := range hosts {
402
					hm.Store(host)
403
				}
404

405
				blocklist := set.NewSafeSet[string]()
406
				blocklist.Add(mockRawSeedHost.ID)
407
				h := hm.LoadRandomHosts(3, blocklist)
408
				assert.Equal(len(h), 1)
409
			},
410
		},
411
	}
412

413
	for _, tc := range tests {
414
		t.Run(tc.name, func(t *testing.T) {
415
			ctl := gomock.NewController(t)
416
			defer ctl.Finish()
417
			gc := gc.NewMockGC(ctl)
418
			tc.mock(gc.EXPECT())
419

420
			hm, err := newHostManager(mockHostGCConfig, gc)
421
			if err != nil {
422
				t.Fatal(err)
423
			}
424

425
			tc.expect(t, hm, tc.hosts)
426
		})
427
	}
428
}
429

430
func TestHostManager_RunGC(t *testing.T) {
431
	tests := []struct {
432
		name   string
433
		mock   func(m *gc.MockGCMockRecorder)
434
		expect func(t *testing.T, hostManager HostManager, mockHost *Host, mockPeer *Peer)
435
	}{
436
		{
437
			name: "host reclaimed",
438
			mock: func(m *gc.MockGCMockRecorder) {
439
				m.Add(gomock.Any()).Return(nil).Times(1)
440
			},
441
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host, mockPeer *Peer) {
442
				assert := assert.New(t)
443
				hostManager.Store(mockHost)
444
				err := hostManager.RunGC()
445
				assert.NoError(err)
446

447
				_, loaded := hostManager.Load(mockHost.ID)
448
				assert.Equal(loaded, false)
449
			},
450
		},
451
		{
452
			name: "host has peers",
453
			mock: func(m *gc.MockGCMockRecorder) {
454
				m.Add(gomock.Any()).Return(nil).Times(1)
455
			},
456
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host, mockPeer *Peer) {
457
				assert := assert.New(t)
458
				hostManager.Store(mockHost)
459
				mockHost.StorePeer(mockPeer)
460
				err := hostManager.RunGC()
461
				assert.NoError(err)
462

463
				host, loaded := hostManager.Load(mockHost.ID)
464
				assert.Equal(loaded, true)
465
				assert.Equal(host.ID, mockHost.ID)
466
			},
467
		},
468
		{
469
			name: "host has upload peers",
470
			mock: func(m *gc.MockGCMockRecorder) {
471
				m.Add(gomock.Any()).Return(nil).Times(1)
472
			},
473
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host, mockPeer *Peer) {
474
				assert := assert.New(t)
475
				hostManager.Store(mockHost)
476
				mockHost.StorePeer(mockPeer)
477
				mockHost.PeerCount.Add(0)
478
				mockHost.ConcurrentUploadCount.Add(1)
479
				err := hostManager.RunGC()
480
				assert.NoError(err)
481

482
				host, loaded := hostManager.Load(mockHost.ID)
483
				assert.Equal(loaded, true)
484
				assert.Equal(host.ID, mockHost.ID)
485
			},
486
		},
487
		{
488
			name: "host is seed peer",
489
			mock: func(m *gc.MockGCMockRecorder) {
490
				m.Add(gomock.Any()).Return(nil).Times(1)
491
			},
492
			expect: func(t *testing.T, hostManager HostManager, mockHost *Host, mockPeer *Peer) {
493
				assert := assert.New(t)
494
				mockSeedHost := NewHost(
495
					mockRawSeedHost.ID, mockRawSeedHost.IP, mockRawSeedHost.Hostname,
496
					mockRawSeedHost.Port, mockRawSeedHost.DownloadPort, mockRawSeedHost.Type)
497
				hostManager.Store(mockSeedHost)
498
				err := hostManager.RunGC()
499
				assert.NoError(err)
500

501
				host, loaded := hostManager.Load(mockSeedHost.ID)
502
				assert.Equal(loaded, true)
503
				assert.Equal(host.ID, mockSeedHost.ID)
504
			},
505
		},
506
	}
507

508
	for _, tc := range tests {
509
		t.Run(tc.name, func(t *testing.T) {
510
			ctl := gomock.NewController(t)
511
			defer ctl.Finish()
512
			gc := gc.NewMockGC(ctl)
513
			tc.mock(gc.EXPECT())
514

515
			mockHost := NewHost(
516
				mockRawHost.ID, mockRawHost.IP, mockRawHost.Hostname,
517
				mockRawHost.Port, mockRawHost.DownloadPort, mockRawHost.Type)
518
			mockTask := NewTask(mockTaskID, mockTaskURL, mockTaskTag, mockTaskApplication, commonv2.TaskType_DFDAEMON, mockTaskFilteredQueryParams, mockTaskHeader, mockTaskBackToSourceLimit)
519
			mockPeer := NewPeer(mockPeerID, mockResourceConfig, mockTask, mockHost)
520
			hostManager, err := newHostManager(mockHostGCConfig, gc)
521
			if err != nil {
522
				t.Fatal(err)
523
			}
524

525
			tc.expect(t, hostManager, mockHost, mockPeer)
526
		})
527
	}
528
}
529