AsyncSnapshotWriterImpl implements AsyncSnapshotWriter { @Override @CheckReturnValue public boolean offer(Entry<? extends Data, ? extends Data> entry) { int partitionId = partitionService.getPartitionId(entry.getKey()); int length = entry.getKey().totalSize() + entry.getValue().totalSize() - 2 * HeapData.TYPE_OFFSET; if (length > usableChunkCapacity) { return putAsyncToMap(partitionId, () -> { byte[] data = new byte[serializedByteArrayHeader.length + length + valueTerminator.length]; totalKeys++; int offset = 0; System.arraycopy(serializedByteArrayHeader, 0, data, offset, serializedByteArrayHeader.length); offset += serializedByteArrayHeader.length - Bits.INT_SIZE_IN_BYTES; Bits.writeInt(data, offset, length + valueTerminator.length, useBigEndian); offset += Bits.INT_SIZE_IN_BYTES; copyWithoutHeader(entry.getKey(), data, offset); offset += entry.getKey().totalSize() - HeapData.TYPE_OFFSET; copyWithoutHeader(entry.getValue(), data, offset); offset += entry.getValue().totalSize() - HeapData.TYPE_OFFSET; System.arraycopy(valueTerminator, 0, data, offset, valueTerminator.length); return new HeapData(data); }); } CustomByteArrayOutputStream buffer = buffers[partitionId]; if (buffer.size() + length + valueTerminator.length > buffer.capacityLimit && !flushPartition(partitionId)) { return false; } writeWithoutHeader(entry.getKey(), buffer); writeWithoutHeader(entry.getValue(), buffer); totalKeys++; return true; } AsyncSnapshotWriterImpl(NodeEngine nodeEngine,
                                   SnapshotContext snapshotContext,
                                   String vertexName,
                                   int memberIndex,
                                   int memberCount,
                                   SerializationService serializationService);   AsyncSnapshotWriterImpl(int chunkSize,
                            NodeEngine nodeEngine,
                            SnapshotContext snapshotContext,
                            String vertexName,
                            int memberIndex,
                            int memberCount,
                            SerializationService serializationService); @Override @CheckReturnValue boolean offer(Entry<? extends Data, ? extends Data> entry); @Override @CheckReturnValue boolean flushAndResetMap(); @Override void resetStats(); @Override boolean hasPendingAsyncOps(); @Override Throwable getError(); @Override boolean isEmpty(); @Override long getTotalPayloadBytes(); @Override long getTotalKeys(); @Override long getTotalChunks(); static final int DEFAULT_CHUNK_SIZE; }
@Test public void when_cannotAutoFlush_then_offerReturnsFalse() { writer.numConcurrentAsyncOps.set(JetService.MAX_PARALLEL_ASYNC_OPS); Entry<Data, Data> entry = entry(serialize("k"), serialize("v")); int entriesInChunk = (writer.usableChunkCapacity - writer.serializedByteArrayHeader.length) / serializedLength(entry); assertTrue("entriesInChunk=" + entriesInChunk, entriesInChunk > 1 && entriesInChunk < 10); for (int i = 0; i < entriesInChunk; i++) { assertTrue(writer.offer(entry)); } assertFalse("offer should not have succeeded, too many parallel operations", writer.offer(entry)); writer.numConcurrentAsyncOps.set(0); assertTrue("offer should have succeeded", writer.offer(entry)); assertTargetMapEntry("k", 0, serializedLength(entry) * entriesInChunk); }