Added KeyRangeBuffer to elevate missed keys when switching partitions

cnc/src/main/java/org/btcollider/cnc/dto/KeyRange.java

@@ -39,4 +39,35 @@ public class KeyRange {
 	public String toString() {
 		return "KeyRange [start=" + start + ", end=" + end + ", total=" + getTotal() +"]";
 	}
+
+	@Override
+	public int hashCode() {
+		final int prime = 31;
+		int result = 1;
+		result = prime * result + ((end == null) ? 0 : end.hashCode());
+		result = prime * result + ((start == null) ? 0 : start.hashCode());
+		return result;
+	}
+
+	@Override
+	public boolean equals(Object obj) {
+		if (this == obj)
+			return true;
+		if (obj == null)
+			return false;
+		if (getClass() != obj.getClass())
+			return false;
+		KeyRange other = (KeyRange) obj;
+		if (end == null) {
+			if (other.end != null)
+				return false;
+		} else if (!end.equals(other.end))
+			return false;
+		if (start == null) {
+			if (other.start != null)
+				return false;
+		} else if (!start.equals(other.start))
+			return false;
+		return true;
+	}
 }

cnc/src/main/java/org/btcollider/cnc/keysrv/impl/multistage/KeyRangeBuffer.java

@@ -0,0 +1,50 @@
+package org.btcollider.cnc.keysrv.impl.multistage;
+
+import java.util.HashMap;
+import java.util.Map.Entry;
+
+import org.btcollider.cnc.dto.KeyRange;
+
+/**
+ * The MultiStagedServer may switch partitions while there are still KeyRanges
+ * from the old partition that are currently worked on.
+ * 
+ * This buffer stores the state of unfinished KeyRanges from old partitions in
+ * order to handle them correctly.
+ * 
+ * Aborted KeyRanges from this Buffer should be always preferred to getting new
+ * ranges.
+ * 
+ * @author armin
+ */
+public class KeyRangeBuffer {
+	private HashMap<KeyRange, Long> keyRanges;
+	private Long maxWorkSpan;
+
+	public KeyRangeBuffer(Long maxWorkSpan) {
+		this.keyRanges = new HashMap<>();
+		this.maxWorkSpan = maxWorkSpan;
+	}
+	
+	public void setInWork(KeyRange keyRange) {
+		this.keyRanges.put(keyRange, System.currentTimeMillis());
+	}
+	
+	public void setSearched(KeyRange keyRange) {
+		keyRanges.computeIfPresent(keyRange, null);
+	}
+
+	public KeyRange getRange() {
+		long curTime = System.currentTimeMillis();
+		
+		KeyRange retracedRange = null;
+		for (Entry<KeyRange, Long> kr: keyRanges.entrySet()) {
+			if ( (curTime - kr.getValue()) >= maxWorkSpan ) {
+				retracedRange = kr.getKey();
+				break;
+			}
+		}
+	
+		return retracedRange;
+	}
+}

cnc/src/main/java/org/btcollider/cnc/keysrv/impl/multistage/MultiStagedServer.java

@@ -20,6 +20,7 @@ public class MultiStagedServer implements KeyServer {
 	private KeyTreeServer primKTS;
 	private KeyTreeServer secKTS;
 	private KeyRange partition;
+	private KeyRangeBuffer keyRangeBuffer;
 	private int remainingDepth;
 	private int subIndex;
 	private long maxWorkSpan;
@@ -39,13 +40,19 @@ public class MultiStagedServer implements KeyServer {
 		// enabled KTConcierge service.
 		// Needs subindex+1 because starting bit always set to 1
 		this.secKTS = new KeyTreeServer(subIndex+1, remainingDepth, maxWorkSpan);
+
+		this.keyRangeBuffer = new KeyRangeBuffer(maxWorkSpan);
 	}
 
 	@Override
 	public KeyRange getRange() {
 		if (partition == null) return null; // We are finished searching the whole space
 		
-		KeyRange kr = secKTS.getRange(); // Get a random range from the subspace
+		KeyRange kr = null;
+		
+		if ( (kr = keyRangeBuffer.getRange()) == null) { // Get retracted KR from buffer if one exists
+			secKTS.getRange(); // else: get a random range from the subspace
+		}
 		
 		if (kr == null) { // sub-range finished
 			log.info("Subrange {} finished. Opening new partition.", partition);
@@ -69,6 +76,8 @@ public class MultiStagedServer implements KeyServer {
 
 	@Override
 	public void setInWork(KeyRange keyRange) {
+		keyRangeBuffer.setInWork(keyRange);
+
 		// strip out partition bits -> get subspace KeyRange only
 		long mask = (long) Math.pow(2, (subIndex+1)) - 1;
 		// add leading bit
@@ -80,6 +89,8 @@ public class MultiStagedServer implements KeyServer {
 
 	@Override
 	public void setSearched(KeyRange keyRange) {
+		keyRangeBuffer.setSearched(keyRange);
+
 		// strip out partition bits -> get subspace KeyRange only
 		long mask = (long) Math.pow(2, (subIndex+1)) - 1;
 		// add leading bit