github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/stochastik/txn.go

// Copyright 2020 WHTCORPS INC, Inc.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package stochastik

import (
	"bytes"
	"context"
	"fmt"
	"runtime/trace"
	"strings"
	"sync/atomic"

	"github.com/opentracing/opentracing-go"
	"github.com/whtcorpsinc/BerolinaSQL/terror"
	"github.com/whtcorpsinc/errors"
	"github.com/whtcorpsinc/failpoint"
	"github.com/whtcorpsinc/fidelpb/go-binlog"
	"github.com/whtcorpsinc/milevadb/blockcodec"
	"github.com/whtcorpsinc/milevadb/causetstore/einsteindb/oracle"
	"github.com/whtcorpsinc/milevadb/config"
	"github.com/whtcorpsinc/milevadb/ekv"
	"github.com/whtcorpsinc/milevadb/soliton/logutil"
	"github.com/whtcorpsinc/milevadb/stochastikctx"
	"github.com/whtcorpsinc/milevadb/stochastikctx/binloginfo"
	"go.uber.org/zap"
)

// TxnState wraps ekv.Transaction to provide a new ekv.Transaction.
// 1. It holds all memex related modification in the buffer before flush to the txn,
// so if execute memex meets error, the txn won't be made dirty.
// 2. It's a lazy transaction, that means it's a txnFuture before StartTS() is really need.
type TxnState struct {
	// States of a TxnState should be one of the followings:
	// Invalid: ekv.Transaction == nil && txnFuture == nil
	// Pending: ekv.Transaction == nil && txnFuture != nil
	// Valid:	ekv.Transaction != nil && txnFuture == nil
	ekv.Transaction
	txnFuture *txnFuture

	initCnt       int
	stagingHandle ekv.StagingHandle
	mutations     map[int64]*binlog.TableMutation
}

func (st *TxnState) init() {
	st.mutations = make(map[int64]*binlog.TableMutation)
}

func (st *TxnState) initStmtBuf() {
	if st.Transaction == nil {
		return
	}
	buf := st.Transaction.GetMemBuffer()
	st.initCnt = buf.Len()
	st.stagingHandle = buf.Staging()
}

// countHint is estimated count of mutations.
func (st *TxnState) countHint() int {
	if st.stagingHandle == ekv.InvalidStagingHandle {
		return 0
	}
	return st.Transaction.GetMemBuffer().Len() - st.initCnt
}

func (st *TxnState) flushStmtBuf() {
	if st.stagingHandle == ekv.InvalidStagingHandle {
		return
	}
	buf := st.Transaction.GetMemBuffer()
	buf.Release(st.stagingHandle)
	st.initCnt = buf.Len()
}

func (st *TxnState) cleanupStmtBuf() {
	if st.stagingHandle == ekv.InvalidStagingHandle {
		return
	}
	buf := st.Transaction.GetMemBuffer()
	buf.Cleanup(st.stagingHandle)
	st.initCnt = buf.Len()
}

// Size implements the MemBuffer interface.
func (st *TxnState) Size() int {
	if st.Transaction == nil {
		return 0
	}
	return st.Transaction.Size()
}

// Valid implements the ekv.Transaction interface.
func (st *TxnState) Valid() bool {
	return st.Transaction != nil && st.Transaction.Valid()
}

func (st *TxnState) pending() bool {
	return st.Transaction == nil && st.txnFuture != nil
}

func (st *TxnState) validOrPending() bool {
	return st.txnFuture != nil || st.Valid()
}

func (st *TxnState) String() string {
	if st.Transaction != nil {
		return st.Transaction.String()
	}
	if st.txnFuture != nil {
		return "txnFuture"
	}
	return "invalid transaction"
}

// GoString implements the "%#v" format for fmt.Printf.
func (st *TxnState) GoString() string {
	var s strings.Builder
	s.WriteString("Txn{")
	if st.pending() {
		s.WriteString("state=pending")
	} else if st.Valid() {
		s.WriteString("state=valid")
		fmt.Fprintf(&s, ", txnStartTS=%d", st.Transaction.StartTS())
		if len(st.mutations) > 0 {
			fmt.Fprintf(&s, ", len(mutations)=%d, %#v", len(st.mutations), st.mutations)
		}
	} else {
		s.WriteString("state=invalid")
	}

	s.WriteString("}")
	return s.String()
}

func (st *TxnState) changeInvalidToValid(txn ekv.Transaction) {
	st.Transaction = txn
	st.initStmtBuf()
	st.txnFuture = nil
}

