github.com/gnattishness/bazel-go-ethereum@v0.0.0-20190929123618-7022a154f56d/les/distributor_test.go (about)

     1  // Copyright 2017 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package les
    18  
    19  import (
    20  	"math/rand"
    21  	"sync"
    22  	"testing"
    23  	"time"
    24  
    25  	"github.com/ethereum/go-ethereum/common/mclock"
    26  )
    27  
    28  type testDistReq struct {
    29  	cost, procTime, order uint64
    30  	canSendTo             map[*testDistPeer]struct{}
    31  }
    32  
    33  func (r *testDistReq) getCost(dp distPeer) uint64 {
    34  	return r.cost
    35  }
    36  
    37  func (r *testDistReq) canSend(dp distPeer) bool {
    38  	_, ok := r.canSendTo[dp.(*testDistPeer)]
    39  	return ok
    40  }
    41  
    42  func (r *testDistReq) request(dp distPeer) func() {
    43  	return func() { dp.(*testDistPeer).send(r) }
    44  }
    45  
    46  type testDistPeer struct {
    47  	sent    []*testDistReq
    48  	sumCost uint64
    49  	lock    sync.RWMutex
    50  }
    51  
    52  func (p *testDistPeer) send(r *testDistReq) {
    53  	p.lock.Lock()
    54  	defer p.lock.Unlock()
    55  
    56  	p.sent = append(p.sent, r)
    57  	p.sumCost += r.cost
    58  }
    59  
    60  func (p *testDistPeer) worker(t *testing.T, checkOrder bool, stop chan struct{}) {
    61  	var last uint64
    62  	for {
    63  		wait := time.Millisecond
    64  		p.lock.Lock()
    65  		if len(p.sent) > 0 {
    66  			rq := p.sent[0]
    67  			wait = time.Duration(rq.procTime)
    68  			p.sumCost -= rq.cost
    69  			if checkOrder {
    70  				if rq.order <= last {
    71  					t.Errorf("Requests processed in wrong order")
    72  				}
    73  				last = rq.order
    74  			}
    75  			p.sent = p.sent[1:]
    76  		}
    77  		p.lock.Unlock()
    78  		select {
    79  		case <-stop:
    80  			return
    81  		case <-time.After(wait):
    82  		}
    83  	}
    84  }
    85  
    86  const (
    87  	testDistBufLimit       = 10000000
    88  	testDistMaxCost        = 1000000
    89  	testDistPeerCount      = 5
    90  	testDistReqCount       = 5000
    91  	testDistMaxResendCount = 3
    92  )
    93  
    94  func (p *testDistPeer) waitBefore(cost uint64) (time.Duration, float64) {
    95  	p.lock.RLock()
    96  	sumCost := p.sumCost + cost
    97  	p.lock.RUnlock()
    98  	if sumCost < testDistBufLimit {
    99  		return 0, float64(testDistBufLimit-sumCost) / float64(testDistBufLimit)
   100  	}
   101  	return time.Duration(sumCost - testDistBufLimit), 0
   102  }
   103  
   104  func (p *testDistPeer) canQueue() bool {
   105  	return true
   106  }
   107  
   108  func (p *testDistPeer) queueSend(f func()) {
   109  	f()
   110  }
   111  
   112  func TestRequestDistributor(t *testing.T) {
   113  	testRequestDistributor(t, false)
   114  }
   115  
   116  func TestRequestDistributorResend(t *testing.T) {
   117  	testRequestDistributor(t, true)
   118  }
   119  
   120  func testRequestDistributor(t *testing.T, resend bool) {
   121  	stop := make(chan struct{})
   122  	defer close(stop)
   123  
   124  	dist := newRequestDistributor(nil, &mclock.System{})
   125  	var peers [testDistPeerCount]*testDistPeer
   126  	for i := range peers {
   127  		peers[i] = &testDistPeer{}
   128  		go peers[i].worker(t, !resend, stop)
   129  		dist.registerTestPeer(peers[i])
   130  	}
   131  
   132  	var wg sync.WaitGroup
   133  
   134  	for i := 1; i <= testDistReqCount; i++ {
   135  		cost := uint64(rand.Int63n(testDistMaxCost))
   136  		procTime := uint64(rand.Int63n(int64(cost + 1)))
   137  		rq := &testDistReq{
   138  			cost:      cost,
   139  			procTime:  procTime,
   140  			order:     uint64(i),
   141  			canSendTo: make(map[*testDistPeer]struct{}),
   142  		}
   143  		for _, peer := range peers {
   144  			if rand.Intn(2) != 0 {
   145  				rq.canSendTo[peer] = struct{}{}
   146  			}
   147  		}
   148  
   149  		wg.Add(1)
   150  		req := &distReq{
   151  			getCost: rq.getCost,
   152  			canSend: rq.canSend,
   153  			request: rq.request,
   154  		}
   155  		chn := dist.queue(req)
   156  		go func() {
   157  			cnt := 1
   158  			if resend && len(rq.canSendTo) != 0 {
   159  				cnt = rand.Intn(testDistMaxResendCount) + 1
   160  			}
   161  			for i := 0; i < cnt; i++ {
   162  				if i != 0 {
   163  					chn = dist.queue(req)
   164  				}
   165  				p := <-chn
   166  				if p == nil {
   167  					if len(rq.canSendTo) != 0 {
   168  						t.Errorf("Request that could have been sent was dropped")
   169  					}
   170  				} else {
   171  					peer := p.(*testDistPeer)
   172  					if _, ok := rq.canSendTo[peer]; !ok {
   173  						t.Errorf("Request sent to wrong peer")
   174  					}
   175  				}
   176  			}
   177  			wg.Done()
   178  		}()
   179  		if rand.Intn(1000) == 0 {
   180  			time.Sleep(time.Duration(rand.Intn(5000000)))
   181  		}
   182  	}
   183  
   184  	wg.Wait()
   185  }