Fixes bug with excessive rate limiting

beacon-chain/sync/BUILD.bazel

@@ -111,6 +111,7 @@ go_test(
         "subscriber_beacon_blocks_test.go",
         "subscriber_committee_index_beacon_attestation_test.go",
         "subscriber_test.go",
+        "sync_test.go",
         "validate_aggregate_proof_test.go",
         "validate_attester_slashing_test.go",
         "validate_beacon_blocks_test.go",

beacon-chain/sync/rpc_beacon_blocks_by_range.go

@@ -61,7 +61,6 @@ func (r *Service) beaconBlocksByRangeRPCHandler(ctx context.Context, msg interfa
 	)
 	for startSlot <= endReqSlot {
 		remainingBucketCapacity = r.blocksRateLimiter.Remaining(stream.Conn().RemotePeer().String())
-
 		if int64(allowedBlocksPerSecond) > remainingBucketCapacity {
 			r.p2p.Peers().IncrementBadResponses(stream.Conn().RemotePeer())
 			if r.p2p.Peers().IsBad(stream.Conn().RemotePeer()) {
@@ -75,7 +74,6 @@ func (r *Service) beaconBlocksByRangeRPCHandler(ctx context.Context, msg interfa
 			r.writeErrorResponseToStream(responseCodeInvalidRequest, rateLimitedError, stream)
 			return errors.New(rateLimitedError)
 		}
-		r.blocksRateLimiter.Add(stream.Conn().RemotePeer().String(), int64(allowedBlocksPerSecond))
 
 		// TODO(3147): Update this with reasonable constraints.
 		if endSlot-startSlot > rangeLimit || m.Step == 0 {
@@ -89,6 +87,12 @@ func (r *Service) beaconBlocksByRangeRPCHandler(ctx context.Context, msg interfa
 			return err
 		}
 
+		// Decrease allowed blocks capacity by the number of streamed blocks.
+		if startSlot <= endSlot {
+			r.blocksRateLimiter.Add(
+				stream.Conn().RemotePeer().String(), int64(1+(endSlot-startSlot)/m.Step))
+		}
+
 		// Recalculate start and end slots for the next batch to be returned to the remote peer.
 		startSlot = endSlot + m.Step
 		endSlot = startSlot + (m.Step * (uint64(allowedBlocksPerSecond) - 1))

beacon-chain/sync/rpc_beacon_blocks_by_range_test.go

@@ -33,13 +33,14 @@ func TestBeaconBlocksRPCHandler_ReturnsBlocks(t *testing.T) {
 	}
 
 	// Populate the database with blocks that would match the request.
-	for i := req.StartSlot; i < req.StartSlot+(req.Step*req.Count); i++ {
+	for i := req.StartSlot; i < req.StartSlot+(req.Step*req.Count); i += req.Step {
 		if err := d.SaveBlock(context.Background(), &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: i}}); err != nil {
 			t.Fatal(err)
 		}
 	}
 
-	r := &Service{p2p: p1, db: d, blocksRateLimiter: leakybucket.NewCollector(10000, 10000, false)}
+	// Start service with 160 as allowed blocks capacity (and almost zero capacity recovery).
+	r := &Service{p2p: p1, db: d, blocksRateLimiter: leakybucket.NewCollector(0.000001, int64(req.Count*10), false)}
 	pcl := protocol.ID("/testing")
 
 	var wg sync.WaitGroup
@@ -68,6 +69,13 @@ func TestBeaconBlocksRPCHandler_ReturnsBlocks(t *testing.T) {
 		t.Errorf("Unexpected error: %v", err)
 	}
 
+	// Make sure that rate limiter doesn't limit capacity exceedingly.
+	remainingCapacity := r.blocksRateLimiter.Remaining(p2.PeerID().String())
+	expectedCapacity := int64(req.Count*10 - req.Count)
+	if remainingCapacity != expectedCapacity {
+		t.Fatalf("Unexpected rate limiting capacity, expected: %v, got: %v", expectedCapacity, remainingCapacity)
+	}
+
 	if testutil.WaitTimeout(&wg, 1*time.Second) {
 		t.Fatal("Did not receive stream within 1 sec")
 	}

beacon-chain/sync/sync_test.go

@@ -0,0 +1,19 @@
+package sync
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/sirupsen/logrus"
+)
+
+func TestMain(m *testing.M) {
+	logrus.SetLevel(logrus.DebugLevel)
+	logrus.SetOutput(ioutil.Discard)
+
+	allowedBlocksPerSecond = 64
+	allowedBlocksBurst = int64(10 * allowedBlocksPerSecond)
+
+	os.Exit(m.Run())
+}