6.深入TiDB：乐观事务

func (s *session) NewTxn(ctx context.Context) error {
	if err := s.checkBeforeNewTxn(ctx); err != nil {
		return err
	}
	// 开启事务
	txn, err := s.store.BeginWithOption(tikv.DefaultStartTSOption().SetTxnScope(s.sessionVars.CheckAndGetTxnScope()))
	if err != nil {
		return err
	}
	...
	return nil
}

func (s *KVStore) Begin() (*KVTxn, error) {
	return s.BeginWithOption(DefaultStartTSOption())
}
 
func (s *KVStore) BeginWithOption(options StartTSOption) (*KVTxn, error) {
	return newTiKVTxnWithOptions(s, options)
}

func newTiKVTxnWithOptions(store *KVStore, options StartTSOption) (*KVTxn, error) {
	if options.TxnScope == "" {
		options.TxnScope = oracle.GlobalTxnScope
	}
	// 去PD服务获取一个时间戳
	startTS, err := ExtractStartTS(store, options)
	if err != nil {
		return nil, errors.Trace(err)
	}
	snapshot := newTiKVSnapshot(store, startTS, store.nextReplicaReadSeed())
	newTiKVTxn := &KVTxn{
		snapshot:  snapshot,
		us:        unionstore.NewUnionStore(snapshot),
		store:     store,
		startTS:   startTS,
		startTime: time.Now(),
		valid:     true,
		vars:      tikv.DefaultVars,
		scope:     options.TxnScope,
	}
	return newTiKVTxn, nil
}

func (t *TableCommon) AddRecord(sctx sessionctx.Context, r []types.Datum, opts ...table.AddRecordOption) (recordID kv.Handle, err error) {
    ...
	var setPresume bool
	skipCheck := sctx.GetSessionVars().StmtCtx.BatchCheck
	if (t.meta.IsCommonHandle || t.meta.PKIsHandle) && !skipCheck && !opt.SkipHandleCheck {
		// 如果是 LazyCheck ，那么只读取本地缓存判断是否存在
		if sctx.GetSessionVars().LazyCheckKeyNotExists() {
			var v []byte
			//只读取本地缓存判断是否存在
			v, err = txn.GetMemBuffer().Get(ctx, key)
			if err != nil {
				setPresume = true
			}
			if err == nil && len(v) == 0 {
				err = kv.ErrNotExist
			}
		} else {
			//否则会通过rpc请求tikv从集群中校验数据是否存在
			_, err = txn.Get(ctx, key)
		}
		if err == nil {
			handleStr := getDuplicateErrorHandleString(t, recordID, r)
			return recordID, kv.ErrKeyExists.FastGenByArgs(handleStr, "PRIMARY")
		} else if !kv.ErrNotExist.Equal(err) {
			return recordID, err
		}
	} 
	if setPresume {
		// 表示假定数据不存在
		err = memBuffer.SetWithFlags(key, value, kv.SetPresumeKeyNotExists)
	} else {
        //将 Key-Value 写到当前事务的缓存中
		err = memBuffer.Set(key, value)
	}
	if err != nil {
		return nil, err
	}
    ...
}

func newTwoPhaseCommitter(txn *KVTxn, sessionID uint64) (*twoPhaseCommitter, error) {
   return &twoPhaseCommitter{
      store:         txn.store,
      txn:           txn,
      startTS:       txn.StartTS(),
      sessionID:     sessionID,
      regionTxnSize: map[uint64]int{},
      ttlManager: ttlManager{
         ch: make(chan struct{}),
      },
      isPessimistic: txn.IsPessimistic(),
      binlog:        txn.binlog,
   }, nil
}

