github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/storage/latch/scheduler_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  	"bytes"
    20  	"math/rand"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/sourcegraph/conc"
    25  	"github.com/vescale/zgraph/storage/kv"
    26  )
    27  
    28  func TestWithConcurrency(t *testing.T) {
    29  	sched := NewScheduler(7)
    30  	sched.Run()
    31  	defer sched.Close()
    32  	rand.Seed(time.Now().Unix())
    33  
    34  	ch := make(chan []kv.Key, 100)
    35  	const workerCount = 10
    36  	var wg conc.WaitGroup
    37  	for i := 0; i < workerCount; i++ {
    38  		wg.Go(func() {
    39  			for txn := range ch {
    40  				lock := sched.Lock(getTso(), txn)
    41  				if lock.IsStale() {
    42  					// Should restart the transaction or return error
    43  				} else {
    44  					lock.SetCommitVer(getTso())
    45  					// Do 2pc
    46  				}
    47  				sched.UnLock(lock)
    48  			}
    49  		})
    50  	}
    51  
    52  	for i := 0; i < 999; i++ {
    53  		ch <- generate()
    54  	}
    55  	close(ch)
    56  
    57  	wg.Wait()
    58  }
    59  
    60  // generate generates something like:
    61  // {[]byte("a"), []byte("b"), []byte("c")}
    62  // {[]byte("a"), []byte("d"), []byte("e"), []byte("f")}
    63  // {[]byte("e"), []byte("f"), []byte("g"), []byte("h")}
    64  // The data should not repeat in the sequence.
    65  func generate() []kv.Key {
    66  	table := []byte{'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'}
    67  	ret := make([]kv.Key, 0, 5)
    68  	chance := []int{100, 60, 40, 20}
    69  	for i := 0; i < len(chance); i++ {
    70  		needMore := rand.Intn(100) < chance[i]
    71  		if needMore {
    72  			randBytes := []byte{table[rand.Intn(len(table))]}
    73  			if !contains(randBytes, ret) {
    74  				ret = append(ret, randBytes)
    75  			}
    76  		}
    77  	}
    78  	return ret
    79  }
    80  
    81  func contains(x []byte, set []kv.Key) bool {
    82  	for _, y := range set {
    83  		if bytes.Equal(x, y) {
    84  			return true
    85  		}
    86  	}
    87  	return false
    88  }