les: add more tests

les/balance.go

@@ -67,7 +67,7 @@ type balanceCallback struct {
 // init initializes balanceTracker
 func (bt *balanceTracker) init(clock mclock.Clock, capacity uint64) {
 	bt.clock = clock
-	bt.initTime = clock.Now()
+	bt.initTime, bt.lastUpdate = clock.Now(), clock.Now() // Init timestamps
 	for i := range bt.callbackIndex {
 		bt.callbackIndex[i] = -1
 	}

les/balance_test.go

@@ -0,0 +1,260 @@
+// Copyright 2019 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 les
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common/mclock"
+)
+
+func TestSetBalance(t *testing.T) {
+	var clock = &mclock.Simulated{}
+	var inputs = []struct {
+		pos uint64
+		neg uint64
+	}{
+		{1000, 0},
+		{0, 1000},
+		{1000, 1000},
+	}
+
+	tracker := balanceTracker{}
+	tracker.init(clock, 1000)
+	defer tracker.stop(clock.Now())
+
+	for _, i := range inputs {
+		tracker.setBalance(i.pos, i.neg)
+		pos, neg := tracker.getBalance(clock.Now())
+		if pos != i.pos {
+			t.Fatalf("Positive balance mismatch, want %v, got %v", i.pos, pos)
+		}
+		if neg != i.neg {
+			t.Fatalf("Negative balance mismatch, want %v, got %v", i.neg, neg)
+		}
+	}
+}
+
+func TestBalanceTimeCost(t *testing.T) {
+	var (
+		clock   = &mclock.Simulated{}
+		tracker = balanceTracker{}
+	)
+	tracker.init(clock, 1000)
+	defer tracker.stop(clock.Now())
+	tracker.setFactors(false, 1, 1)
+	tracker.setFactors(true, 1, 1)
+
+	tracker.setBalance(uint64(time.Minute), 0) // 1 minute time allowance
+
+	var inputs = []struct {
+		runTime time.Duration
+		expPos  uint64
+		expNeg  uint64
+	}{
+		{time.Second, uint64(time.Second * 59), 0},
+		{0, uint64(time.Second * 59), 0},
+		{time.Second * 59, 0, 0},
+		{time.Second, 0, uint64(time.Second)},
+	}
+	for _, i := range inputs {
+		clock.Run(i.runTime)
+		if pos, _ := tracker.getBalance(clock.Now()); pos != i.expPos {
+			t.Fatalf("Positive balance mismatch, want %v, got %v", i.expPos, pos)
+		}
+		if _, neg := tracker.getBalance(clock.Now()); neg != i.expNeg {
+			t.Fatalf("Negative balance mismatch, want %v, got %v", i.expNeg, neg)
+		}
+	}
+
+	tracker.setBalance(uint64(time.Minute), 0) // Refill 1 minute time allowance
+	for _, i := range inputs {
+		clock.Run(i.runTime)
+		if pos, _ := tracker.getBalance(clock.Now()); pos != i.expPos {
+			t.Fatalf("Positive balance mismatch, want %v, got %v", i.expPos, pos)
+		}
+		if _, neg := tracker.getBalance(clock.Now()); neg != i.expNeg {
+			t.Fatalf("Negative balance mismatch, want %v, got %v", i.expNeg, neg)
+		}
+	}
+}
+
+func TestBalanceReqCost(t *testing.T) {
+	var (
+		clock   = &mclock.Simulated{}
+		tracker = balanceTracker{}
+	)
+	tracker.init(clock, 1000)
+	defer tracker.stop(clock.Now())
+	tracker.setFactors(false, 1, 1)
+	tracker.setFactors(true, 1, 1)
+
+	tracker.setBalance(uint64(time.Minute), 0) // 1 minute time serving time allowance
+	var inputs = []struct {
+		reqCost uint64
+		expPos  uint64
+		expNeg  uint64
+	}{
+		{uint64(time.Second), uint64(time.Second * 59), 0},
+		{0, uint64(time.Second * 59), 0},
+		{uint64(time.Second * 59), 0, 0},
+		{uint64(time.Second), 0, uint64(time.Second)},
+	}
+	for _, i := range inputs {
+		tracker.requestCost(i.reqCost)
+		if pos, _ := tracker.getBalance(clock.Now()); pos != i.expPos {
+			t.Fatalf("Positive balance mismatch, want %v, got %v", i.expPos, pos)
+		}
+		if _, neg := tracker.getBalance(clock.Now()); neg != i.expNeg {
+			t.Fatalf("Negative balance mismatch, want %v, got %v", i.expNeg, neg)
+		}
+	}
+}
+
+func TestBalanceToPriority(t *testing.T) {
+	var (
+		clock   = &mclock.Simulated{}
+		tracker = balanceTracker{}
+	)
+	tracker.init(clock, 1000) // cap = 1000
+	defer tracker.stop(clock.Now())
+	tracker.setFactors(false, 1, 1)
+	tracker.setFactors(true, 1, 1)
+
+	var inputs = []struct {
+		pos      uint64
+		neg      uint64
+		priority int64
+	}{
+		{1000, 0, ^int64(1)},
+		{2000, 0, ^int64(2)}, // Higher balance, lower priority value
+		{0, 0, 0},
+		{0, 1000, 1000},
+	}
+	for _, i := range inputs {
+		tracker.setBalance(i.pos, i.neg)
+		priority := tracker.getPriority(clock.Now())
+		if priority != i.priority {
+			t.Fatalf("Priority mismatch, want %v, got %v", i.priority, priority)
+		}
+	}
+}
+
+func TestEstimatedPriority(t *testing.T) {
+	var (
+		clock   = &mclock.Simulated{}
+		tracker = balanceTracker{}
+	)
+	tracker.init(clock, 1000000000) // cap = 1000,000,000
+	defer tracker.stop(clock.Now())
+	tracker.setFactors(false, 1, 1)
+	tracker.setFactors(true, 1, 1)
+
+	tracker.setBalance(uint64(time.Minute), 0)
+	var inputs = []struct {
+		runTime    time.Duration // time cost
+		futureTime time.Duration // diff of future time
+		reqCost    uint64        // single request cost
+		priority   int64         // expected estimated priority
+	}{
+		{time.Second, time.Second, 0, ^int64(58)},
+		{0, time.Second, 0, ^int64(58)},
+
+		// 2 seconds time cost, 1 second estimated time cost, 10^9 request cost,
+		// 10^9 estimated request cost per second.
+		{time.Second, time.Second, 1000000000, ^int64(55)},
+
+		// 3 seconds time cost, 3 second estimated time cost, 10^9*2 request cost,
+		// 4*10^9 estimated request cost.
+		{time.Second, 3*time.Second, 1000000000, ^int64(48)},
+
+		// All positive balance is used up
+		{time.Second*55, 0, 0, 0},
+
+		// 1 minute estimated time cost, 4/58 * 10^9 estimated request cost per sec.
+		{0, time.Minute, 0, int64(time.Minute) + int64(time.Second) * 120/29},
+	}
+	for _, i := range inputs {
+		clock.Run(i.runTime)
+		tracker.requestCost(i.reqCost)
+		priority := tracker.estimatedPriority(clock.Now()+mclock.AbsTime(i.futureTime), true)
+		if priority != i.priority {
+			t.Fatalf("Estimated priority mismatch, want %v, got %v", i.priority, priority)
+		}
+	}
+}
+
+func TestCallbackChecking(t *testing.T) {
+	var (
+		clock   = &mclock.Simulated{}
+		tracker = balanceTracker{}
+	)
+	tracker.init(clock, 1000000) // cap = 1000,000
+	defer tracker.stop(clock.Now())
+	tracker.setFactors(false, 1, 1)
+	tracker.setFactors(true, 1, 1)
+
+	var inputs = []struct{
+		priority int64
+		expDiff  time.Duration
+	} {
+		{^int64(500), time.Millisecond*500},
+		{0, time.Second},
+		{int64(time.Second), 2*time.Second},
+	}
+	tracker.setBalance(uint64(time.Second), 0)
+	for _, i := range inputs {
+		diff, _ := tracker.timeUntil(i.priority)
+		if diff != i.expDiff {
+			t.Fatalf("Time difference mismatch, want %v, got %v", i.expDiff, diff)
+		}
+	}
+}
+
+func TestCallback(t *testing.T) {
+	var (
+		clock   = &mclock.Simulated{}
+		tracker = balanceTracker{}
+	)
+	tracker.init(clock, 1000) // cap = 1000
+	defer tracker.stop(clock.Now())
+	tracker.setFactors(false, 1, 1)
+	tracker.setFactors(true, 1, 1)
+
+	callCh := make(chan struct{}, 1)
+	tracker.setBalance(uint64(time.Minute), 0)
+	tracker.addCallback(balanceCallbackZero, 0, func() {callCh <- struct{}{}})
+
+	clock.Run(time.Minute)
+	select {
+	case <-callCh:
+	case <-time.NewTimer(time.Second).C:
+		t.Fatalf("Callback hasn't been called yet")
+	}
+
+	tracker.setBalance(uint64(time.Minute), 0)
+	tracker.addCallback(balanceCallbackZero, 0, func() {callCh <- struct{}{}})
+	tracker.removeCallback(balanceCallbackZero)
+
+	clock.Run(time.Minute)
+	select {
+	case <-callCh:
+		t.Fatalf("Callback shouldn't be called")
+	case <-time.NewTimer(time.Millisecond*100).C:
+	}
+}

les/clientpool.go

@@ -44,14 +44,14 @@ const (
 	// freeConnectedBias is applied to already connected clients when new client
 	// is "free client". So that already connected client won't be kicked out very
 	// soon.
-	freeConnectedBias = time.Minute * 5
+	freeConnectedBias = time.Minute * 3
 
 	// priorityConnectedBias is applied to already connected clients when new client
 	// is "priority client". The value is smaller than freeConnectedBias, so that
 	// we can ensure most of new priority client can be accepted when pool is full.
 	// But if the balance of priority client is very small, there is no reason for
 	// very high priority.
-	priorityConnectedBias = time.Minute
+	priorityConnectedBias = time.Second * 30
 )
 
 // clientPool implements a client database that assigns a priority to each client
@@ -291,14 +291,16 @@ func (f *clientPool) connect(peer clientPeer, capacity uint64) bool {
 			f.dropClient(c, f.clock.Now(), true)
 		}
 	}