func (c *twoPhaseCommitter) initKeysAndMutations() error {
	var size, putCnt, delCnt, lockCnt, checkCnt int

	txn := c.txn
	// 当前事务的数据都存放在 memBuf 中
	// memBuffer里的 key 是有序排列
	memBuf := txn.GetMemBuffer()
	sizeHint := txn.us.GetMemBuffer().Len()
	c.mutations = newMemBufferMutations(sizeHint, memBuf)
	c.isPessimistic = txn.IsPessimistic()
	filter := txn.kvFilter

	var err error
	// 遍历 memBuffer 可以顺序的收集到事务里需要修改的 key
	for it := memBuf.IterWithFlags(nil, nil); it.Valid(); err = it.Next() {
		_ = err
		key := it.Key()
		flags := it.Flags()
		var value []byte
		var op pb.Op

		if !it.HasValue() {
			...
		} else {
			value = it.Value()
			...
		}

		var isPessimistic bool
		if flags.HasLocked() {
			isPessimistic = c.isPessimistic
		}
		c.mutations.Push(op, isPessimistic, it.Handle())
		size += len(key) + len(value)

		if len(c.primaryKey) == 0 && op != pb.Op_CheckNotExists {
			c.primaryKey = key
		}
	}
	...

	commitDetail := &util.CommitDetails{WriteSize: size, WriteKeys: c.mutations.Len()}
	metrics.TiKVTxnWriteKVCountHistogram.Observe(float64(commitDetail.WriteKeys))
	metrics.TiKVTxnWriteSizeHistogram.Observe(float64(commitDetail.WriteSize))
	c.hasNoNeedCommitKeys = checkCnt > 0
	// 计算事务的 TTL 时间
	c.lockTTL = txnLockTTL(txn.startTime, size)
	c.priority = txn.priority.ToPB()
	c.syncLog = txn.syncLog
	c.resourceGroupTag = txn.resourceGroupTag
	c.setDetail(commitDetail)
	return nil
}

func txnLockTTL(startTime time.Time, txnSize int) uint64 { 
	lockTTL := defaultLockTTL
	// 当事务大小大于16KB
	if txnSize >= txnCommitBatchSize {
		sizeMiB := float64(txnSize) / bytesPerMiB
		// 6000 * 事务大小平方根
		lockTTL = uint64(float64(ttlFactor) * math.Sqrt(sizeMiB))
		//最小为3s
		if lockTTL < defaultLockTTL {
			lockTTL = defaultLockTTL
		}
		//最大为20s
		if lockTTL > ManagedLockTTL {
			lockTTL = ManagedLockTTL
		}
	}
 
	elapsed := time.Since(startTime) / time.Millisecond
	return lockTTL + uint64(elapsed)
}

func (c *twoPhaseCommitter) prewriteMutations(bo *Backoffer, mutations CommitterMutations) error {
	... 
	return c.doActionOnMutations(bo, actionPrewrite{}, mutations)
}

func (c *twoPhaseCommitter) doActionOnMutations(bo *Backoffer, action twoPhaseCommitAction, mutations CommitterMutations) error {
	if mutations.Len() == 0 {
		return nil
	}
	//按照region分组对mutations进行分组
	groups, err := c.groupMutations(bo, mutations)
	if err != nil {
		return errors.Trace(err)
	}
	...
	//进一步的分批处理
	return c.doActionOnGroupMutations(bo, action, groups)
}

func (c *twoPhaseCommitter) doActionOnGroupMutations(bo *Backoffer, action twoPhaseCommitAction, groups []groupedMutations) error {
	action.tiKVTxnRegionsNumHistogram().Observe(float64(len(groups))) 
	... 
	batchBuilder := newBatched(c.primary())
	//每个分组内按16KB大小再分批
	for _, group := range groups {
		batchBuilder.appendBatchMutationsBySize(group.region, group.mutations, sizeFunc, txnCommitBatchSize)
	}
	firstIsPrimary := batchBuilder.setPrimary()

	actionCommit, actionIsCommit := action.(actionCommit)
	...
	//commit先同步的提交primary key所在的batch
	if firstIsPrimary &&
		((actionIsCommit && !c.isAsyncCommit()) || actionIsCleanup || actionIsPessimiticLock) {
		// primary should be committed(not async commit)/cleanup/pessimistically locked first
		err = c.doActionOnBatches(bo, action, batchBuilder.primaryBatch())
		... 
		//提交完之后将primary key所在的batch移除
		batchBuilder.forgetPrimary()
	}
	// Already spawned a goroutine for async commit transaction.
	// 其它的key由go routine后台异步的提交
	if actionIsCommit && !actionCommit.retry && !c.isAsyncCommit() {
		secondaryBo := retry.NewBackofferWithVars(context.Background(), CommitSecondaryMaxBackoff, c.txn.vars)
		go func() {
			... 
			e := c.doActionOnBatches(secondaryBo, action, batchBuilder.allBatches())
			...
		}()
	} else {
		//执行 prewrite
		err = c.doActionOnBatches(bo, action, batchBuilder.allBatches())
	}
	return errors.Trace(err)
}

