github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/storage/latch/latch_test.go (about)

     1  // Copyright 2022 zGraph Authors. All rights reserved.
     2  //
     3  // Copyright 2018 PingCAP, Inc.
     4  //
     5  // Licensed under the Apache License, Version 2.0 (the "License");
     6  // you may not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //     http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package latch
    17  
    18  import (
    19  	"sync/atomic"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/stretchr/testify/assert"
    24  	"github.com/vescale/zgraph/storage/kv"
    25  )
    26  
    27  var baseTso = uint64(time.Now().UnixNano())
    28  
    29  func getTso() kv.Version {
    30  	return kv.Version(atomic.AddUint64(&baseTso, uint64(1)))
    31  }
    32  
    33  func TestWakeUp(t *testing.T) {
    34  	assert := assert.New(t)
    35  	latches := NewLatches(256)
    36  	keysA := []kv.Key{
    37  		[]byte("a"), []byte("b"), []byte("c")}
    38  	startVerA := getTso()
    39  	lockA := latches.genLock(startVerA, keysA)
    40  
    41  	keysB := []kv.Key{[]byte("d"), []byte("e"), []byte("a"), []byte("c")}
    42  	startVerB := getTso()
    43  	lockB := latches.genLock(startVerB, keysB)
    44  
    45  	// A acquire lock success.
    46  	result := latches.acquire(lockA)
    47  	assert.Equal(result, acquireSuccess)
    48  
    49  	// B acquire lock failed.
    50  	result = latches.acquire(lockB)
    51  	assert.Equal(result, acquireLocked)
    52  
    53  	// A release lock, and get wakeup list.
    54  	commitTSA := getTso()
    55  	wakeupList := make([]*Lock, 0)
    56  	lockA.SetCommitVer(commitTSA)
    57  	wakeupList = latches.release(lockA, wakeupList)
    58  	assert.Equal(wakeupList[0].startVer, startVerB)
    59  
    60  	// B acquire failed since startVerB has stale for some keys.
    61  	result = latches.acquire(lockB)
    62  	assert.Equal(result, acquireStale)
    63  
    64  	// B release lock since it received a stale.
    65  	wakeupList = latches.release(lockB, wakeupList)
    66  	assert.Equal(0, len(wakeupList))
    67  
    68  	// B restart:get a new startVer.
    69  	startVerB = getTso()
    70  	lockB = latches.genLock(startVerB, keysB)
    71  	result = latches.acquire(lockB)
    72  	assert.Equal(result, acquireSuccess)
    73  }
    74  
    75  func TestFirstAcquireFailedWithStale(t *testing.T) {
    76  	assert := assert.New(t)
    77  	latches := NewLatches(256)
    78  
    79  	keys := []kv.Key{
    80  		[]byte("a"), []byte("b"), []byte("c")}
    81  	startVerA := getTso()
    82  	lockA := latches.genLock(startVerA, keys)
    83  	startVerB := getTso()
    84  	lockB := latches.genLock(startVerA, keys)
    85  
    86  	// acquire lockA success
    87  	result := latches.acquire(lockA)
    88  	assert.Equal(result, acquireSuccess)
    89  
    90  	// release lockA
    91  	commitVerA := getTso()
    92  	wakeupList := make([]*Lock, 0)
    93  	lockA.SetCommitVer(commitVerA)
    94  	latches.release(lockA, wakeupList)
    95  
    96  	assert.Greater(commitVerA, startVerB)
    97  	// acquire lockB first time, should be failed with stale since commitVerA > startVerB
    98  	result = latches.acquire(lockB)
    99  	assert.Equal(result, acquireStale)
   100  	latches.release(lockB, wakeupList)
   101  }
   102  
   103  func TestRecycle(t *testing.T) {
   104  	assert := assert.New(t)
   105  	latches := NewLatches(8)
   106  	startVer := getTso()
   107  	lock := latches.genLock(startVer, []kv.Key{
   108  		[]byte("a"), []byte("b"),
   109  	})
   110  	lock1 := latches.genLock(startVer, []kv.Key{
   111  		[]byte("b"), []byte("c"),
   112  	})
   113  	assert.Equal(latches.acquire(lock), acquireSuccess)
   114  	assert.Equal(latches.acquire(lock1), acquireLocked)
   115  	lock.SetCommitVer(startVer + 1)
   116  	var wakeupList []*Lock
   117  	latches.release(lock, wakeupList)
   118  	// Release lock will grant latch to lock1 automatically,
   119  	// so release lock1 is called here.
   120  	latches.release(lock1, wakeupList)
   121  
   122  	lock2 := latches.genLock(startVer+3, []kv.Key{
   123  		[]byte("b"), []byte("c"),
   124  	})
   125  	assert.Equal(latches.acquire(lock2), acquireSuccess)
   126  	wakeupList = wakeupList[:0]
   127  	latches.release(lock2, wakeupList)
   128  
   129  	allEmpty := true
   130  	for i := 0; i < len(latches.slots); i++ {
   131  		latch := &latches.slots[i]
   132  		if latch.queue != nil {
   133  			allEmpty = false
   134  		}
   135  	}
   136  	assert.False(allEmpty)
   137  
   138  	currentTS := kv.Version(time.Now().Add(expireDuration).UnixNano()) + 3
   139  	latches.recycle(currentTS)
   140  
   141  	for i := 0; i < len(latches.slots); i++ {
   142  		latch := &latches.slots[i]
   143  		assert.Nil(latch.queue)
   144  	}
   145  }