-	if e.priority {
+	// Register new client to connection queue.
-		e.balanceTracker.addCallback(balanceCallbackZero, 0, func() { f.balanceExhausted(id) })
-	}
 	f.connectedMap[id] = e
 	f.connectedQueue.Push(e)
 	f.connectedCap += e.capacity
+	// If the current client is a paid client, notify it to update the capacity.
+	// And also monitor the status of client, downgrade it to normal client if
+	// positive balance is used up.
 	if e.capacity != f.freeClientCap {
 		e.peer.updateCapacity(e.capacity)
+		e.balanceTracker.addCallback(balanceCallbackZero, 0, func() { f.balanceExhausted(id) })
 	}
 	totalConnectedGauge.Update(int64(f.connectedCap))
 	clientConnectedMeter.Mark(1)
@@ -313,14 +315,14 @@ func (f *clientPool) disconnect(p clientPeer) {
 	f.lock.Lock()
 	defer f.lock.Unlock()
 
+	// Short circuit if client pool is already closed.
 	if f.closed {
 		return
 	}
-	id := p.ID()
 	// Short circuit if the peer hasn't been registered.
-	e := f.connectedMap[id]
+	e := f.connectedMap[p.ID()]
 	if e == nil {
-		log.Debug("Client not connected", "address", p.freeClientId(), "id", peerIdToString(id))
+		log.Debug("Client not connected", "address", p.freeClientId(), "id", peerIdToString(p.ID()))
 		return
 	}
 	f.dropClient(e, f.clock.Now(), false)
@@ -383,6 +385,7 @@ func (f *clientPool) balanceExhausted(id enode.ID) {
 		c.capacity = f.freeClientCap
 		c.peer.updateCapacity(c.capacity)
 	}
+	f.ndb.delPB(id)
 }
 
 // setConnLimit sets the maximum number and total capacity of connected clients,