func (action actionPrewrite) handleSingleBatch(c *twoPhaseCommitter, bo *Backoffer, batch batchMutations) (err error) {
	...
	// 获取事务的大小
	txnSize := uint64(c.regionTxnSize[batch.region.id])
	// 因为region的缺失导致的重试，所以不知道事务大小，这里重置事务大小为最大值
	if action.retry {
		txnSize = math.MaxUint64
	} 
	tBegin := time.Now()
	attempts := 0
	// 构建 Request
	req := c.buildPrewriteRequest(batch, txnSize)
	// 构建 RegionRequestSender
	sender := NewRegionRequestSender(c.store.regionCache, c.store.GetTiKVClient())
	 
	for {
		//尝试次数
		attempts++
		// 如果请求超过了1分钟，那么打印一条日志
		if time.Since(tBegin) > slowRequestThreshold {
			logutil.BgLogger().Warn("slow prewrite request", zap.Uint64("startTS", c.startTS), zap.Stringer("region", &batch.region), zap.Int("attempts", attempts))
			tBegin = time.Now()
		}
		//发送请求
		resp, err := sender.SendReq(bo, req, batch.region, client.ReadTimeoutShort)
		// Unexpected error occurs, return it
		if err != nil {
			return errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return errors.Trace(err)
		}
		// 如果遇到了regionError， 则需要重新调用doActionOnMutations重新分组，重新尝试
		if regionErr != nil {
			...
			err = c.doActionOnMutations(bo, actionPrewrite{true}, batch.mutations)
			return errors.Trace(err)
		}

		if resp.Resp == nil {
			return errors.Trace(tikverr.ErrBodyMissing)
		}
		prewriteResp := resp.Resp.(*pb.PrewriteResponse)
		keyErrs := prewriteResp.GetErrors()
		if len(keyErrs) == 0 {
			//如果没有keyError，并且Batch是primary，则启动一个tllManager
			if batch.isPrimary { 
				// 如果事务大于32M，那么开启ttlManager定时发送TxnHeartBeat心跳
				if int64(c.txnSize) > config.GetGlobalConfig().TiKVClient.TTLRefreshedTxnSize &&
					prewriteResp.OnePcCommitTs == 0 {
					c.run(c, nil)
				}
			} 
			...
			return nil
		}
		var locks []*Lock
		for _, keyErr := range keyErrs { 
			// 该key已存在
			if alreadyExist := keyErr.GetAlreadyExist(); alreadyExist != nil {
				e := &tikverr.ErrKeyExist{AlreadyExist: alreadyExist}
				err = c.extractKeyExistsErr(e)
				if err != nil {
					atomic.StoreUint32(&c.prewriteFailed, 1)
				}
				return err
			} 
			// 从 keyErr 中抽取出冲突的lock
			lock, err1 := extractLockFromKeyErr(keyErr)
			if err1 != nil {
				atomic.StoreUint32(&c.prewriteFailed, 1)
				return errors.Trace(err1)
			} 
			locks = append(locks, lock)
		}
		start := time.Now()
		//尝试解决这些locks，获取锁的过期时间
		msBeforeExpired, err := c.store.lockResolver.resolveLocksForWrite(bo, c.startTS, locks)
		if err != nil {
			return errors.Trace(err)
		}
		atomic.AddInt64(&c.getDetail().ResolveLockTime, int64(time.Since(start)))
		if msBeforeExpired > 0 {
			// 过期时间大于0，那么sleep等待
			err = bo.BackoffWithCfgAndMaxSleep(retry.BoTxnLock, int(msBeforeExpired), errors.Errorf("2PC prewrite lockedKeys: %d", len(locks)))
			if err != nil {
				return errors.Trace(err)
			}
		}
	}
}

func (actionCommit) handleSingleBatch(c *twoPhaseCommitter, bo *Backoffer, batch batchMutations) error {
	req := tikvrpc.NewRequest(...) 
	sender := NewRegionRequestSender(c.store.regionCache, c.store.GetTiKVClient())
	for {
		// 重试次数
		attempts++
		...
		//向 tikv 发起提交请求
		resp, err := sender.SendReq(bo, req, batch.region, client.ReadTimeoutShort)
		...
		// 如果遇到了regionError， 则需要重新调用doActionOnMutations重新分组，重新尝试
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return errors.Trace(err)
		}
		if regionErr != nil {
			...
			// 重新分组
			err = c.doActionOnMutations(bo, actionCommit{true}, batch.mutations)
			return errors.Trace(err)
		} 
		commitResp := resp.Resp.(*pb.CommitResponse)
		 
		if keyErr := commitResp.GetError(); keyErr != nil {
			if rejected := keyErr.GetCommitTsExpired(); rejected != nil {
				// 重新获取 commitTS
				commitTS, err := c.store.getTimestampWithRetry(bo, c.txn.GetScope())
				...
				continue
			}
			...
		}
		break
	} 
	return nil
}