Change the API of FlightMap so that we have separate methods for

truncating multiple entries and removing an individual entry

This will solve the problem where if we have gaps in the map
then some entries will never be removed.

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/ConsensusModule.java

@@ -73,7 +73,7 @@ public class ConsensusModule
     private final RaftMachine raftMachine;
     private final DelayedRenewableTimeoutService raftTimeoutService;
     private final RaftMembershipManager raftMembershipManager;
-    private final InFlightMap<Long,RaftLogEntry> inFlightMap = new InFlightMap<>();
+    private final InFlightMap<RaftLogEntry> inFlightMap = new InFlightMap<>();
 
     public ConsensusModule( MemberId myself, final PlatformModule platformModule, Outbound<MemberId,RaftMessages.RaftMessage> outbound,
             File clusterStateDirectory, CoreTopologyService discoveryService )
@@ -193,7 +193,7 @@ public RaftMembershipManager raftMembershipManager()
         return raftMembershipManager;
     }
 
-    public InFlightMap<Long,RaftLogEntry> inFlightMap()
+    public InFlightMap<RaftLogEntry> inFlightMap()
     {
         return inFlightMap;
     }

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/RaftMachine.java

@@ -97,7 +97,7 @@ public RaftMachine( MemberId myself, StateStorage<TermState> termStorage,
                         Outbound<MemberId, RaftMessages.RaftMessage> outbound,
                         LogProvider logProvider, RaftMembershipManager membershipManager,
                         RaftLogShippingManager logShipping,
-                        InFlightMap<Long, RaftLogEntry> inFlightMap,
+                        InFlightMap<RaftLogEntry> inFlightMap,
                         Monitors monitors )
     {
         this.myself = myself;

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/log/segmented/InFlightMap.java

@@ -19,14 +19,14 @@
  */
 package org.neo4j.coreedge.core.consensus.log.segmented;
 
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.SortedMap;
+import java.util.concurrent.ConcurrentSkipListMap;
 
 import static java.lang.String.format;
 
-public class InFlightMap<K, V>
+public class InFlightMap<V>
 {
-    private final Map<K,V> map = new ConcurrentHashMap<>();
+    private final SortedMap<Long,V> map = new ConcurrentSkipListMap<>();
 
     /**
      * Adds a new mapping.
@@ -35,7 +35,7 @@ public class InFlightMap<K, V>
      * @param value The value corresponding to the key provided.
      * @throws IllegalArgumentException if a mapping for the key already exists
      */
-    public void register( K key, V value )
+    public void put( Long key, V value )
     {
         V previousValue = map.putIfAbsent( key, value );
 
@@ -52,19 +52,35 @@ public void register( K key, V value )
      * @param key The key to use for retrieving the value from the map
      * @return the value for this key, otherwise null.
      */
-    public V retrieve( K key )
+    public V get( Long key )
     {
         return map.get( key );
     }
 
     /**
-     * Attempts to unregister this object from the map.
+     * Attempts to remove this object from the map.
      *
      * @param key The object to attempt unregistering.
      * @return true if the attempt to unregister was successful, otherwise false if this object was not found.
      */
-    public boolean unregister( K key )
+    public boolean remove( Long key )
     {
         return map.remove( key ) != null;
     }
+
+    /**
+     * Attempts to remove all objects at this key or higher from the map.
+     *
+     * @param key The object to attempt unregistering.
+     */
+    public void truncate( Long key )
+    {
+        map.tailMap( key ).keySet().forEach( map::remove );
+    }
+
+    @Override
+    public String toString()
+    {
+        return String.format( "InFlightMap{map=%s}", map );
+    }
 }

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/outcome/AppendLogEntry.java

@@ -49,9 +49,9 @@ public void applyTo( RaftLog raftLog, Log log ) throws IOException
     }
 
     @Override
-    public void applyTo( InFlightMap<Long, RaftLogEntry> inFlightMap, Log log ) throws IOException
+    public void applyTo( InFlightMap<RaftLogEntry> inFlightMap, Log log ) throws IOException
     {
-        inFlightMap.register( index, entry );
+        inFlightMap.put( index, entry );
     }
 
     @Override

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/outcome/BatchAppendLogEntries.java

@@ -62,11 +62,11 @@ public void applyTo( RaftLog raftLog, Log log ) throws IOException
     }
 
     @Override
