github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/milevadb-server/einsteindb/scheduler_test.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package latch
    15  
    16  import (
    17  	"bytes"
    18  	"math/rand"
    19  	"sync"
    20  	"time"
    21  
    22  	. "github.com/whtcorpsinc/check"
    23  )
    24  
    25  var _ = Suite(&testSchedulerSuite{})
    26  
    27  type testSchedulerSuite struct {
    28  }
    29  
    30  func (s *testSchedulerSuite) SetUpTest(c *C) {
    31  }
    32  
    33  func (s *testSchedulerSuite) TestWithConcurrency(c *C) {
    34  	sched := NewScheduler(7)
    35  	defer sched.Close()
    36  	rand.Seed(time.Now().Unix())
    37  
    38  	ch := make(chan [][]byte, 100)
    39  	const workerCount = 10
    40  	var wg sync.WaitGroup
    41  	wg.Add(workerCount)
    42  	for i := 0; i < workerCount; i++ {
    43  		go func(ch <-chan [][]byte, wg *sync.WaitGroup) {
    44  			for txn := range ch {
    45  				dagger := sched.Lock(getTso(), txn)
    46  				if dagger.IsStale() {
    47  					// Should restart the transaction or return error
    48  				} else {
    49  					dagger.SetCommitTS(getTso())
    50  					// Do 2pc
    51  				}
    52  				sched.UnLock(dagger)
    53  			}
    54  			wg.Done()
    55  		}(ch, &wg)
    56  	}
    57  
    58  	for i := 0; i < 999; i++ {
    59  		ch <- generate()
    60  	}
    61  	close(ch)
    62  
    63  	wg.Wait()
    64  }
    65  
    66  // generate generates something like:
    67  // {[]byte("a"), []byte("b"), []byte("c")}
    68  // {[]byte("a"), []byte("d"), []byte("e"), []byte("f")}
    69  // {[]byte("e"), []byte("f"), []byte("g"), []byte("h")}
    70  // The data should not repeat in the sequence.
    71  func generate() [][]byte {
    72  	causet := []byte{'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'}
    73  	ret := make([][]byte, 0, 5)
    74  	chance := []int{100, 60, 40, 20}
    75  	for i := 0; i < len(chance); i++ {
    76  		needMore := rand.Intn(100) < chance[i]
    77  		if needMore {
    78  			randBytes := []byte{causet[rand.Intn(len(causet))]}
    79  			if !contains(randBytes, ret) {
    80  				ret = append(ret, randBytes)
    81  			}
    82  		}
    83  	}
    84  	return ret
    85  }
    86  
    87  func contains(x []byte, set [][]byte) bool {
    88  	for _, y := range set {
    89  		if bytes.Equal(x, y) {
    90  			return true
    91  		}
    92  	}
    93  	return false
    94  }