@@ -545,6 +548,7 @@ type nodeDB struct {
 	nbEvictCallBack func(mclock.AbsTime, negBalance) bool // Callback to determine whether the negative balance can be evicted.
 	clock           mclock.Clock
 	closeCh         chan struct{}
+	cleanupHook     func() // Test hook used for testing
 }
 
 func newNodeDB(db ethdb.Database, clock mclock.Clock) *nodeDB {
@@ -572,7 +576,9 @@ func (db *nodeDB) key(id []byte, neg bool) []byte {
 	if neg {
 		prefix = negativeBalancePrefix
 	}
-	db.auxbuf = db.auxbuf[:0]
+	if len(prefix)+len(db.verbuf)+len(id) > len(db.auxbuf) {
+		db.auxbuf = append(db.auxbuf, make([]byte, len(prefix)+len(db.verbuf)+len(id)-len(db.auxbuf))...)
+	}
 	copy(db.auxbuf[:len(db.verbuf)], db.verbuf[:])
 	copy(db.auxbuf[len(db.verbuf):len(db.verbuf)+len(prefix)], prefix)
 	copy(db.auxbuf[len(prefix)+len(db.verbuf):len(prefix)+len(db.verbuf)+len(id)], id)
@@ -693,5 +699,9 @@ func (db *nodeDB) expireNodes() {
 			db.db.Delete(iter.Key())
 		}
 	}
+	// Invoke testing hook if it's not nil.
+	if db.cleanupHook != nil {
+		db.cleanupHook()
+	}
 	log.Debug("Expire nodes", "visited", visited, "deleted", deleted, "elapsed", common.PrettyDuration(time.Since(start)))
 }

les/clientpool_test.go

