github.com/KinWaiYuen/client-go/v2@v2.5.4/internal/latch/latch_test.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/latch/latch_test.go
    19  //
    20  
    21  // Copyright 2018 PingCAP, Inc.
    22  //
    23  // Licensed under the Apache License, Version 2.0 (the "License");
    24  // you may not use this file except in compliance with the License.
    25  // You may obtain a copy of the License at
    26  //
    27  //     http://www.apache.org/licenses/LICENSE-2.0
    28  //
    29  // Unless required by applicable law or agreed to in writing, software
    30  // distributed under the License is distributed on an "AS IS" BASIS,
    31  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    32  // See the License for the specific language governing permissions and
    33  // limitations under the License.
    34  
    35  package latch
    36  
    37  import (
    38  	"testing"
    39  	"time"
    40  
    41  	"github.com/KinWaiYuen/client-go/v2/oracle"
    42  	"github.com/stretchr/testify/assert"
    43  	"go.uber.org/atomic"
    44  )
    45  
    46  var tso atomic.Uint64
    47  
    48  func getTso() uint64 {
    49  	return tso.Inc()
    50  }
    51  
    52  func newLock(latches *Latches, keys [][]byte) (startTs uint64, lock *Lock) {
    53  	startTs = getTso()
    54  	lock = latches.genLock(startTs, keys)
    55  	return
    56  }
    57  
    58  func TestWakeUp(t *testing.T) {
    59  	latches := NewLatches(256)
    60  
    61  	keysA := [][]byte{[]byte("a"), []byte("b"), []byte("c")}
    62  	_, lockA := newLock(latches, keysA)
    63  
    64  	keysB := [][]byte{[]byte("d"), []byte("e"), []byte("a"), []byte("c")}
    65  	startTSB, lockB := newLock(latches, keysB)
    66  
    67  	// A acquire lock success.
    68  	result := latches.acquire(lockA)
    69  	assert.Equal(t, acquireSuccess, result)
    70  
    71  	// B acquire lock failed.
    72  	result = latches.acquire(lockB)
    73  	assert.Equal(t, acquireLocked, result)
    74  
    75  	// A release lock, and get wakeup list.
    76  	commitTSA := getTso()
    77  	wakeupList := make([]*Lock, 0)
    78  	lockA.SetCommitTS(commitTSA)
    79  	wakeupList = latches.release(lockA, wakeupList)
    80  	assert.Equal(t, startTSB, wakeupList[0].startTS)
    81  
    82  	// B acquire failed since startTSB has stale for some keys.
    83  	result = latches.acquire(lockB)
    84  	assert.Equal(t, acquireStale, result)
    85  
    86  	// B release lock since it received a stale.
    87  	wakeupList = latches.release(lockB, wakeupList)
    88  	assert.Len(t, wakeupList, 0)
    89  
    90  	// B restart:get a new startTS.
    91  	startTSB = getTso()
    92  	lockB = latches.genLock(startTSB, keysB)
    93  	result = latches.acquire(lockB)
    94  	assert.Equal(t, acquireSuccess, result)
    95  }
    96  
    97  func TestFirstAcquireFailedWithStale(t *testing.T) {
    98  	latches := NewLatches(256)
    99  
   100  	keys := [][]byte{[]byte("a"), []byte("b"), []byte("c")}
   101  	_, lockA := newLock(latches, keys)
   102  	startTSB, lockB := newLock(latches, keys)
   103  
   104  	// acquire lockA success
   105  	result := latches.acquire(lockA)
   106  	assert.Equal(t, acquireSuccess, result)
   107  
   108  	// release lockA
   109  	commitTSA := getTso()
   110  	wakeupList := make([]*Lock, 0)
   111  	lockA.SetCommitTS(commitTSA)
   112  	latches.release(lockA, wakeupList)
   113  	assert.Greater(t, commitTSA, startTSB)
   114  
   115  	// acquire lockB first time, should be failed with stale since commitTSA > startTSB
   116  	result = latches.acquire(lockB)
   117  	assert.Equal(t, acquireStale, result)
   118  	latches.release(lockB, wakeupList)
   119  }
   120  
   121  func TestRecycle(t *testing.T) {
   122  	latches := NewLatches(8)
   123  	now := time.Now()
   124  	startTS := oracle.GoTimeToTS(now)
   125  	lock := latches.genLock(startTS, [][]byte{
   126  		[]byte("a"), []byte("b"),
   127  	})
   128  	lock1 := latches.genLock(startTS, [][]byte{
   129  		[]byte("b"), []byte("c"),
   130  	})
   131  	assert.Equal(t, acquireSuccess, latches.acquire(lock))
   132  	assert.Equal(t, acquireLocked, latches.acquire(lock1))
   133  	lock.SetCommitTS(startTS + 1)
   134  	var wakeupList []*Lock
   135  	latches.release(lock, wakeupList)
   136  	// Release lock will grant latch to lock1 automatically,
   137  	// so release lock1 is called here.
   138  	latches.release(lock1, wakeupList)
   139  
   140  	lock2 := latches.genLock(startTS+3, [][]byte{
   141  		[]byte("b"), []byte("c"),
   142  	})
   143  	assert.Equal(t, acquireSuccess, latches.acquire(lock2))
   144  	wakeupList = wakeupList[:0]
   145  	latches.release(lock2, wakeupList)
   146  
   147  	allEmpty := true
   148  	for i := 0; i < len(latches.slots); i++ {
   149  		latch := &latches.slots[i]
   150  		if latch.queue != nil {
   151  			allEmpty = false
   152  		}
   153  	}
   154  	assert.False(t, allEmpty)
   155  
   156  	currentTS := oracle.GoTimeToTS(now.Add(expireDuration)) + 3
   157  	latches.recycle(currentTS)
   158  
   159  	for i := 0; i < len(latches.slots); i++ {
   160  		latch := &latches.slots[i]
   161  		assert.Nil(t, latch.queue)
   162  	}
   163  }