github.com/KinWaiYuen/client-go/v2@v2.5.4/kv/kv.go (about)

     1  // Copyright 2021 TiKV Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // NOTE: The code in this file is based on code from the
    16  // TiDB project, licensed under the Apache License v 2.0
    17  //
    18  // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/kv/kv.go
    19  //
    20  
    21  package kv
    22  
    23  import (
    24  	"math"
    25  	"sync"
    26  	"time"
    27  
    28  	tikverr "github.com/KinWaiYuen/client-go/v2/error"
    29  	"github.com/KinWaiYuen/client-go/v2/util"
    30  )
    31  
    32  // ReturnedValue pairs the Value and AlreadyLocked flag for PessimisticLock return values result.
    33  type ReturnedValue struct {
    34  	Value         []byte
    35  	AlreadyLocked bool
    36  }
    37  
    38  // Used for pessimistic lock wait time
    39  // these two constants are special for lock protocol with tikv
    40  // math.MaxInt64 means always wait, 0 means nowait, others meaning lock wait in milliseconds
    41  const (
    42  	LockAlwaysWait = int64(math.MaxInt64)
    43  	LockNoWait     = int64(0)
    44  )
    45  
    46  type lockWaitTimeInMs struct {
    47  	value int64
    48  }
    49  
    50  func defaultLockWaitTime() *lockWaitTimeInMs {
    51  	return &lockWaitTimeInMs{value: LockAlwaysWait}
    52  }
    53  
    54  // LockCtx contains information for LockKeys method.
    55  type LockCtx struct {
    56  	Killed                *uint32
    57  	ForUpdateTS           uint64
    58  	lockWaitTime          *lockWaitTimeInMs
    59  	WaitStartTime         time.Time
    60  	PessimisticLockWaited *int32
    61  	LockKeysDuration      *int64
    62  	LockKeysCount         *int32
    63  	ReturnValues          bool
    64  	Values                map[string]ReturnedValue
    65  	ValuesLock            sync.Mutex
    66  	LockExpired           *uint32
    67  	Stats                 *util.LockKeysDetails
    68  	ResourceGroupTag      []byte
    69  	OnDeadlock            func(*tikverr.ErrDeadlock)
    70  }
    71  
    72  // LockWaitTime returns lockWaitTimeInMs
    73  func (ctx *LockCtx) LockWaitTime() int64 {
    74  	if ctx.lockWaitTime == nil {
    75  		ctx.lockWaitTime = defaultLockWaitTime()
    76  	}
    77  	return ctx.lockWaitTime.value
    78  }
    79  
    80  // NewLockCtx creates a LockCtx.
    81  func NewLockCtx(forUpdateTS uint64, lockWaitTime int64, waitStartTime time.Time) *LockCtx {
    82  	return &LockCtx{
    83  		ForUpdateTS:   forUpdateTS,
    84  		lockWaitTime:  &lockWaitTimeInMs{value: lockWaitTime},
    85  		WaitStartTime: waitStartTime,
    86  	}
    87  }
    88  
    89  // InitReturnValues creates the map to store returned value.
    90  func (ctx *LockCtx) InitReturnValues(valueLen int) {
    91  	ctx.ReturnValues = true
    92  	ctx.Values = make(map[string]ReturnedValue, valueLen)
    93  }
    94  
    95  // GetValueNotLocked returns a value if the key is not already locked.
    96  // (nil, false) means already locked.
    97  func (ctx *LockCtx) GetValueNotLocked(key []byte) ([]byte, bool) {
    98  	rv := ctx.Values[string(key)]
    99  	if !rv.AlreadyLocked {
   100  		return rv.Value, true
   101  	}
   102  	return nil, false
   103  }
   104  
   105  // IterateValuesNotLocked applies f to all key-values that are not already
   106  // locked.
   107  func (ctx *LockCtx) IterateValuesNotLocked(f func([]byte, []byte)) {
   108  	ctx.ValuesLock.Lock()
   109  	defer ctx.ValuesLock.Unlock()
   110  	for key, val := range ctx.Values {
   111  		if !val.AlreadyLocked {
   112  			f([]byte(key), val.Value)
   113  		}
   114  	}
   115  }