func (st *TxnState) changeInvalidToPending(future *txnFuture) {
	st.Transaction = nil
	st.txnFuture = future
}

func (st *TxnState) changePendingToValid(ctx context.Context) error {
	if st.txnFuture == nil {
		return errors.New("transaction future is not set")
	}

	future := st.txnFuture
	st.txnFuture = nil

	defer trace.StartRegion(ctx, "WaitTsoFuture").End()
	txn, err := future.wait()
	if err != nil {
		st.Transaction = nil
		return err
	}
	st.Transaction = txn
	st.initStmtBuf()
	return nil
}

func (st *TxnState) changeToInvalid() {
	if st.stagingHandle != ekv.InvalidStagingHandle {
		st.Transaction.GetMemBuffer().Cleanup(st.stagingHandle)
	}
	st.stagingHandle = ekv.InvalidStagingHandle
	st.Transaction = nil
	st.txnFuture = nil
}

var hasMockAutoIncIDRetry = int64(0)

func enableMockAutoIncIDRetry() {
	atomic.StoreInt64(&hasMockAutoIncIDRetry, 1)
}

func mockAutoIncIDRetry() bool {
	return atomic.LoadInt64(&hasMockAutoIncIDRetry) == 1
}

var mockAutoRandIDRetryCount = int64(0)

func needMockAutoRandIDRetry() bool {
	return atomic.LoadInt64(&mockAutoRandIDRetryCount) > 0
}

func decreaseMockAutoRandIDRetryCount() {
	atomic.AddInt64(&mockAutoRandIDRetryCount, -1)
}

// ResetMockAutoRandIDRetryCount set the number of occurrences of
// `ekv.ErrTxnRetryable` when calling TxnState.Commit().
func ResetMockAutoRandIDRetryCount(failTimes int64) {
	atomic.StoreInt64(&mockAutoRandIDRetryCount, failTimes)
}

// Commit overrides the Transaction interface.
func (st *TxnState) Commit(ctx context.Context) error {
	defer st.reset()
	if len(st.mutations) != 0 || st.countHint() != 0 {
		logutil.BgLogger().Error("the code should never run here",
			zap.String("TxnState", st.GoString()),
			zap.Int("staging handler", int(st.stagingHandle)),
			zap.Stack("something must be wrong"))
		return errors.Trace(ekv.ErrInvalidTxn)
	}

	// mockCommitError8942 is used for PR #8942.
	failpoint.Inject("mockCommitError8942", func(val failpoint.Value) {
		if val.(bool) {
			failpoint.Return(ekv.ErrTxnRetryable)
		}
	})

	// mockCommitRetryForAutoIncID is used to mock an commit retry for adjustAutoIncrementCauset.
	failpoint.Inject("mockCommitRetryForAutoIncID", func(val failpoint.Value) {
		if val.(bool) && !mockAutoIncIDRetry() {
			enableMockAutoIncIDRetry()
			failpoint.Return(ekv.ErrTxnRetryable)
		}
	})

	failpoint.Inject("mockCommitRetryForAutoRandID", func(val failpoint.Value) {
		if val.(bool) && needMockAutoRandIDRetry() {
			decreaseMockAutoRandIDRetryCount()
			failpoint.Return(ekv.ErrTxnRetryable)
		}
	})

	return st.Transaction.Commit(ctx)
}

// Rollback overrides the Transaction interface.
func (st *TxnState) Rollback() error {
	defer st.reset()
	return st.Transaction.Rollback()
}

func (st *TxnState) reset() {
	st.cleanup()
	st.changeToInvalid()
}

func (st *TxnState) cleanup() {
	st.cleanupStmtBuf()
	st.initStmtBuf()
	for key := range st.mutations {
		delete(st.mutations, key)
	}
}

// KeysNeedToLock returns the keys need to be locked.
func (st *TxnState) KeysNeedToLock() ([]ekv.Key, error) {
	if st.stagingHandle == ekv.InvalidStagingHandle {
		return nil, nil
	}
	keys := make([]ekv.Key, 0, st.countHint())
	buf := st.Transaction.GetMemBuffer()
	buf.InspectStage(st.stagingHandle, func(k ekv.Key, flags ekv.KeyFlags, v []byte) {
		if !keyNeedToLock(k, v, flags) {
			return
		}
		keys = append(keys, k)
	})
	return keys, nil
}

