go-ethereum/les/freeclient_test.go

// Copyright 2017 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.

// Package light implements on-demand retrieval capable state and chain objects
// for the Ethereum Light Client.
package les

import (
	"fmt"
	"math/rand"
	"sync"
	"testing"
	"time"

	"github.com/ethereum/go-ethereum/common/mclock"
	"github.com/ethereum/go-ethereum/ethdb"
)

func TestFreeClientPoolL10C100(t *testing.T) {
	testFreeClientPool(t, 10, 100)
}

func TestFreeClientPoolL40C200(t *testing.T) {
	testFreeClientPool(t, 40, 200)
}

func TestFreeClientPoolL100C300(t *testing.T) {
	testFreeClientPool(t, 100, 300)
}

const testFreeClientPoolTicks = 500000

func testFreeClientPool(t *testing.T, connLimit, clientCount int) {
	quit := make(chan struct{})
	clock := mclock.NewSimulatedClock(time.Second, 10)
	defer clock.Stop()
	var wg sync.WaitGroup
	db := ethdb.NewMemDatabase()

	pool := newFreeClientPool(db, connLimit, 10000, quit, &wg, clock)

	connected := make([]bool, clientCount)
	connTicks := make([]int, clientCount)
	disconnCh := make(chan int)
	tickCh := clock.After(time.Second)
	tickCounter := 0

	peerId := func(i int) string {
		return fmt.Sprintf("test peer #%d", i)
	}

	disconnFn := func(i int) func() {
		return func() {
			disconnCh <- i
		}
	}

	// pool should accept new peers up to its connected limit
	for i := 0; i < connLimit; i++ {
		if pool.connect(peerId(i), disconnFn(i)) {
			connected[i] = true
		} else {
			t.Errorf("Test peer #%d rejected", i)
		}
	}
	// since all accepted peers are new and should not be kicked out, the next one should be rejected
	if pool.connect(peerId(connLimit), disconnFn(connLimit)) {
		connected[connLimit] = true
		t.Errorf("Peer accepted over connected limit")
	}

	// randomly connect and disconnect peers, expect to have a similar total connection time at the end
loop:
	for {
		select {
		case <-tickCh:
			i := rand.Intn(clientCount)
			if connected[i] {
				pool.disconnect(peerId(i))
				connected[i] = false
				connTicks[i] += tickCounter
			} else {
				if pool.connect(peerId(i), disconnFn(i)) {
					connected[i] = true
					connTicks[i] -= tickCounter
				}
			}
			tickCounter++
			if tickCounter >= testFreeClientPoolTicks {
				break loop
			}
			tickCh = clock.After(time.Second)
		case i := <-disconnCh:
			pool.disconnect(peerId(i))
			if connected[i] {
				connTicks[i] += tickCounter
				connected[i] = false
			}
		}
	}

	expTicks := testFreeClientPoolTicks * connLimit / clientCount
	expMin := expTicks - expTicks/10
	expMax := expTicks + expTicks/10

	// check if the total connected time of peers are all in the expected range
	for i, c := range connected {
		if c {
			connTicks[i] += tickCounter
		}
		if connTicks[i] < expMin || connTicks[i] > expMax {
			t.Errorf("Total connected time of test node #%d (%d) outside expected range (%d to %d)", i, connTicks[i], expMin, expMax)
		}
	}

	// a previously unknown peer should be accepted now
	if !pool.connect("newPeer", func() {}) {
		t.Errorf("Previously unknown peer rejected")
	}

	// close and restart pool
	close(quit)
	wg.Wait()
	quit2 := make(chan struct{})
	var wg2 sync.WaitGroup
	pool = newFreeClientPool(db, connLimit, 10000, quit2, &wg2, clock)

	// try connecting all known peers (connLimit should be filled up)
	for i := 0; i < clientCount; i++ {
		pool.connect(peerId(i), func() {})
	}
	// expect pool to remember known nodes and kick out one of them to accept a new one
	if !pool.connect("newPeer2", func() {}) {
		t.Errorf("Previously unknown peer rejected after restarting pool")
	}
	close(quit2)
}