@@ -11,7 +11,7 @@ use crate::{
11
11
use derive_new:: new;
12
12
use fail:: fail_point;
13
13
use futures:: { prelude:: * , stream:: BoxStream } ;
14
- use std:: { iter, ops:: RangeBounds , sync:: Arc } ;
14
+ use std:: { iter, ops:: RangeBounds , sync:: Arc , time :: Instant } ;
15
15
use tikv_client_proto:: { kvrpcpb, pdpb:: Timestamp } ;
16
16
use tokio:: { sync:: RwLock , time:: Duration } ;
17
17
@@ -60,6 +60,7 @@ pub struct Transaction<PdC: PdClient = PdRpcClient> {
60
60
rpc : Arc < PdC > ,
61
61
options : TransactionOptions ,
62
62
is_heartbeat_started : bool ,
63
+ start_instant : Instant ,
63
64
}
64
65
65
66
impl < PdC : PdClient > Transaction < PdC > {
@@ -80,6 +81,7 @@ impl<PdC: PdClient> Transaction<PdC> {
80
81
rpc,
81
82
options,
82
83
is_heartbeat_started : false ,
84
+ start_instant : std:: time:: Instant :: now ( ) ,
83
85
}
84
86
}
85
87
@@ -550,6 +552,8 @@ impl<PdC: PdClient> Transaction<PdC> {
550
552
self . timestamp . clone ( ) ,
551
553
self . rpc . clone ( ) ,
552
554
self . options . clone ( ) ,
555
+ self . buffer . get_write_size ( ) as u64 ,
556
+ self . start_instant ,
553
557
)
554
558
. commit ( )
555
559
. await ;
@@ -603,6 +607,8 @@ impl<PdC: PdClient> Transaction<PdC> {
603
607
self . timestamp . clone ( ) ,
604
608
self . rpc . clone ( ) ,
605
609
self . options . clone ( ) ,
610
+ self . buffer . get_write_size ( ) as u64 ,
611
+ self . start_instant ,
606
612
)
607
613
. rollback ( )
608
614
. await ;
@@ -617,6 +623,7 @@ impl<PdC: PdClient> Transaction<PdC> {
617
623
/// Send a heart beat message to keep the transaction alive on the server and update its TTL.
618
624
///
619
625
/// Returns the TTL set on the transaction's locks by TiKV.
626
+ #[ doc( hidden) ]
620
627
pub async fn send_heart_beat ( & mut self ) -> Result < u64 > {
621
628
self . check_allow_operation ( ) . await ?;
622
629
let primary_key = match self . buffer . get_primary_key ( ) {
@@ -745,7 +752,7 @@ impl<PdC: PdClient> Transaction<PdC> {
745
752
}
746
753
747
754
async fn start_auto_heartbeat ( & mut self ) {
748
- if !self . options . auto_heartbeat || self . is_heartbeat_started {
755
+ if !self . options . heartbeat_option . is_auto_heartbeat ( ) || self . is_heartbeat_started {
749
756
return ;
750
757
}
751
758
self . is_heartbeat_started = true ;
@@ -758,10 +765,14 @@ impl<PdC: PdClient> Transaction<PdC> {
758
765
let start_ts = self . timestamp . clone ( ) ;
759
766
let region_backoff = self . options . retry_options . region_backoff . clone ( ) ;
760
767
let rpc = self . rpc . clone ( ) ;
768
+ let heartbeat_interval = match self . options . heartbeat_option {
769
+ HeartbeatOption :: NoHeartbeat => DEFAULT_HEARTBEAT_INTERVAL ,
770
+ HeartbeatOption :: FixedTime ( heartbeat_interval) => heartbeat_interval,
771
+ } ;
761
772
762
773
let heartbeat_task = async move {
763
774
loop {
764
- tokio:: time:: sleep ( DEFAULT_HEARTBEAT_INTERVAL ) . await ;
775
+ tokio:: time:: sleep ( heartbeat_interval ) . await ;
765
776
{
766
777
let status = status. read ( ) . await ;
767
778
if matches ! (
@@ -819,6 +830,17 @@ impl<PdC: PdClient> Drop for Transaction<PdC> {
819
830
}
820
831
}
821
832
833
+ /// The default max TTL of a lock in milliseconds. Also called `ManagedLockTTL` in TiDB.
834
+ const MAX_TTL : u64 = 20000 ;
835
+ /// The default TTL of a lock in milliseconds.
836
+ const DEFAULT_LOCK_TTL : u64 = 3000 ;
837
+ /// The default heartbeat interval
838
+ const DEFAULT_HEARTBEAT_INTERVAL : Duration = Duration :: from_millis ( MAX_TTL / 2 ) ;
839
+ /// TiKV recommends each RPC packet should be less than around 1MB. We keep KV size of
840
+ /// each request below 16KB.
841
+ const TXN_COMMIT_BATCH_SIZE : u64 = 16 * 1024 ;
842
+ const TTL_FACTOR : f64 = 6000.0 ;
843
+
822
844
/// Optimistic or pessimistic transaction.
823
845
#[ derive( Clone , PartialEq , Debug ) ]
824
846
pub enum TransactionKind {
@@ -844,8 +866,14 @@ pub struct TransactionOptions {
844
866
retry_options : RetryOptions ,
845
867
/// What to do if the transaction is dropped without an attempt to commit or rollback
846
868
check_level : CheckLevel ,
847
- /// Whether heartbeat will be sent automatically
848
- auto_heartbeat : bool ,
869
+ #[ doc( hidden) ]
870
+ heartbeat_option : HeartbeatOption ,
871
+ }
872
+
873
+ #[ derive( Clone , PartialEq , Debug ) ]
874
+ pub enum HeartbeatOption {
875
+ NoHeartbeat ,
876
+ FixedTime ( Duration ) ,
849
877
}
850
878
851
879
impl Default for TransactionOptions {
@@ -864,7 +892,7 @@ impl TransactionOptions {
864
892
read_only : false ,
865
893
retry_options : RetryOptions :: default_optimistic ( ) ,
866
894
check_level : CheckLevel :: Panic ,
867
- auto_heartbeat : true ,
895
+ heartbeat_option : HeartbeatOption :: FixedTime ( DEFAULT_HEARTBEAT_INTERVAL ) ,
868
896
}
869
897
}
870
898
@@ -877,7 +905,7 @@ impl TransactionOptions {
877
905
read_only : false ,
878
906
retry_options : RetryOptions :: default_pessimistic ( ) ,
879
907
check_level : CheckLevel :: Panic ,
880
- auto_heartbeat : true ,
908
+ heartbeat_option : HeartbeatOption :: FixedTime ( DEFAULT_HEARTBEAT_INTERVAL ) ,
881
909
}
882
910
}
883
911
@@ -935,10 +963,8 @@ impl TransactionOptions {
935
963
}
936
964
}
937
965
938
- /// The transaction should not automatically send heartbeat messages to TiKV to keep itself
939
- // alive.
940
- pub fn no_auto_hearbeat ( mut self ) -> TransactionOptions {
941
- self . auto_heartbeat = false ;
966
+ pub fn heartbeat_option ( mut self , heartbeat_option : HeartbeatOption ) -> TransactionOptions {
967
+ self . heartbeat_option = heartbeat_option;
942
968
self
943
969
}
944
970
@@ -967,10 +993,11 @@ pub enum CheckLevel {
967
993
None ,
968
994
}
969
995
970
- /// The default TTL of a lock in milliseconds.
971
- const DEFAULT_LOCK_TTL : u64 = 3000 ;
972
- /// The default heartbeat interval.
973
- const DEFAULT_HEARTBEAT_INTERVAL : Duration = Duration :: from_secs ( 1 ) ;
996
+ impl HeartbeatOption {
997
+ pub fn is_auto_heartbeat ( & self ) -> bool {
998
+ !matches ! ( self , HeartbeatOption :: NoHeartbeat )
999
+ }
1000
+ }
974
1001
975
1002
/// A struct wrapping the details of two-phase commit protocol (2PC).
976
1003
///
@@ -988,6 +1015,8 @@ struct Committer<PdC: PdClient = PdRpcClient> {
988
1015
options : TransactionOptions ,
989
1016
#[ new( default ) ]
990
1017
undetermined : bool ,
1018
+ write_size : u64 ,
1019
+ start_instant : Instant ,
991
1020
}
992
1021
993
1022
impl < PdC : PdClient > Committer < PdC > {
@@ -1026,20 +1055,20 @@ impl<PdC: PdClient> Committer<PdC> {
1026
1055
1027
1056
async fn prewrite ( & mut self ) -> Result < Option < Timestamp > > {
1028
1057
let primary_lock = self . primary_key . clone ( ) . unwrap ( ) ;
1029
- // FIXME: calculate TTL for big transactions
1030
- let lock_ttl = DEFAULT_LOCK_TTL ;
1058
+ let elapsed = self . start_instant . elapsed ( ) . as_millis ( ) as u64 ;
1059
+ let lock_ttl = self . calc_txn_lock_ttl ( ) ;
1031
1060
let mut request = match & self . options . kind {
1032
1061
TransactionKind :: Optimistic => new_prewrite_request (
1033
1062
self . mutations . clone ( ) ,
1034
1063
primary_lock,
1035
1064
self . start_version . clone ( ) ,
1036
- lock_ttl,
1065
+ lock_ttl + elapsed ,
1037
1066
) ,
1038
1067
TransactionKind :: Pessimistic ( for_update_ts) => new_pessimistic_prewrite_request (
1039
1068
self . mutations . clone ( ) ,
1040
1069
primary_lock,
1041
1070
self . start_version . clone ( ) ,
1042
- lock_ttl,
1071
+ lock_ttl + elapsed ,
1043
1072
for_update_ts. clone ( ) ,
1044
1073
) ,
1045
1074
} ;
@@ -1173,6 +1202,16 @@ impl<PdC: PdClient> Committer<PdC> {
1173
1202
}
1174
1203
Ok ( ( ) )
1175
1204
}
1205
+
1206
+ fn calc_txn_lock_ttl ( & mut self ) -> u64 {
1207
+ let mut lock_ttl = DEFAULT_LOCK_TTL ;
1208
+ if self . write_size > TXN_COMMIT_BATCH_SIZE {
1209
+ let size_mb = self . write_size as f64 / 1024.0 / 1024.0 ;
1210
+ lock_ttl = ( TTL_FACTOR * size_mb. sqrt ( ) ) as u64 ;
1211
+ lock_ttl = lock_ttl. min ( MAX_TTL ) . max ( DEFAULT_LOCK_TTL ) ;
1212
+ }
1213
+ lock_ttl
1214
+ }
1176
1215
}
1177
1216
1178
1217
#[ derive( PartialEq ) ]
@@ -1197,6 +1236,7 @@ enum TransactionStatus {
1197
1236
mod tests {
1198
1237
use crate :: {
1199
1238
mock:: { MockKvClient , MockPdClient } ,
1239
+ transaction:: HeartbeatOption ,
1200
1240
Transaction , TransactionOptions ,
1201
1241
} ;
1202
1242
use fail:: FailScenario ;
@@ -1207,40 +1247,42 @@ mod tests {
1207
1247
atomic:: { AtomicUsize , Ordering } ,
1208
1248
Arc ,
1209
1249
} ,
1250
+ time:: Duration ,
1210
1251
} ;
1211
1252
use tikv_client_proto:: { kvrpcpb, pdpb:: Timestamp } ;
1212
1253
1213
1254
#[ tokio:: test]
1214
1255
async fn test_optimistic_heartbeat ( ) -> Result < ( ) , io:: Error > {
1215
1256
let scenario = FailScenario :: setup ( ) ;
1216
- fail:: cfg ( "after-prewrite" , "sleep(10000 )" ) . unwrap ( ) ;
1257
+ fail:: cfg ( "after-prewrite" , "sleep(1500 )" ) . unwrap ( ) ;
1217
1258
let heartbeats = Arc :: new ( AtomicUsize :: new ( 0 ) ) ;
1218
1259
let heartbeats_cloned = heartbeats. clone ( ) ;
1219
1260
let pd_client = Arc :: new ( MockPdClient :: new ( MockKvClient :: with_dispatch_hook (
1220
1261
move |req : & dyn Any | {
1221
- if let Some ( _heartbeat ) = req. downcast_ref :: < kvrpcpb:: TxnHeartBeatRequest > ( ) {
1262
+ if let Some ( _ ) = req. downcast_ref :: < kvrpcpb:: TxnHeartBeatRequest > ( ) {
1222
1263
heartbeats_cloned. fetch_add ( 1 , Ordering :: SeqCst ) ;
1223
- return Ok ( Box :: new ( kvrpcpb:: TxnHeartBeatResponse :: default ( ) ) as Box < dyn Any > ) ;
1224
- } else if let Some ( _prewrite) = req. downcast_ref :: < kvrpcpb:: PrewriteRequest > ( ) {
1225
- return Ok ( Box :: new ( kvrpcpb:: PrewriteResponse :: default ( ) ) as Box < dyn Any > ) ;
1264
+ Ok ( Box :: new ( kvrpcpb:: TxnHeartBeatResponse :: default ( ) ) as Box < dyn Any > )
1265
+ } else if let Some ( _) = req. downcast_ref :: < kvrpcpb:: PrewriteRequest > ( ) {
1266
+ Ok ( Box :: new ( kvrpcpb:: PrewriteResponse :: default ( ) ) as Box < dyn Any > )
1267
+ } else {
1268
+ Ok ( Box :: new ( kvrpcpb:: CommitResponse :: default ( ) ) as Box < dyn Any > )
1226
1269
}
1227
- Ok ( Box :: new ( kvrpcpb:: CommitResponse :: default ( ) ) as Box < dyn Any > )
1228
1270
} ,
1229
1271
) ) ) ;
1230
1272
let key1 = "key1" . to_owned ( ) ;
1231
1273
let mut heartbeat_txn = Transaction :: new (
1232
1274
Timestamp :: default ( ) ,
1233
1275
pd_client,
1234
- TransactionOptions :: new_optimistic ( ) ,
1276
+ TransactionOptions :: new_optimistic ( )
1277
+ . heartbeat_option ( HeartbeatOption :: FixedTime ( Duration :: from_secs ( 1 ) ) ) ,
1235
1278
) ;
1236
1279
heartbeat_txn. put ( key1. clone ( ) , "foo" ) . await . unwrap ( ) ;
1237
1280
let heartbeat_txn_handle = tokio:: task:: spawn_blocking ( move || {
1238
1281
assert ! ( futures:: executor:: block_on( heartbeat_txn. commit( ) ) . is_ok( ) )
1239
1282
} ) ;
1240
1283
assert_eq ! ( heartbeats. load( Ordering :: SeqCst ) , 0 ) ;
1241
- tokio:: time:: sleep ( tokio:: time:: Duration :: from_secs ( 5 ) ) . await ;
1242
1284
heartbeat_txn_handle. await . unwrap ( ) ;
1243
- assert ! ( heartbeats. load( Ordering :: SeqCst ) >= 1 ) ;
1285
+ assert_eq ! ( heartbeats. load( Ordering :: SeqCst ) , 1 ) ;
1244
1286
scenario. teardown ( ) ;
1245
1287
Ok ( ( ) )
1246
1288
}
@@ -1251,35 +1293,33 @@ mod tests {
1251
1293
let heartbeats_cloned = heartbeats. clone ( ) ;
1252
1294
let pd_client = Arc :: new ( MockPdClient :: new ( MockKvClient :: with_dispatch_hook (
1253
1295
move |req : & dyn Any | {
1254
- if let Some ( _heartbeat ) = req. downcast_ref :: < kvrpcpb:: TxnHeartBeatRequest > ( ) {
1296
+ if let Some ( _ ) = req. downcast_ref :: < kvrpcpb:: TxnHeartBeatRequest > ( ) {
1255
1297
heartbeats_cloned. fetch_add ( 1 , Ordering :: SeqCst ) ;
1256
- return Ok ( Box :: new ( kvrpcpb:: TxnHeartBeatResponse :: default ( ) ) as Box < dyn Any > ) ;
1257
- } else if let Some ( _prewrite) = req. downcast_ref :: < kvrpcpb:: PrewriteRequest > ( ) {
1258
- return Ok ( Box :: new ( kvrpcpb:: PrewriteResponse :: default ( ) ) as Box < dyn Any > ) ;
1259
- } else if let Some ( _pessimistic_lock) =
1260
- req. downcast_ref :: < kvrpcpb:: PessimisticLockRequest > ( )
1261
- {
1262
- return Ok (
1263
- Box :: new ( kvrpcpb:: PessimisticLockResponse :: default ( ) ) as Box < dyn Any >
1264
- ) ;
1298
+ Ok ( Box :: new ( kvrpcpb:: TxnHeartBeatResponse :: default ( ) ) as Box < dyn Any > )
1299
+ } else if let Some ( _) = req. downcast_ref :: < kvrpcpb:: PrewriteRequest > ( ) {
1300
+ Ok ( Box :: new ( kvrpcpb:: PrewriteResponse :: default ( ) ) as Box < dyn Any > )
1301
+ } else if let Some ( _) = req. downcast_ref :: < kvrpcpb:: PessimisticLockRequest > ( ) {
1302
+ Ok ( Box :: new ( kvrpcpb:: PessimisticLockResponse :: default ( ) ) as Box < dyn Any > )
1303
+ } else {
1304
+ Ok ( Box :: new ( kvrpcpb:: CommitResponse :: default ( ) ) as Box < dyn Any > )
1265
1305
}
1266
- Ok ( Box :: new ( kvrpcpb:: CommitResponse :: default ( ) ) as Box < dyn Any > )
1267
1306
} ,
1268
1307
) ) ) ;
1269
1308
let key1 = "key1" . to_owned ( ) ;
1270
1309
let mut heartbeat_txn = Transaction :: new (
1271
1310
Timestamp :: default ( ) ,
1272
1311
pd_client,
1273
- TransactionOptions :: new_pessimistic ( ) ,
1312
+ TransactionOptions :: new_pessimistic ( )
1313
+ . heartbeat_option ( HeartbeatOption :: FixedTime ( Duration :: from_secs ( 1 ) ) ) ,
1274
1314
) ;
1275
1315
heartbeat_txn. put ( key1. clone ( ) , "foo" ) . await . unwrap ( ) ;
1276
1316
assert_eq ! ( heartbeats. load( Ordering :: SeqCst ) , 0 ) ;
1277
- tokio:: time:: sleep ( tokio:: time:: Duration :: from_millis ( 5000 ) ) . await ;
1317
+ tokio:: time:: sleep ( tokio:: time:: Duration :: from_millis ( 1500 ) ) . await ;
1318
+ assert_eq ! ( heartbeats. load( Ordering :: SeqCst ) , 1 ) ;
1278
1319
let heartbeat_txn_handle = tokio:: spawn ( async move {
1279
1320
assert ! ( heartbeat_txn. commit( ) . await . is_ok( ) ) ;
1280
1321
} ) ;
1281
1322
heartbeat_txn_handle. await . unwrap ( ) ;
1282
- assert ! ( heartbeats. load( Ordering :: SeqCst ) > 1 ) ;
1283
1323
Ok ( ( ) )
1284
1324
}
1285
1325
}
0 commit comments