-    public void applyTo( InFlightMap<Long, RaftLogEntry> inFlightMap, Log log )
+    public void applyTo( InFlightMap<RaftLogEntry> inFlightMap, Log log )
     {
         for ( int i = offset; i < entries.length; i++ )
         {
-            inFlightMap.register( baseIndex + i , entries[i]);
+            inFlightMap.put( baseIndex + i , entries[i]);
         }
     }
 

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/outcome/RaftLogCommand.java

@@ -38,5 +38,5 @@ interface Handler
 
     void applyTo( RaftLog raftLog, Log log ) throws IOException;
 
-    void applyTo( InFlightMap<Long,RaftLogEntry> inFlightMap, Log log ) throws IOException;
+    void applyTo( InFlightMap<RaftLogEntry> inFlightMap, Log log ) throws IOException;
 }

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/outcome/TruncateLogCommand.java

@@ -49,15 +49,10 @@ public void applyTo( RaftLog raftLog, Log log ) throws IOException
     }
 
     @Override
-    public void applyTo( InFlightMap<Long,RaftLogEntry> inFlightMap, Log log ) throws IOException
+    public void applyTo( InFlightMap<RaftLogEntry> inFlightMap, Log log ) throws IOException
     {
-        long truncateIndex = fromIndex;
-        log.debug( "Start truncating in-flight-map from index %d. Current map:%n%s", truncateIndex, inFlightMap );
-        while ( inFlightMap.unregister( truncateIndex ) )
-        {
-            truncateIndex++;
-        }
-        log.debug( "End truncating in-flight-map at index %d", truncateIndex - 1 );
+        log.debug( "Start truncating in-flight-map from index %d. Current map:%n%s", fromIndex, inFlightMap );
+        inFlightMap.truncate( fromIndex );
     }
 
     @Override

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/shipping/RaftLogShipper.java

@@ -111,7 +111,7 @@ public enum Timeouts implements RenewableTimeoutService.TimeoutName
     private final long retryTimeMillis;
     private final int catchupBatchSize;
     private final int maxAllowedShippingLag;
-    private final InFlightMap<Long, RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
 
     private DelayedRenewableTimeoutService timeoutService;
     private RenewableTimeout timeout;
