@@ -4,6 +4,7 @@ package utils_test
4
4
5
5
import (
6
6
"context"
7
+ "math"
7
8
"sync"
8
9
"testing"
9
10
@@ -14,7 +15,7 @@ import (
14
15
15
16
func TestCheckGCSafepoint (t * testing.T ) {
16
17
ctx := context .Background ()
17
- pdClient := & mockSafePoint {safepoint : 2333 }
18
+ pdClient := & mockSafePoint {safepoint : 2333 , services : make ( map [ string ] uint64 ) }
18
19
{
19
20
err := utils .CheckGCSafePoint (ctx , pdClient , 2333 + 1 )
20
21
require .NoError (t , err )
@@ -34,23 +35,78 @@ func TestCheckGCSafepoint(t *testing.T) {
34
35
}
35
36
}
36
37
38
+ func TestCheckUpdateServiceSafepoint (t * testing.T ) {
39
+ ctx := context .Background ()
40
+ pdClient := & mockSafePoint {safepoint : 2333 , services : make (map [string ]uint64 )}
41
+ {
42
+ // nothing happened, because current safepoint is large than servicee safepoint.
43
+ err := utils .UpdateServiceSafePoint (ctx , pdClient , utils.BRServiceSafePoint {
44
+ "BR_SERVICE" ,
45
+ 1 ,
46
+ 1 ,
47
+ })
48
+ require .NoError (t , err )
49
+ curSafePoint , err := pdClient .UpdateGCSafePoint (ctx , 0 )
50
+ require .NoError (t , err )
51
+ require .Equal (t , uint64 (2333 ), curSafePoint )
52
+ }
53
+ {
54
+ // register br service safepoint
55
+ err := utils .UpdateServiceSafePoint (ctx , pdClient , utils.BRServiceSafePoint {
56
+ "BR_SERVICE" ,
57
+ 1 ,
58
+ 2334 ,
59
+ })
60
+ require .NoError (t , err )
61
+ curSafePoint , find := pdClient .GetServiceSafePoint ("BR_SERVICE" )
62
+ // update with new safepoint - 1.
63
+ require .Equal (t , uint64 (2333 ), curSafePoint )
64
+ require .True (t , find )
65
+ }
66
+ {
67
+ // remove br service safepoint
68
+ err := utils .UpdateServiceSafePoint (ctx , pdClient , utils.BRServiceSafePoint {
69
+ "BR_SERVICE" ,
70
+ 0 ,
71
+ math .MaxUint64 ,
72
+ })
73
+ require .NoError (t , err )
74
+ _ , find := pdClient .GetServiceSafePoint ("BR_SERVICE" )
75
+ require .False (t , find )
76
+ }
77
+ }
78
+
37
79
type mockSafePoint struct {
38
80
sync.Mutex
39
81
pd.Client
82
+ services map [string ]uint64
40
83
safepoint uint64
41
84
minServiceSafepoint uint64
42
85
}
43
86
87
+ func (m * mockSafePoint ) GetServiceSafePoint (serviceID string ) (uint64 , bool ) {
88
+ m .Lock ()
89
+ defer m .Unlock ()
90
+ safepoint , ok := m .services [serviceID ]
91
+ return safepoint , ok
92
+ }
93
+
44
94
func (m * mockSafePoint ) UpdateServiceGCSafePoint (ctx context.Context , serviceID string , ttl int64 , safePoint uint64 ) (uint64 , error ) {
45
95
m .Lock ()
46
96
defer m .Unlock ()
47
97
98
+ if ttl <= 0 {
99
+ delete (m .services , serviceID )
100
+ return 0 , nil
101
+ }
102
+
48
103
if m .safepoint > safePoint {
49
104
return m .safepoint , nil
50
105
}
51
106
if m .minServiceSafepoint == 0 || m .minServiceSafepoint > safePoint {
52
107
m .minServiceSafepoint = safePoint
53
108
}
109
+ m .services [serviceID ] = safePoint
54
110
return m .minServiceSafepoint , nil
55
111
}
56
112
@@ -65,7 +121,7 @@ func (m *mockSafePoint) UpdateGCSafePoint(ctx context.Context, safePoint uint64)
65
121
}
66
122
67
123
func TestStartServiceSafePointKeeper (t * testing.T ) {
68
- pdClient := & mockSafePoint {safepoint : 2333 }
124
+ pdClient := & mockSafePoint {safepoint : 2333 , services : make ( map [ string ] uint64 ) }
69
125
70
126
cases := []struct {
71
127
sp utils.BRServiceSafePoint
0 commit comments