func keyNeedToLock(k, v []byte, flags ekv.KeyFlags) bool {
	isTableKey := bytes.HasPrefix(k, blockcodec.TablePrefix())
	if !isTableKey {
		// spacetime key always need to dagger.
		return true
	}
	if flags.HasPresumeKeyNotExists() {
		return true
	}
	isDelete := len(v) == 0
	if isDelete {
		// only need to delete event key.
		return k[10] == 'r'
	}
	if blockcodec.IsUntouchedIndexKValue(k, v) {
		return false
	}
	isNonUniqueIndex := blockcodec.IsIndexKey(k) && len(v) == 1
	// Put event key and unique index need to dagger.
	return !isNonUniqueIndex
}

func getBinlogMutation(ctx stochastikctx.Context, blockID int64) *binlog.TableMutation {
	bin := binloginfo.GetPrewriteValue(ctx, true)
	for i := range bin.Mutations {
		if bin.Mutations[i].TableId == blockID {
			return &bin.Mutations[i]
		}
	}
	idx := len(bin.Mutations)
	bin.Mutations = append(bin.Mutations, binlog.TableMutation{TableId: blockID})
	return &bin.Mutations[idx]
}

func mergeToMutation(m1, m2 *binlog.TableMutation) {
	m1.InsertedRows = append(m1.InsertedRows, m2.InsertedRows...)
	m1.UFIDelatedRows = append(m1.UFIDelatedRows, m2.UFIDelatedRows...)
	m1.DeletedIds = append(m1.DeletedIds, m2.DeletedIds...)
	m1.DeletedPks = append(m1.DeletedPks, m2.DeletedPks...)
	m1.DeletedRows = append(m1.DeletedRows, m2.DeletedRows...)
	m1.Sequence = append(m1.Sequence, m2.Sequence...)
}

type txnFailFuture struct{}

func (txnFailFuture) Wait() (uint64, error) {
	return 0, errors.New("mock get timestamp fail")
}

// txnFuture is a promise, which promises to return a txn in future.
type txnFuture struct {
	future      oracle.Future
	causetstore ekv.CausetStorage
}

func (tf *txnFuture) wait() (ekv.Transaction, error) {
	startTS, err := tf.future.Wait()
	if err == nil {
		return tf.causetstore.BeginWithStartTS(startTS)
	} else if config.GetGlobalConfig().CausetStore == "entangledstore" {
		return nil, err
	}

	logutil.BgLogger().Warn("wait tso failed", zap.Error(err))
	// It would retry get timestamp.
	return tf.causetstore.Begin()
}

func (s *stochastik) getTxnFuture(ctx context.Context) *txnFuture {
	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
		span1 := span.Tracer().StartSpan("stochastik.getTxnFuture", opentracing.ChildOf(span.Context()))
		defer span1.Finish()
		ctx = opentracing.ContextWithSpan(ctx, span1)
	}

	oracleStore := s.causetstore.GetOracle()
	var tsFuture oracle.Future
	if s.stochastikVars.LowResolutionTSO {
		tsFuture = oracleStore.GetLowResolutionTimestampAsync(ctx)
	} else {
		tsFuture = oracleStore.GetTimestampAsync(ctx)
	}
	ret := &txnFuture{future: tsFuture, causetstore: s.causetstore}
	failpoint.InjectContext(ctx, "mockGetTSFail", func() {
		ret.future = txnFailFuture{}
	})
	return ret
}

// HasDirtyContent checks whether there's dirty uFIDelate on the given causet.
// Put this function here is to avoid cycle import.
func (s *stochastik) HasDirtyContent(tid int64) bool {
	if s.txn.Transaction == nil {
		return false
	}
	seekKey := blockcodec.EncodeTablePrefix(tid)
	it, err := s.txn.GetMemBuffer().Iter(seekKey, nil)
	terror.Log(err)
	return it.Valid() && bytes.HasPrefix(it.Key(), seekKey)
}

// StmtCommit implements the stochastikctx.Context interface.
func (s *stochastik) StmtCommit() {
	defer func() {
		s.txn.cleanup()
	}()

	st := &s.txn
	st.flushStmtBuf()

	// Need to flush binlog.
	for blockID, delta := range st.mutations {
		mutation := getBinlogMutation(s, blockID)
		mergeToMutation(mutation, delta)
	}
}

// StmtRollback implements the stochastikctx.Context interface.
func (s *stochastik) StmtRollback() {
	s.txn.cleanup()
}

// StmtGetMutation implements the stochastikctx.Context interface.
func (s *stochastik) StmtGetMutation(blockID int64) *binlog.TableMutation {
	st := &s.txn
	if _, ok := st.mutations[blockID]; !ok {
		st.mutations[blockID] = &binlog.TableMutation{TableId: blockID}
	}
	return st.mutations[blockID]
}