@@ -36,21 +36,68 @@ package tikv_test
3636
3737import (
3838 "context"
39+ "sync"
3940 "testing"
4041 "time"
4142
43+ "github.com/google/uuid"
44+ "github.com/prometheus/client_golang/prometheus"
45+ dto "github.com/prometheus/client_model/go"
4246 "github.com/stretchr/testify/suite"
47+ tikverr "github.com/tikv/client-go/v2/error"
48+ "github.com/tikv/client-go/v2/metrics"
4349 "github.com/tikv/client-go/v2/oracle"
4450 "github.com/tikv/client-go/v2/tikv"
51+ "github.com/tikv/client-go/v2/txnkv/transaction"
52+ "github.com/tikv/client-go/v2/util"
53+ "github.com/tikv/client-go/v2/util/intest"
4554)
4655
56+ type mockCommitTSOracle struct {
57+ oracle.Oracle
58+ mu struct {
59+ sync.Mutex
60+ startTS uint64
61+ commitTSOffsets []uint64
62+ }
63+ }
64+
65+ func (o * mockCommitTSOracle ) ResetMock (startTS uint64 , commitTSOffsets []uint64 ) {
66+ o .mu .Lock ()
67+ defer o .mu .Unlock ()
68+ o .mu .startTS = startTS
69+ o .mu .commitTSOffsets = commitTSOffsets
70+ }
71+
72+ func (o * mockCommitTSOracle ) GetTimestamp (ctx context.Context , option * oracle.Option ) (uint64 , error ) {
73+ if ts := ctx .Value (transaction .CtxInGetTimestampForCommitKey ); ts != nil {
74+ o .mu .Lock ()
75+ startTS , ok := ts .(uint64 )
76+ if ! ok || startTS != o .mu .startTS {
77+ o .mu .Unlock ()
78+ } else {
79+ defer o .mu .Unlock ()
80+ if len (o .mu .commitTSOffsets ) == 0 {
81+ panic ("empty mock oracle ts set" )
82+ }
83+ offset := o .mu .commitTSOffsets [0 ]
84+ o .mu .commitTSOffsets = o .mu .commitTSOffsets [1 :]
85+ return o .mu .startTS + offset , nil
86+ }
87+ }
88+ return o .Oracle .GetTimestamp (ctx , option )
89+ }
90+
4791type testOptionSuite struct {
4892 suite.Suite
49- store tikv.StoreProbe
93+ mockOracle * mockCommitTSOracle
94+ store tikv.StoreProbe
5095}
5196
5297func (s * testOptionSuite ) SetupSuite () {
5398 s .store = tikv.StoreProbe {KVStore : NewTestStore (s .T ())}
99+ s .mockOracle = & mockCommitTSOracle {Oracle : s .store .GetOracle ()}
100+ s .store .SetOracle (s .mockOracle )
54101}
55102
56103func (s * testOptionSuite ) TearDownSuite () {
@@ -61,36 +108,128 @@ func TestOption(t *testing.T) {
61108 suite .Run (t , new (testOptionSuite ))
62109}
63110
111+ func (s * testOptionSuite ) MetricValues (m any ) * dto.Metric {
112+ metric := & dto.Metric {}
113+ s .NoError (m .(prometheus.Metric ).Write (metric ))
114+ return metric
115+ }
116+
117+ func (s * testOptionSuite ) GetHistogramMetricSampleCount (m any ) uint64 {
118+ return s .MetricValues (m ).GetHistogram ().GetSampleCount ()
119+ }
120+
64121func (s * testOptionSuite ) TestSetCommitWaitUntilTSO () {
65- getTS1 := func (startTS uint64 ) uint64 { return 100 }
66- getTS2 := func (startTS uint64 ) uint64 {
67- return startTS + 200
122+ if * withTiKV {
123+ s .T ().Skip ("TestSetCommitWaitUntilTSO only runs on local mock storage" )
68124 }
69- getTS3 := func (startTS uint64 ) uint64 {
70- now := oracle .GetTimeFromTS (startTS )
71- return oracle .GoTimeToTS (now .Add (10 * time .Second ))
72- }
73- type testCase struct {
74- getTS func (startTS uint64 ) uint64
75- noError bool
125+ origInTest := intest .InTest
126+ defer func () {
127+ intest .InTest = origInTest
128+ }()
129+ intest .InTest = true
130+ doWithCollectHistSamplesInc := func (do func () error , metrics []any ) ([]uint64 , error ) {
131+ cnt := make ([]uint64 , len (metrics ))
132+ for i , metric := range metrics {
133+ cnt [i ] = s .GetHistogramMetricSampleCount (metric )
134+ }
135+ err := do ()
136+ for i , metric := range metrics {
137+ cnt [i ] = s .GetHistogramMetricSampleCount (metric ) - cnt [i ]
138+ }
139+ return cnt , err
76140 }
77141
78- for _ , tc := range []testCase {
79- {getTS1 , true },
80- {getTS2 , true },
81- {getTS3 , false },
142+ for _ , tc := range []struct {
143+ name string
144+ // We will call `txn.SetCommitWaitUntilTSO(commitWaitTSO + txn.StartTS())`
145+ commitWaitTSO uint64
146+ // Mock the PD TSOs in commit phase as {
147+ // mockCommitTSO[0] + txn.StartTS(), // first attempt
148+ // mockCommitTSO[1] + txn.StartTS(), // second attempt
149+ // ...
150+ // }
151+ mockCommitTSO []uint64
152+ setWaitTimeout time.Duration
153+ err bool
154+ }{
155+ {
156+ name : "no lag commit ts" ,
157+ commitWaitTSO : 1 ,
158+ mockCommitTSO : []uint64 {100 },
159+ },
160+ {
161+ name : "lag, retry once and success" ,
162+ commitWaitTSO : 200 ,
163+ mockCommitTSO : []uint64 {100 , 201 },
164+ },
165+ {
166+ name : "lag, retry twice and success" ,
167+ commitWaitTSO : 300 ,
168+ mockCommitTSO : []uint64 {100 , 200 , 301 },
169+ },
170+ {
171+ name : "lag too much, fail directly" ,
172+ commitWaitTSO : oracle .ComposeTS (10000 , 0 ),
173+ mockCommitTSO : []uint64 {100 },
174+ err : true ,
175+ },
176+ {
177+ name : "lag, retry but timeout" ,
178+ commitWaitTSO : 100 ,
179+ mockCommitTSO : []uint64 {10 , 20 },
180+ setWaitTimeout : time .Millisecond ,
181+ err : true ,
182+ },
82183 } {
83- txn , err := s .store .Begin ()
84- s .NoError (err )
85- commitUntil := tc .getTS (txn .StartTS ())
86- txn .KVTxn .SetCommitWaitUntilTSO (commitUntil )
87- s .NoError (txn .Set ([]byte ("somekey" ), []byte ("somevalue" )))
88- err = txn .Commit (context .Background ())
89- s .Equal (tc .noError , err == nil )
90-
91- if tc .noError {
92- s .Greater (txn .CommitTS (), uint64 (100 ))
93- }
184+ s .Run (tc .name , func () {
185+ txn , err := s .store .Begin ()
186+ s .NoError (err )
187+ s .NoError (txn .Set ([]byte ("somekey:" + uuid .NewString ()), []byte ("somevalue" )))
188+
189+ txn .SetCommitWaitUntilTSO (tc .commitWaitTSO + txn .StartTS ())
190+ if tc .setWaitTimeout > 0 {
191+ txn .SetCommitWaitUntilTSOTimeout (tc .setWaitTimeout )
192+ }
193+ var commitDetail * util.CommitDetails
194+ inc , err := doWithCollectHistSamplesInc (func () error {
195+ s .mockOracle .ResetMock (txn .StartTS (), tc .mockCommitTSO )
196+ defer s .mockOracle .ResetMock (0 , nil )
197+ return txn .Commit (context .WithValue (context .Background (), util .CommitDetailCtxKey , & commitDetail ))
198+ }, []any {
199+ // ok metrics
200+ metrics .LagCommitTSWaitHistogramWithOK ,
201+ metrics .LagCommitTSAttemptHistogramWithOK ,
202+ // err metrics
203+ metrics .LagCommitTSWaitHistogramWithError ,
204+ metrics .LagCommitTSAttemptHistogramWithError ,
205+ })
206+
207+ s .Equal (txn .StartTS ()+ tc .commitWaitTSO , txn .GetCommitWaitUntilTSO ())
208+ if ! tc .err {
209+ s .NoError (err )
210+ s .Equal (txn .CommitTS (), txn .StartTS ()+ tc .mockCommitTSO [len (tc .mockCommitTSO )- 1 ])
211+ if len (tc .mockCommitTSO ) == 1 {
212+ // if fetch TSO once and then success, there is no lag metrics
213+ s .Equal ([]uint64 {0 , 0 , 0 , 0 }, inc )
214+ s .Equal (util.CommitTSLagDetails {}, commitDetail .LagDetails )
215+ } else {
216+ s .Equal ([]uint64 {1 , 1 , 0 , 0 }, inc )
217+ s .NotZero (txn .KVTxn .GetCommitWaitUntilTSO ())
218+ lagWaitTime := commitDetail .LagDetails .WaitTime
219+ s .Positive (lagWaitTime )
220+ s .Equal (util.CommitTSLagDetails {
221+ WaitTime : lagWaitTime ,
222+ BackoffCnt : len (tc .mockCommitTSO ) - 1 ,
223+ FirstLagTS : txn .StartTS () + tc .mockCommitTSO [0 ],
224+ WaitUntilTS : txn .GetCommitWaitUntilTSO (),
225+ }, commitDetail .LagDetails )
226+ }
227+ } else {
228+ s .Error (err )
229+ s .True (tikverr .IsErrorCommitTSLag (err ))
230+ s .Equal ([]uint64 {0 , 0 , 1 , 1 }, inc )
231+ }
232+ })
94233 }
95234}
96235
0 commit comments