@@ -54,7 +54,7 @@ func TestClientPoolL100C300P20(t *testing.T) {
 	testClientPool(t, 100, 300, 20, false)
 }
 
-const testClientPoolTicks = 500000
+const testClientPoolTicks = 100000
 
 type poolTestPeer int
 
@@ -134,11 +134,11 @@ func testClientPool(t *testing.T, connLimit, clientCount, paidCount int, randomD
 	}
 
 	expTicks := testClientPoolTicks/2*connLimit/clientCount + testClientPoolTicks/2*(connLimit-paidCount)/(clientCount-paidCount)
-	expMin := expTicks - expTicks/10
+	expMin := expTicks - expTicks/5
-	expMax := expTicks + expTicks/10
+	expMax := expTicks + expTicks/5
 	paidTicks := testClientPoolTicks/2*connLimit/clientCount + testClientPoolTicks/2
-	paidMin := paidTicks - paidTicks/10
+	paidMin := paidTicks - paidTicks/5
-	paidMax := paidTicks + paidTicks/10
+	paidMax := paidTicks + paidTicks/5
 
 	// check if the total connected time of peers are all in the expected range
 	for i, c := range connected {
@@ -352,6 +352,35 @@ func TestPositiveBalanceCalculation(t *testing.T) {
 	}
 }
 
+func TestDowngradePriorityClient(t *testing.T) {
+	var (
+		clock  mclock.Simulated
+		db     = rawdb.NewMemoryDatabase()
+		kicked = make(chan int, 10)
+	)
+	removeFn := func(id enode.ID) { kicked <- int(id[0]) } // Noop
+	pool := newClientPool(db, 1, &clock, removeFn)
+	defer pool.stop()
+	pool.setLimits(10, uint64(10)) // Total capacity limit is 10
+	pool.setPriceFactors(priceFactors{1, 0, 1}, priceFactors{1, 0, 1})
+
+	pool.addBalance(poolTestPeer(0).ID(), uint64(time.Minute), false)
+	pool.connect(poolTestPeer(0), 10)
+	clock.Run(time.Minute)             // All positive balance should be used up.
+	time.Sleep(300 * time.Millisecond) // Ensure the callback is called
+
+	pb := pool.ndb.getOrNewPB(poolTestPeer(0).ID())
+	if pb.value != 0 {
+		t.Fatalf("Positive balance mismatch, want %v, got %v", 0, pb.value)
+	}
+
+	pool.addBalance(poolTestPeer(0).ID(), uint64(time.Minute), false)
+	pb = pool.ndb.getOrNewPB(poolTestPeer(0).ID())
+	if pb.value != uint64(time.Minute) {
+		t.Fatalf("Positive balance mismatch, want %v, got %v", uint64(time.Minute), pb.value)
+	}
+}
+
 func TestNegativeBalanceCalculation(t *testing.T) {
 	var (
 		clock  mclock.Simulated
@@ -443,7 +472,10 @@ func TestNodeDB(t *testing.T) {
 }
 
 func TestNodeDBExpiration(t *testing.T) {
-	var iterated int
+	var (
+		iterated int
+		done     = make(chan struct{}, 1)
+	)
 	callback := func(now mclock.AbsTime, b negBalance) bool {
 		iterated += 1
 		return true
@@ -452,6 +484,7 @@ func TestNodeDBExpiration(t *testing.T) {
 	ndb := newNodeDB(rawdb.NewMemoryDatabase(), clock)
 	defer ndb.close()
 	ndb.nbEvictCallBack = callback
+	ndb.cleanupHook = func() { done <- struct{}{} }
 
 	var cases = []struct {
 		ip      string
@@ -466,7 +499,11 @@ func TestNodeDBExpiration(t *testing.T) {
 		ndb.setNB(c.ip, c.balance)
 	}
 	clock.Run(time.Hour + time.Minute)
-	time.Sleep(300 * time.Millisecond)
+	select {
+	case <-done:
+	case <-time.NewTimer(time.Second).C:
+		t.Fatalf("timeout")
+	}
 	if iterated != 4 {
 		t.Fatalf("Failed to evict useless negative balances, want %v, got %d", 4, iterated)
 	}
@@ -475,7 +512,11 @@ func TestNodeDBExpiration(t *testing.T) {
 		ndb.setNB(c.ip, c.balance)
 	}
 	clock.Run(time.Hour + time.Minute)
-	time.Sleep(300 * time.Millisecond)
+	select {
+	case <-done:
+	case <-time.NewTimer(time.Second).C:
+		t.Fatalf("timeout")
+	}
 	if iterated != 8 {
 		t.Fatalf("Failed to evict useless negative balances, want %v, got %d", 4, iterated)
 	}