@@ -124,7 +124,7 @@ public enum Timeouts implements RenewableTimeoutService.TimeoutName
     RaftLogShipper( Outbound<MemberId, RaftMessages.RaftMessage> outbound, LogProvider logProvider,
                     ReadableRaftLog raftLog, Clock clock,
                     MemberId leader, MemberId follower, long leaderTerm, long leaderCommit, long retryTimeMillis,
-                    int catchupBatchSize, int maxAllowedShippingLag, InFlightMap<Long, RaftLogEntry> inFlightMap )
+                    int catchupBatchSize, int maxAllowedShippingLag, InFlightMap<RaftLogEntry> inFlightMap )
     {
         this.outbound = outbound;
         this.catchupBatchSize = catchupBatchSize;

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/shipping/RaftLogShippingManager.java

@@ -51,7 +51,7 @@ public class RaftLogShippingManager extends LifecycleAdapter implements RaftMemb
     private final long retryTimeMillis;
     private final int catchupBatchSize;
     private final int maxAllowedShippingLag;
-    private final InFlightMap<Long,RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
 
     private Map<MemberId,RaftLogShipper> logShippers = new HashMap<>();
     private LeaderContext lastLeaderContext;
@@ -63,7 +63,7 @@ public RaftLogShippingManager( Outbound<MemberId,RaftMessages.RaftMessage> outbo
                                    ReadableRaftLog raftLog,
                                    Clock clock, MemberId myself, RaftMembership membership, long retryTimeMillis,
                                    int catchupBatchSize, int maxAllowedShippingLag,
-                                   InFlightMap<Long, RaftLogEntry> inFlightMap )
+                                   InFlightMap<RaftLogEntry> inFlightMap )
     {
         this.outbound = outbound;
         this.logProvider = logProvider;

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/consensus/state/RaftState.java

@@ -46,7 +46,7 @@ public class RaftState implements ReadableRaftState
     private final RaftMembership membership;
     private final Log log;
     private final RaftLog entryLog;
-    private final InFlightMap<Long,RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
 
     private TermState termState;
     private VoteState voteState;
@@ -63,7 +63,7 @@ public RaftState( MemberId myself,
                       RaftMembership membership,
                       RaftLog entryLog,
                       StateStorage<VoteState> voteStorage,
-                      InFlightMap<Long, RaftLogEntry> inFlightMap, LogProvider logProvider )
+                      InFlightMap<RaftLogEntry> inFlightMap, LogProvider logProvider )
     {
         this.myself = myself;
         this.termStorage = termStorage;

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/state/CommandApplicationProcess.java

@@ -56,7 +56,7 @@ public class CommandApplicationProcess extends LifecycleAdapter
     private final ProgressTracker progressTracker;
     private final SessionTracker sessionTracker;
     private final Supplier<DatabaseHealth> dbHealth;
-    private final InFlightMap<Long,RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
     private final Log log;
     private final CoreStateApplier applier;
     private final RaftLogCommitIndexMonitor commitIndexMonitor;
@@ -81,7 +81,7 @@ public CommandApplicationProcess(
             StateStorage<Long> lastFlushedStorage,
             SessionTracker sessionTracker,
             CoreStateApplier applier,
-            InFlightMap<Long, RaftLogEntry> inFlightMap,
+            InFlightMap<RaftLogEntry> inFlightMap,
             Monitors monitors )
     {
         this.coreStateMachines = coreStateMachines;

enterprise/core-edge/src/main/java/org/neo4j/coreedge/core/state/InFlightLogEntryReader.java

@@ -31,13 +31,13 @@
 public class InFlightLogEntryReader implements AutoCloseable
 {
     private final ReadableRaftLog raftLog;
-    private final InFlightMap<Long, RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
     private final boolean removeReadIndexFromMap;
 
     private RaftLogCursor cursor;
     private boolean useInFlightMap = true;
 
-    public InFlightLogEntryReader( ReadableRaftLog raftLog, InFlightMap<Long,RaftLogEntry> inFlightMap,
+    public InFlightLogEntryReader( ReadableRaftLog raftLog, InFlightMap<RaftLogEntry> inFlightMap,
             boolean removeReadIndexFromMap )
     {
         this.raftLog = raftLog;
@@ -51,7 +51,7 @@ public RaftLogEntry get( long logIndex ) throws IOException
 
         if ( useInFlightMap )
         {
-            entry = inFlightMap.retrieve( logIndex );
+            entry = inFlightMap.get( logIndex );
         }
 
         if ( entry == null )
@@ -62,7 +62,7 @@ public RaftLogEntry get( long logIndex ) throws IOException
 
         if ( removeReadIndexFromMap )
         {
-            inFlightMap.unregister( logIndex );
+            inFlightMap.remove( logIndex );
         }
 
         return entry;

enterprise/core-edge/src/test/java/org/neo4j/coreedge/core/consensus/RaftMachineBuilder.java

@@ -77,7 +77,7 @@ public class RaftMachineBuilder
             new InMemoryStateStorage<>( new RaftMembershipState() );
     private Monitors monitors = new Monitors();
     private CommitListener commitListener = commitIndex -> {};
-    private final InFlightMap<Long, RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
 
     public RaftMachineBuilder( MemberId member, int expectedClusterSize, RaftGroup.Builder memberSetBuilder )
     {

enterprise/core-edge/src/test/java/org/neo4j/coreedge/core/consensus/explorer/ComparableRaftState.java

@@ -53,11 +53,11 @@ public class ComparableRaftState implements ReadableRaftState
     private long lastLogIndexBeforeWeBecameLeader = -1;
     private FollowerStates followerStates = new FollowerStates<>();
     protected final RaftLog entryLog;
-    private final InFlightMap<Long,RaftLogEntry> inFlightMap;
+    private final InFlightMap<RaftLogEntry> inFlightMap;
     private long commitIndex = -1;
 
     ComparableRaftState( MemberId myself, Set votingMembers, Set replicationMembers,
-                         RaftLog entryLog, InFlightMap<Long, RaftLogEntry> inFlightMap, LogProvider logProvider )
+                         RaftLog entryLog, InFlightMap<RaftLogEntry> inFlightMap, LogProvider logProvider )
     {
         this.myself = myself;
         this.votingMembers = votingMembers;

enterprise/core-edge/src/test/java/org/neo4j/coreedge/core/consensus/log/segmented/InFlightLogEntriesCacheTest.java

@@ -21,8 +21,8 @@
 
 import org.junit.Test;
 
-import java.util.LinkedList;
-import java.util.List;
+import java.util.HashMap;
+import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -34,61 +34,58 @@ public class InFlightLogEntriesCacheTest
     @Test
     public void shouldRegisterAndUnregisterValues() throws Exception
     {
-        InFlightMap<Object,Object> entries = new InFlightMap<>();
+        InFlightMap<Object> entries = new InFlightMap<>();
 
-        List<Object> logEntryList = new LinkedList<>();
-        logEntryList.add( new Object() );
+        Map<Long, Object> logEntryList = new HashMap<>();
+        logEntryList.put(1L, new Object() );
 
-        for ( Object registeredEntry : logEntryList )
+        for ( Map.Entry<Long, Object> entry : logEntryList.entrySet() )
         {
-            entries.register( registeredEntry, registeredEntry );
+            entries.put( entry.getKey(), entry.getValue() );
         }
 
-        for ( Object expected : logEntryList )
+        for ( Map.Entry<Long, Object> entry : logEntryList.entrySet() )
         {
-            Object retrieved = entries.retrieve( expected );
-            assertEquals( expected, retrieved );
+            Object retrieved = entries.get( entry.getKey() );
+            assertEquals( entry.getValue(), retrieved );
         }
 
-        Object unexpected = new Object();
-        Object shouldBeNull = entries.retrieve( unexpected );
+        Long unexpected = 2L;
+        Object shouldBeNull = entries.get( unexpected );
         assertNull( shouldBeNull );
 
-        for ( Object expected : logEntryList )
+        for ( Map.Entry<Long, Object> entry : logEntryList.entrySet() )
         {
-            boolean wasThere = entries.unregister( expected );
+            boolean wasThere = entries.remove( entry.getKey() );
             assertEquals( true, wasThere );
         }
-
-        boolean shouldBeFalse = entries.unregister( unexpected );
-        assertFalse( shouldBeFalse );
     }
 
     @Test( expected = IllegalArgumentException.class )
     public void shouldNotReinsertValues() throws Exception
     {
-        InFlightMap<Object,Object> entries = new InFlightMap<>();
+        InFlightMap<Object> entries = new InFlightMap<>();
         Object addedObject = new Object();
-        entries.register( addedObject, addedObject );
-        entries.register( addedObject, addedObject );
+        entries.put( 1L, addedObject );
+        entries.put( 1L, addedObject );
     }
 
     @Test
     public void shouldNotReplaceRegisteredValues() throws Exception
     {
-        InFlightMap<Object,Object> cache = new InFlightMap<>();
+        InFlightMap<Object> cache = new InFlightMap<>();
         Object first = new Object();
         Object second = new Object();
 
         try
         {
-            cache.register( first, first );
-            cache.register( first, second );
+            cache.put( 1L, first );
+            cache.put( 1L, second );
             fail("Should not allow silent replacement of values");
         }
         catch ( IllegalArgumentException e )
         {
-            assertEquals( first, cache.retrieve( first ) );
+            assertEquals( first, cache.get( 1L ) );
         }
     }
 }

enterprise/core-edge/src/test/java/org/neo4j/coreedge/core/consensus/outcome/BatchAppendLogEntriesTest.java

@@ -90,14 +90,14 @@ public void applyTo() throws Exception
 
         BatchAppendLogEntries batchAppend = new BatchAppendLogEntries( baseIndex, offset, entries );
 
-        InFlightMap<Long,RaftLogEntry> cache = new InFlightMap<>();
+        InFlightMap<RaftLogEntry> cache = new InFlightMap<>();
 
         //when
         batchAppend.applyTo( cache, log );
 
         //then
-        assertNull( cache.retrieve( 0L ) );
-        assertNotNull( cache.retrieve( 1L ) );
-        assertNotNull( cache.retrieve( 2L ) );
+        assertNull( cache.get( 0L ) );
+        assertNotNull( cache.get( 1L ) );
+        assertNotNull( cache.get( 2L ) );
     }
 }