SnapshotTracker implements AutoCloseable { boolean addChunk(final int chunkIndex, final byte[] chunk, final OptionalInt maybeLastChunkHashCode) throws InvalidChunkException, IOException { log.debug("addChunk: chunkIndex={}, lastChunkIndex={}, collectedChunks.size={}, lastChunkHashCode={}", chunkIndex, lastChunkIndex, count, this.lastChunkHashCode); if (sealed) { throw new InvalidChunkException("Invalid chunk received with chunkIndex " + chunkIndex + " all chunks already received"); } if (lastChunkIndex + 1 != chunkIndex) { throw new InvalidChunkException("Expected chunkIndex " + (lastChunkIndex + 1) + " got " + chunkIndex); } if (maybeLastChunkHashCode.isPresent() && maybeLastChunkHashCode.getAsInt() != this.lastChunkHashCode) { throw new InvalidChunkException("The hash code of the recorded last chunk does not match " + "the senders hash code, expected " + this.lastChunkHashCode + " was " + maybeLastChunkHashCode.getAsInt()); } bufferedStream.write(chunk); count += chunk.length; sealed = chunkIndex == totalChunks; lastChunkIndex = chunkIndex; this.lastChunkHashCode = Arrays.hashCode(chunk); return sealed; } SnapshotTracker(final Logger log, final int totalChunks, final String leaderId, final RaftActorContext context); @Override void close();  }
@Test(expected = SnapshotTracker.InvalidChunkException.class) public void testAddChunkWhenAlreadySealed() throws IOException { try (SnapshotTracker tracker = new SnapshotTracker(LOG, 2, "leader", mockContext)) { tracker.addChunk(1, chunk1, OptionalInt.empty()); tracker.addChunk(2, chunk2, OptionalInt.empty()); tracker.addChunk(3, chunk3, OptionalInt.empty()); } }
@Test(expected = SnapshotTracker.InvalidChunkException.class) public void testInvalidFirstChunkIndex() throws IOException { try (SnapshotTracker tracker = new SnapshotTracker(LOG, 2, "leader", mockContext)) { tracker.addChunk(LeaderInstallSnapshotState.FIRST_CHUNK_INDEX - 1, chunk1, OptionalInt.empty()); } }
@Test(expected = SnapshotTracker.InvalidChunkException.class) public void testOutOfSequenceChunk() throws IOException { try (SnapshotTracker tracker = new SnapshotTracker(LOG, 2, "leader", mockContext)) { tracker.addChunk(1, chunk1, OptionalInt.empty()); tracker.addChunk(3, chunk3, OptionalInt.empty()); } }
@Test(expected = SnapshotTracker.InvalidChunkException.class) public void testInvalidLastChunkHashCode() throws IOException { try (SnapshotTracker tracker = new SnapshotTracker(LOG, 2, "leader", mockContext)) { tracker.addChunk(1, chunk1, OptionalInt.of(LeaderInstallSnapshotState.INITIAL_LAST_CHUNK_HASH_CODE)); tracker.addChunk(2, chunk2, OptionalInt.of(1)); } }