github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/les/distributor_test.go (about)

     1  // Package light implements on-demand retrieval capable state and chain objects
     2  // for the Ethereum Light Client.
     3  package les
     4  
     5  import (
     6  	"math/rand"
     7  	"sync"
     8  	"testing"
     9  	"time"
    10  )
    11  
    12  type testDistReq struct {
    13  	cost, procTime, order uint64
    14  	canSendTo             map[*testDistPeer]struct{}
    15  }
    16  
    17  func (r *testDistReq) getCost(dp distPeer) uint64 {
    18  	return r.cost
    19  }
    20  
    21  func (r *testDistReq) canSend(dp distPeer) bool {
    22  	_, ok := r.canSendTo[dp.(*testDistPeer)]
    23  	return ok
    24  }
    25  
    26  func (r *testDistReq) request(dp distPeer) func() {
    27  	return func() { dp.(*testDistPeer).send(r) }
    28  }
    29  
    30  type testDistPeer struct {
    31  	sent    []*testDistReq
    32  	sumCost uint64
    33  	lock    sync.RWMutex
    34  }
    35  
    36  func (p *testDistPeer) send(r *testDistReq) {
    37  	p.lock.Lock()
    38  	defer p.lock.Unlock()
    39  
    40  	p.sent = append(p.sent, r)
    41  	p.sumCost += r.cost
    42  }
    43  
    44  func (p *testDistPeer) worker(t *testing.T, checkOrder bool, stop chan struct{}) {
    45  	var last uint64
    46  	for {
    47  		wait := time.Millisecond
    48  		p.lock.Lock()
    49  		if len(p.sent) > 0 {
    50  			rq := p.sent[0]
    51  			wait = time.Duration(rq.procTime)
    52  			p.sumCost -= rq.cost
    53  			if checkOrder {
    54  				if rq.order <= last {
    55  					t.Errorf("Requests processed in wrong order")
    56  				}
    57  				last = rq.order
    58  			}
    59  			p.sent = p.sent[1:]
    60  		}
    61  		p.lock.Unlock()
    62  		select {
    63  		case <-stop:
    64  			return
    65  		case <-time.After(wait):
    66  		}
    67  	}
    68  }
    69  
    70  const (
    71  	testDistBufLimit       = 10000000
    72  	testDistMaxCost        = 1000000
    73  	testDistPeerCount      = 5
    74  	testDistReqCount       = 50000
    75  	testDistMaxResendCount = 3
    76  )
    77  
    78  func (p *testDistPeer) waitBefore(cost uint64) (time.Duration, float64) {
    79  	p.lock.RLock()
    80  	sumCost := p.sumCost + cost
    81  	p.lock.RUnlock()
    82  	if sumCost < testDistBufLimit {
    83  		return 0, float64(testDistBufLimit-sumCost) / float64(testDistBufLimit)
    84  	} else {
    85  		return time.Duration(sumCost - testDistBufLimit), 0
    86  	}
    87  }
    88  
    89  func (p *testDistPeer) canQueue() bool {
    90  	return true
    91  }
    92  
    93  func (p *testDistPeer) queueSend(f func()) {
    94  	f()
    95  }
    96  
    97  func TestRequestDistributor(t *testing.T) {
    98  	testRequestDistributor(t, false)
    99  }
   100  
   101  func TestRequestDistributorResend(t *testing.T) {
   102  	testRequestDistributor(t, true)
   103  }
   104  
   105  func testRequestDistributor(t *testing.T, resend bool) {
   106  	stop := make(chan struct{})
   107  	defer close(stop)
   108  
   109  	dist := newRequestDistributor(nil, stop)
   110  	var peers [testDistPeerCount]*testDistPeer
   111  	for i := range peers {
   112  		peers[i] = &testDistPeer{}
   113  		go peers[i].worker(t, !resend, stop)
   114  		dist.registerTestPeer(peers[i])
   115  	}
   116  
   117  	var wg sync.WaitGroup
   118  
   119  	for i := 1; i <= testDistReqCount; i++ {
   120  		cost := uint64(rand.Int63n(testDistMaxCost))
   121  		procTime := uint64(rand.Int63n(int64(cost + 1)))
   122  		rq := &testDistReq{
   123  			cost:      cost,
   124  			procTime:  procTime,
   125  			order:     uint64(i),
   126  			canSendTo: make(map[*testDistPeer]struct{}),
   127  		}
   128  		for _, peer := range peers {
   129  			if rand.Intn(2) != 0 {
   130  				rq.canSendTo[peer] = struct{}{}
   131  			}
   132  		}
   133  
   134  		wg.Add(1)
   135  		req := &distReq{
   136  			getCost: rq.getCost,
   137  			canSend: rq.canSend,
   138  			request: rq.request,
   139  		}
   140  		chn := dist.queue(req)
   141  		go func() {
   142  			cnt := 1
   143  			if resend && len(rq.canSendTo) != 0 {
   144  				cnt = rand.Intn(testDistMaxResendCount) + 1
   145  			}
   146  			for i := 0; i < cnt; i++ {
   147  				if i != 0 {
   148  					chn = dist.queue(req)
   149  				}
   150  				p := <-chn
   151  				if p == nil {
   152  					if len(rq.canSendTo) != 0 {
   153  						t.Errorf("Request that could have been sent was dropped")
   154  					}
   155  				} else {
   156  					peer := p.(*testDistPeer)
   157  					if _, ok := rq.canSendTo[peer]; !ok {
   158  						t.Errorf("Request sent to wrong peer")
   159  					}
   160  				}
   161  			}
   162  			wg.Done()
   163  		}()
   164  		if rand.Intn(1000) == 0 {
   165  			time.Sleep(time.Duration(rand.Intn(5000000)))
   166  		}
   167  	}
   168  
   169  	wg.Wait()
   170  }