Limit ibft msg queues

config/src/main/java/tech/pegasys/pantheon/config/IbftConfigOptions.java

@@ -21,6 +21,7 @@ public class IbftConfigOptions {
   private static final long DEFAULT_EPOCH_LENGTH = 30_000;
   private static final int DEFAULT_BLOCK_PERIOD_SECONDS = 1;
   private static final int DEFAULT_ROUND_EXPIRY_SECONDS = 1;
+  private static final int DEFAULT_MESSAGE_BUFFER_SIZE = 10_000;
 
   private final JsonObject ibftConfigRoot;
 
@@ -39,4 +40,8 @@ public int getBlockPeriodSeconds() {
   public int getRequestTimeoutSeconds() {
     return ibftConfigRoot.getInteger("requesttimeoutseconds", DEFAULT_ROUND_EXPIRY_SECONDS);
   }
+
+  public int getMessageBufferSize() {
+    return ibftConfigRoot.getInteger("messagebuffersize", DEFAULT_MESSAGE_BUFFER_SIZE);
+  }
 }

config/src/test/java/tech/pegasys/pantheon/config/IbftConfigOptionsTest.java

@@ -26,6 +26,7 @@ public class IbftConfigOptionsTest {
   private static final int EXPECTED_DEFAULT_EPOCH_LENGTH = 30_000;
   private static final int EXPECTED_DEFAULT_BLOCK_PERIOD = 1;
   private static final int EXPECTED_DEFAULT_REQUEST_TIMEOUT = 1;
+  private static final int EXPECTED_DEFAULT_MESSAGE_BUFFER_SIZE = 10_000;
 
   @Test
   public void shouldGetEpochLengthFromConfig() {
@@ -80,6 +81,24 @@ public void shouldGetDefaultRequestTimeoutFromDefaultConfig() {
         .isEqualTo(EXPECTED_DEFAULT_REQUEST_TIMEOUT);
   }
 
+  @Test
+  public void shouldGetMessageBufferSizeFromConfig() {
+    final IbftConfigOptions config = fromConfigOptions(singletonMap("MessageBufferSize", 5_000));
+    assertThat(config.getMessageBufferSize()).isEqualTo(5_000);
+  }
+
+  @Test
+  public void shouldFallbackToDefaultMessageBufferSize() {
+    final IbftConfigOptions config = fromConfigOptions(emptyMap());
+    assertThat(config.getMessageBufferSize()).isEqualTo(EXPECTED_DEFAULT_MESSAGE_BUFFER_SIZE);
+  }
+
+  @Test
+  public void shouldGetDefaultMessageBufferSizeFromDefaultConfig() {
+    assertThat(IbftConfigOptions.DEFAULT.getMessageBufferSize())
+        .isEqualTo(EXPECTED_DEFAULT_MESSAGE_BUFFER_SIZE);
+  }
+
   private IbftConfigOptions fromConfigOptions(final Map<String, Object> ibftConfigOptions) {
     return GenesisConfigFile.fromConfig(
             new JsonObject(singletonMap("config", singletonMap("ibft", ibftConfigOptions))))

consensus/ibft/src/integration-test/java/tech/pegasys/pantheon/consensus/ibft/support/TestContextBuilder.java

@@ -112,6 +112,7 @@ public EventMultiplexer getEventMultiplexer() {
   public static final int EPOCH_LENGTH = 10_000;
   public static final int BLOCK_TIMER_SEC = 3;
   public static final int ROUND_TIMER_SEC = 12;
+  public static final int MESSAGE_BUFFER_SIZE = 10_000;
 
   private Clock clock = Clock.fixed(Instant.MIN, ZoneId.of("UTC"));
   private IbftEventQueue ibftEventQueue = new IbftEventQueue();
@@ -158,7 +159,8 @@ public TestContext build() {
     // Use a stubbed version of the multicaster, to prevent creating PeerConnections etc.
     final StubValidatorMulticaster multicaster = new StubValidatorMulticaster();
 
-    final UniqueMessageMulticaster uniqueMulticaster = new UniqueMessageMulticaster(multicaster);
+    final UniqueMessageMulticaster uniqueMulticaster =
+        new UniqueMessageMulticaster(multicaster, MESSAGE_BUFFER_SIZE);
 
     final Gossiper gossiper = useGossip ? new IbftGossip(uniqueMulticaster) : mock(Gossiper.class);
 
@@ -301,7 +303,8 @@ private static ControllerAndState createControllerAndFinalState(
                 finalState,
                 new IbftRoundFactory(
                     finalState, protocolContext, protocolSchedule, minedBlockObservers),
-                messageValidatorFactory),
+                messageValidatorFactory,
+                MESSAGE_BUFFER_SIZE),
             gossiper,
             new HashMap<>());
 

consensus/ibft/src/main/java/tech/pegasys/pantheon/consensus/ibft/SizeLimitedMap.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.consensus.ibft;
+
+import java.util.LinkedHashMap;
+import java.util.Map.Entry;
+
+/**
+ * Map that is limited to a specified size and will evict oldest entries when the size limit is
+ * reached.
+ */
+public class SizeLimitedMap<K, V> extends LinkedHashMap<K, V> {
+  private final int maxEntries;
+
+  public SizeLimitedMap(final int maxEntries) {
+    this.maxEntries = maxEntries;
+  }
+
+  @Override
+  protected boolean removeEldestEntry(final Entry<K, V> ignored) {
+    return size() > maxEntries;
+  }
+}

consensus/ibft/src/main/java/tech/pegasys/pantheon/consensus/ibft/UniqueMessageMulticaster.java

@@ -12,45 +12,33 @@
  */
 package tech.pegasys.pantheon.consensus.ibft;
 
+import static java.util.Collections.newSetFromMap;
+
 import tech.pegasys.pantheon.consensus.ibft.network.ValidatorMulticaster;
 import tech.pegasys.pantheon.ethereum.core.Address;
 import tech.pegasys.pantheon.ethereum.p2p.api.MessageData;
 
 import java.util.Collection;
 import java.util.Collections;
-import java.util.LinkedHashMap;
-import java.util.Map;
 import java.util.Set;
 
 public class UniqueMessageMulticaster implements ValidatorMulticaster {
-
-  private final int maxSeenMessages;
   private final ValidatorMulticaster multicaster;
-
-  UniqueMessageMulticaster(final ValidatorMulticaster multicaster, final int maxSeenMessages) {
-    this.maxSeenMessages = maxSeenMessages;
-    this.multicaster = multicaster;
-  }
+  private final Set<Integer> seenMessages;
 
   /**
    * Constructor that attaches gossip logic to a set of multicaster
    *
    * @param multicaster Network connections to the remote validators
+   * @param maxSeenMessages Maximum messages to track as seen
    */
-  public UniqueMessageMulticaster(final ValidatorMulticaster multicaster) {
-    this(multicaster, 10_000);
+  public UniqueMessageMulticaster(
+      final ValidatorMulticaster multicaster, final int maxSeenMessages) {
+    this.multicaster = multicaster;
+    // Set that starts evicting members when it hits capacity
+    this.seenMessages = newSetFromMap(new SizeLimitedMap<>(maxSeenMessages));
   }
 
-  // Set that starts evicting members when it hits capacity
-  private final Set<Integer> seenMessages =
-      Collections.newSetFromMap(
-          new LinkedHashMap<Integer, Boolean>() {
-            @Override
-            protected boolean removeEldestEntry(final Map.Entry<Integer, Boolean> eldest) {
-              return size() > maxSeenMessages;
-            }
-          });
-
   @Override
   public void send(final MessageData message) {
     send(message, Collections.emptyList());

consensus/ibft/src/main/java/tech/pegasys/pantheon/consensus/ibft/statemachine/IbftBlockHeightManager.java

@@ -19,6 +19,7 @@
 import tech.pegasys.pantheon.consensus.ibft.BlockTimer;
 import tech.pegasys.pantheon.consensus.ibft.ConsensusRoundIdentifier;
 import tech.pegasys.pantheon.consensus.ibft.RoundTimer;
+import tech.pegasys.pantheon.consensus.ibft.SizeLimitedMap;
 import tech.pegasys.pantheon.consensus.ibft.ibftevent.RoundExpiry;
 import tech.pegasys.pantheon.consensus.ibft.network.IbftMessageTransmitter;
 import tech.pegasys.pantheon.consensus.ibft.payload.CommitPayload;
@@ -43,7 +44,6 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 
-import com.google.common.collect.Maps;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -64,7 +64,7 @@ public class IbftBlockHeightManager implements BlockHeightManager {
   private final BlockTimer blockTimer;
   private final IbftMessageTransmitter transmitter;
   private final MessageFactory messageFactory;
-  private final Map<Integer, RoundState> futureRoundStateBuffer = Maps.newHashMap();
+  private final Map<Integer, RoundState> futureRoundStateBuffer;
   private final NewRoundMessageValidator newRoundMessageValidator;
   private final Clock clock;
   private final Function<ConsensusRoundIdentifier, RoundState> roundStateCreator;
@@ -80,7 +80,8 @@ public IbftBlockHeightManager(
       final RoundChangeManager roundChangeManager,
       final IbftRoundFactory ibftRoundFactory,
       final Clock clock,
-      final MessageValidatorFactory messageValidatorFactory) {
+      final MessageValidatorFactory messageValidatorFactory,
+      final int maxFutureRoundStateBufferSize) {
     this.parentHeader = parentHeader;
     this.roundFactory = ibftRoundFactory;
     this.roundTimer = finalState.getRoundTimer();
@@ -91,6 +92,7 @@ public IbftBlockHeightManager(
     this.roundChangeManager = roundChangeManager;
     this.finalState = finalState;
 
+    futureRoundStateBuffer = new SizeLimitedMap<>(maxFutureRoundStateBufferSize);
     newRoundMessageValidator = messageValidatorFactory.createNewRoundValidator(parentHeader);
 
     roundStateCreator =

consensus/ibft/src/main/java/tech/pegasys/pantheon/consensus/ibft/statemachine/IbftBlockHeightManagerFactory.java

@@ -21,14 +21,17 @@ public class IbftBlockHeightManagerFactory {
   private final IbftRoundFactory roundFactory;
   private final IbftFinalState finalState;
   private final MessageValidatorFactory messageValidatorFactory;
+  private final int messageBufferSize;
 
   public IbftBlockHeightManagerFactory(
       final IbftFinalState finalState,
       final IbftRoundFactory roundFactory,
-      final MessageValidatorFactory messageValidatorFactory) {
+      final MessageValidatorFactory messageValidatorFactory,
+      final int messageBufferSize) {
     this.roundFactory = roundFactory;
     this.finalState = finalState;
     this.messageValidatorFactory = messageValidatorFactory;
+    this.messageBufferSize = messageBufferSize;
   }
 
   public BlockHeightManager create(final BlockHeader parentHeader) {
@@ -52,6 +55,7 @@ private BlockHeightManager createFullBlockHeightManager(final BlockHeader parent
             messageValidatorFactory.createRoundChangeMessageValidator(parentHeader)),
         roundFactory,
         finalState.getClock(),
-        messageValidatorFactory);
+        messageValidatorFactory,
+        messageBufferSize);
   }
 }

consensus/ibft/src/main/java/tech/pegasys/pantheon/consensus/ibft/statemachine/IbftController.java

@@ -17,6 +17,7 @@
 import tech.pegasys.pantheon.consensus.ibft.ConsensusRoundIdentifier;
 import tech.pegasys.pantheon.consensus.ibft.Gossiper;
 import tech.pegasys.pantheon.consensus.ibft.IbftGossip;
+import tech.pegasys.pantheon.consensus.ibft.SizeLimitedMap;
 import tech.pegasys.pantheon.consensus.ibft.ibftevent.BlockTimerExpiry;
 import tech.pegasys.pantheon.consensus.ibft.ibftevent.IbftReceivedMessageEvent;
 import tech.pegasys.pantheon.consensus.ibft.ibftevent.NewChainHead;
@@ -40,7 +41,6 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -58,8 +58,14 @@ public IbftController(
       final Blockchain blockchain,
       final IbftFinalState ibftFinalState,
       final IbftBlockHeightManagerFactory ibftBlockHeightManagerFactory,
-      final IbftGossip gossiper) {
-    this(blockchain, ibftFinalState, ibftBlockHeightManagerFactory, gossiper, Maps.newHashMap());
+      final IbftGossip gossiper,
+      final int messageBufferSize) {
+    this(
+        blockchain,
+        ibftFinalState,
+        ibftBlockHeightManagerFactory,
+        gossiper,
+        new SizeLimitedMap<>(messageBufferSize));
   }
 
   @VisibleForTesting

consensus/ibft/src/main/java/tech/pegasys/pantheon/consensus/ibft/statemachine/RoundChangeManager.java

@@ -13,6 +13,7 @@
 package tech.pegasys.pantheon.consensus.ibft.statemachine;
 
 import tech.pegasys.pantheon.consensus.ibft.ConsensusRoundIdentifier;
+import tech.pegasys.pantheon.consensus.ibft.SizeLimitedMap;
 import tech.pegasys.pantheon.consensus.ibft.payload.RoundChangeCertificate;
 import tech.pegasys.pantheon.consensus.ibft.payload.RoundChangePayload;
 import tech.pegasys.pantheon.consensus.ibft.payload.SignedData;
@@ -44,7 +45,8 @@ public static class RoundChangeStatus {
 
     // Store only 1 round change per round per validator
     @VisibleForTesting
-    final Map<Address, SignedData<RoundChangePayload>> receivedMessages = Maps.newLinkedHashMap();
+    final Map<Address, SignedData<RoundChangePayload>> receivedMessages =
+        new SizeLimitedMap<>(10_000);
 
     private boolean actioned = false;
 

consensus/ibft/src/test/java/tech/pegasys/pantheon/consensus/ibft/SizeLimitedMapTest.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2019 ConsenSys AG.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package tech.pegasys.pantheon.consensus.ibft;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.Test;
+
+public class SizeLimitedMapTest {
+
+  @Test
+  public void evictMessageRecordAtCapacity() {
+    SizeLimitedMap<String, Boolean> map = new SizeLimitedMap<>(5);
+
+    map.put("message1", true);
+    assertThat(map).hasSize(1);
+
+    // add messages so map is at capacity
+    for (int i = 2; i <= 5; i++) {
+      map.put("message" + i, true);
+    }
+    assertThat(map).hasSize(5);
+
+    map.put("message6", false);
+    assertThat(map).hasSize(5);
+    assertThat(map.keySet()).doesNotContain("message1");
+    assertThat(map.keySet()).contains("message2", "message3", "message4", "message5", "message6");
+
+    map.put("message7", true);
+    assertThat(map).hasSize(5);
+    assertThat(map.keySet()).doesNotContain("message1", "message2");
+    assertThat(map.keySet()).contains("message3", "message4", "message5", "message6", "message7");
+  }
+}