diff --git a/les/clientpool.go b/les/clientpool.go index 2df538620b..0b4d1b9612 100644 --- a/les/clientpool.go +++ b/les/clientpool.go @@ -459,6 +459,9 @@ func (f *clientPool) addBalance(id enode.ID, amount uint64, setTotal bool) { defer func() { c.balanceTracker.setBalance(pb.value, negBalance) if !c.priority && pb.value > 0 { + // The capacity should be adjusted based on the requirement, + // but we have no idea about the new capacity, need a second + // call to udpate it. c.priority = true c.balanceTracker.addCallback(balanceCallbackZero, 0, func() { f.balanceExhausted(id) }) } diff --git a/les/clientpool_test.go b/les/clientpool_test.go index af75d7bbda..53973696ca 100644 --- a/les/clientpool_test.go +++ b/les/clientpool_test.go @@ -68,6 +68,14 @@ func (i poolTestPeer) freeClientId() string { func (i poolTestPeer) updateCapacity(uint64) {} +type poolTestPeerWithCap struct { + poolTestPeer + + cap uint64 +} + +func (i *poolTestPeerWithCap) updateCapacity(cap uint64) { i.cap = cap } + func testClientPool(t *testing.T, connLimit, clientCount, paidCount int, randomDisconnect bool) { rand.Seed(time.Now().UnixNano()) var ( @@ -364,11 +372,20 @@ func TestDowngradePriorityClient(t *testing.T) { 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) + p := &poolTestPeerWithCap{ + poolTestPeer: poolTestPeer(0), + } + pool.addBalance(p.ID(), uint64(time.Minute), false) + pool.connect(p, 10) + if p.cap != 10 { + t.Fatalf("The capcacity of priority peer hasn't been updated, got: %d", p.cap) + } + clock.Run(time.Minute) // All positive balance should be used up. time.Sleep(300 * time.Millisecond) // Ensure the callback is called - + if p.cap != 1 { + t.Fatalf("The capcacity of peer should be downgraded, got: %d", p.cap) + } pb := pool.ndb.getOrNewPB(poolTestPeer(0).ID()) if pb.value != 0 { t.Fatalf("Positive balance mismatch, want %v, got %v", 0, pb.value) diff --git a/les/odr_test.go b/les/odr_test.go index 74808b345f..17b963d480 100644 --- a/les/odr_test.go +++ b/les/odr_test.go @@ -202,7 +202,7 @@ func testOdr(t *testing.T, protocol int, expFail uint64, checkCached bool, fn od bhash := rawdb.ReadCanonicalHash(server.db, i) b1 := fn(light.NoOdr, server.db, server.handler.server.chainConfig, server.handler.blockchain, nil, bhash) - ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) b2 := fn(ctx, client.db, client.handler.backend.chainConfig, nil, client.handler.backend.blockchain, bhash) cancel()