diff --git a/cnc/src/main/java/org/btcollider/cnc/keysrv/KeyServer.java b/cnc/src/main/java/org/btcollider/cnc/keysrv/KeyServer.java index ad79762..fd02fdf 100644 --- a/cnc/src/main/java/org/btcollider/cnc/keysrv/KeyServer.java +++ b/cnc/src/main/java/org/btcollider/cnc/keysrv/KeyServer.java @@ -1,6 +1,8 @@ package org.btcollider.cnc.keysrv; -import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import org.btcollider.cnc.dto.KeyRange; public interface KeyServer { @@ -34,9 +36,20 @@ public interface KeyServer { * Dumps the state of the keys managed by this KeyServer to a file. A KeyServer must be able to * read back the status of its keys from its dump file, but not necessarily from other KeyServers. * - * @param file The file to dump the key state to. Path must exist and must be write-able. Not - * null. + * @param out The output stream to write the dump to. Not null. */ - void dump(File file); + void dump(OutputStream out) throws IOException; + + /** + * Recreates the keys of a KeyServer from a dump. A KeyServer must be able to read back the status + * of its keys form a dump file, but not necessarily from other KeyServers. + * + * Use this method only directly after initialization on a fresh KeyServer. Otherwise the state + * might be unexpected. + * + * @param int The input stream to read the dump from. Not null. + * @throws IOException The undumping failed. Keys are in an undefined state. + */ + void undump(InputStream in) throws IOException; } diff --git a/cnc/src/main/java/org/btcollider/cnc/keysrv/keytree/KeyTreeServer.java b/cnc/src/main/java/org/btcollider/cnc/keysrv/keytree/KeyTreeServer.java index 5c6ce1d..10b6961 100644 --- a/cnc/src/main/java/org/btcollider/cnc/keysrv/keytree/KeyTreeServer.java +++ b/cnc/src/main/java/org/btcollider/cnc/keysrv/keytree/KeyTreeServer.java @@ -1,6 +1,8 @@ package org.btcollider.cnc.keysrv.keytree; -import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.util.BitSet; import java.util.LinkedList; import java.util.Queue; @@ -307,7 +309,48 @@ public class KeyTreeServer implements KeyServer { } @Override - public void dump(File file) { + public void dump(OutputStream out) throws IOException { + // Dump the KeyTree depth first + synchronized (this) { + recDump(root, out); + } + } + private void recDump(KeyTree node, OutputStream out) throws IOException { + if (node == null) + return; + + int dumpByte = (node.getValue() ? 0b01 : 0b00) | (node.isSearched() ? 0b10 : 0b00); + + out.write(dumpByte); + + recDump(node.getLeft(), out); + recDump(node.getRight(), out); + } + + @Override + public void undump(InputStream in) throws IOException { + // Read dump depth first + synchronized (this) { + recUndump(root, in); + } + } + + private void recUndump(KeyTree node, InputStream in) throws IOException { + if (node == null) { + assert (in.read() == -1); + return; + } + + int dumpByte = in.read(); + int value = dumpByte & 0b01; + int searched = dumpByte & 0b10; + + assert ((value == 0 && !node.getValue()) || (value == 1 && node.getValue())); + + node.setSearched(searched == 0 ? false : true); + + recUndump(node.getLeft(), in); + recUndump(node.getRight(), in); } } diff --git a/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/KeyRangeBuffer.java b/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/KeyRangeBuffer.java index c7fccee..44f5914 100644 --- a/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/KeyRangeBuffer.java +++ b/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/KeyRangeBuffer.java @@ -2,6 +2,7 @@ package org.btcollider.cnc.keysrv.multistage; import java.util.HashMap; import java.util.Map.Entry; +import java.util.Set; import org.btcollider.cnc.dto.KeyRange; /** @@ -45,4 +46,8 @@ public class KeyRangeBuffer { return retracedRange; } + + protected Set getAllRanges() { + return keyRanges.keySet(); + } } diff --git a/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/MultiStagedServer.java b/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/MultiStagedServer.java index 14f7140..4271e9c 100644 --- a/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/MultiStagedServer.java +++ b/cnc/src/main/java/org/btcollider/cnc/keysrv/multistage/MultiStagedServer.java @@ -1,6 +1,15 @@ package org.btcollider.cnc.keysrv.multistage; -import java.io.File; +import java.io.BufferedReader; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.io.OutputStreamWriter; +import java.io.PrintWriter; +import java.nio.charset.StandardCharsets; import org.btcollider.cnc.dto.KeyRange; import org.btcollider.cnc.keysrv.KeyServer; import org.btcollider.cnc.keysrv.keytree.KeyTreeServer; @@ -16,7 +25,7 @@ import org.slf4j.LoggerFactory; */ public class MultiStagedServer implements KeyServer { final Logger log = LoggerFactory.getLogger(MultiStagedServer.class); - private static final int PARTITION_DEPTH = 15; + private static final int PARTITION_DEPTH = 16; private KeyTreeServer primKTS; private KeyTreeServer secKTS; @@ -108,8 +117,63 @@ public class MultiStagedServer implements KeyServer { } @Override - public void dump(File file) { + public void dump(OutputStream out) throws IOException { + synchronized (this) { + ByteArrayOutputStream primDump = new ByteArrayOutputStream(); + primKTS.dump(primDump); + primDump.flush(); + ByteArrayOutputStream secDump = new ByteArrayOutputStream(); + secKTS.dump(secDump); + secDump.flush(); + + out.write(primDump.toByteArray()); + out.write(0xFF); + out.write(secDump.toByteArray()); + out.write(0xFF); + out.flush(); + + PrintWriter pw = new PrintWriter(new OutputStreamWriter(out, StandardCharsets.US_ASCII)); + pw.println(); + pw.println(partition.getStart() + "," + partition.getEnd()); + + keyRangeBuffer.getAllRanges().forEach(kr -> pw.println(kr.getStart() + "," + kr.getEnd())); + pw.flush(); + + out.flush(); + } + } + + @Override + public void undump(InputStream in) throws IOException { + synchronized (this) { + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + for (int b = in.read(); b != 0xFF; b = in.read()) { + bos.write(b); + } + ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray()); + primKTS.undump(bis); + + bos = new ByteArrayOutputStream(); + for (int b = in.read(); b != 0xFF; b = in.read()) { + bos.write(b); + } + bis = new ByteArrayInputStream(bos.toByteArray()); + secKTS.undump(bis); + + BufferedReader br = new BufferedReader(new InputStreamReader(in, StandardCharsets.US_ASCII)); + br.readLine(); + String[] partitionDump = br.readLine().split(","); + this.partition = + new KeyRange(Long.parseLong(partitionDump[0]), Long.parseLong(partitionDump[1])); + + String krbuf; + while ((krbuf = br.readLine()) != null) { + String[] krbufsplit = krbuf.split(","); + keyRangeBuffer + .setInWork(new KeyRange(Long.parseLong(krbufsplit[0]), Long.parseLong(krbufsplit[1]))); + } + } } }