diff --git a/modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java b/modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java new file mode 100644 index 0000000000000..e9b2a81f65815 --- /dev/null +++ b/modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.ignite.internal.benchmarks.jmh.direct; + +import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import org.apache.ignite.internal.CoreMessagesProvider; +import org.apache.ignite.internal.benchmarks.jmh.runner.JmhIdeBenchmarkRunner; +import org.apache.ignite.internal.direct.DirectMessageReader; +import org.apache.ignite.internal.direct.DirectMessageWriter; +import org.apache.ignite.internal.managers.communication.IgniteMessageFactoryImpl; +import org.apache.ignite.internal.processors.affinity.AffinityTopologyVersion; +import org.apache.ignite.internal.processors.cache.distributed.dht.preloader.GridDhtPartitionsFullMessage; +import org.apache.ignite.internal.processors.cache.distributed.dht.preloader.GroupPartitionIdPair; +import org.apache.ignite.internal.util.typedef.internal.U; +import org.apache.ignite.plugin.extensions.communication.Message; +import org.apache.ignite.plugin.extensions.communication.MessageFactory; +import org.apache.ignite.plugin.extensions.communication.MessageFactoryProvider; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; + +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.apache.ignite.marshaller.Marshallers.jdk; +import static org.openjdk.jmh.annotations.Mode.Throughput; + +/** Benchmarks the {@link DirectMessageReader} compressed-field hot path. */ +@State(Scope.Thread) +@BenchmarkMode(Throughput) +@Warmup(iterations = 5, time = 3, timeUnit = SECONDS) +@Measurement(iterations = 5, time = 3, timeUnit = SECONDS) +@Fork(1) +public class JmhDirectMessageReaderBenchmark { + /** */ + @Param({"30", "500"}) + private int entries; + + /** */ + private DirectMessageReader reader; + + /** Fully serialized compressed message, as received from the network. */ + private ByteBuffer buf; + + /** */ + public static void main(String[] args) throws Exception { + JmhIdeBenchmarkRunner.create() + .benchmarks(JmhDirectMessageReaderBenchmark.class.getName()) + .run(); + } + + /** */ + @Setup + public void setup() { + MessageFactory msgFactory = msgFactory(); + + reader = new DirectMessageReader(msgFactory, null); + + Map> partHistSuppliers = new HashMap<>(); + Map>> partsToReload = new HashMap<>(); + + for (int i = 0; i < entries; i++) { + UUID nodeId = UUID.randomUUID(); + + partHistSuppliers.put(nodeId, Map.of(new GroupPartitionIdPair(i, i + 1), i + 2L)); + partsToReload.put(nodeId, Map.of(i, Set.of(i + 1))); + } + + GridDhtPartitionsFullMessage msg = new GridDhtPartitionsFullMessage(null, null, + new AffinityTopologyVersion(0), partHistSuppliers, partsToReload); + + DirectMessageWriter writer = new DirectMessageWriter(msgFactory); + + buf = ByteBuffer.allocate(4 * 1024 * 1024); + + writer.setBuffer(buf); + + boolean finished = writer.writeMessage(msg, true); + + if (!finished) + throw new IllegalStateException("Message does not fit into the buffer."); + + buf.flip(); + } + + /** Exchange-style compressed message deserialization. */ + @Benchmark + public Message compressedMessage() { + buf.rewind(); + + reader.setBuffer(buf); + + Message msg = reader.readMessage(true); + + reader.reset(); + + return msg; + } + + /** */ + private static MessageFactory msgFactory() { + return new IgniteMessageFactoryImpl(new MessageFactoryProvider[]{ + new CoreMessagesProvider(jdk(), jdk(), U.gridClassLoader())}); + } +} diff --git a/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java b/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java index c5b2f3474fbc6..bde4e531b3407 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java +++ b/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java @@ -23,6 +23,7 @@ import java.util.Map; import java.util.UUID; import java.util.function.Function; +import org.apache.ignite.IgniteException; import org.apache.ignite.internal.direct.state.DirectMessageState; import org.apache.ignite.internal.direct.state.DirectMessageStateItem; import org.apache.ignite.internal.direct.stream.DirectByteBufferStream; @@ -49,6 +50,9 @@ * Message reader implementation. */ public class DirectMessageReader implements MessageReader { + /** Empty buffer to release the payload reference from {@link #tmpReader} streams between compressed reads. */ + private static final ByteBuffer EMPTY_BUF = ByteBuffer.wrap(new byte[0]); + /** State. */ @GridToStringInclude private final DirectMessageState state; @@ -62,6 +66,15 @@ public class DirectMessageReader implements MessageReader { /** Buffer for reading. */ private ByteBuffer buf; + /** + * Cached stream of the current state item. Avoids re-resolving {@code state.item().stream} on every primitive + * read; updated only when the current item changes (buffer set / nested read enter or exit). + */ + private DirectByteBufferStream curStream; + + /** Reusable reader for compressed payload deserialization. Lazily created; thread-confined like the reader itself. */ + private DirectMessageReader tmpReader; + /** Whether last field was fully read. */ private boolean lastRead; @@ -84,7 +97,9 @@ public DirectMessageReader(final MessageFactory msgFactory, IgniteCacheObjectPro @Override public void setBuffer(ByteBuffer buf) { this.buf = buf; - state.item().stream.setBuffer(buf); + curStream = state.item().stream; + + curStream.setBuffer(buf); } /** @@ -98,7 +113,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public byte readByte() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; byte val = stream.readByte(); @@ -109,7 +124,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public short readShort() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; short val = stream.readShort(); @@ -120,7 +135,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public int readInt() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; int val = stream.readInt(); @@ -131,7 +146,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public long readLong() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; long val = stream.readLong(); @@ -142,7 +157,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public float readFloat() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; float val = stream.readFloat(); @@ -153,7 +168,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public double readDouble() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; double val = stream.readDouble(); @@ -164,7 +179,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public char readChar() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; char val = stream.readChar(); @@ -175,7 +190,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public boolean readBoolean() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; boolean val = stream.readBoolean(); @@ -186,7 +201,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public byte[] readByteArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; byte[] arr = stream.readByteArray(); @@ -197,7 +212,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public short[] readShortArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; short[] arr = stream.readShortArray(); @@ -208,7 +223,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public int[] readIntArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; int[] arr = stream.readIntArray(); @@ -219,7 +234,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public long[] readLongArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; long[] arr = stream.readLongArray(); @@ -230,7 +245,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public float[] readFloatArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; float[] arr = stream.readFloatArray(); @@ -241,7 +256,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public double[] readDoubleArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; double[] arr = stream.readDoubleArray(); @@ -252,7 +267,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public char[] readCharArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; char[] arr = stream.readCharArray(); @@ -263,7 +278,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public boolean[] readBooleanArray() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; boolean[] arr = stream.readBooleanArray(); @@ -274,7 +289,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public String readString() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; String val = stream.readString(); @@ -285,7 +300,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public BitSet readBitSet() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; BitSet val = stream.readBitSet(); @@ -296,7 +311,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public UUID readUuid() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; UUID val = stream.readUuid(); @@ -307,7 +322,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public IgniteUuid readIgniteUuid() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; IgniteUuid val = stream.readIgniteUuid(); @@ -318,7 +333,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public AffinityTopologyVersion readAffinityTopologyVersion() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; AffinityTopologyVersion val = stream.readAffinityTopologyVersion(); @@ -329,14 +344,14 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Nullable @Override public T readMessage(boolean compress) { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; T msg; if (compress) msg = readCompressedMessageAndDeserialize( stream, - tmpReader -> tmpReader.state.item().stream.readMessage(tmpReader) + r -> r.state.item().stream.readMessage(r) ); else { msg = stream.readMessage(this); @@ -349,7 +364,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public CacheObject readCacheObject() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; CacheObject val = stream.readCacheObject(); @@ -360,7 +375,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public KeyCacheObject readKeyCacheObject() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; KeyCacheObject key = stream.readKeyCacheObject(); @@ -371,7 +386,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public GridLongList readGridLongList() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; GridLongList ll = stream.readGridLongList(); @@ -382,7 +397,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public T[] readObjectArray(MessageArrayType type) { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; T[] msg = stream.readObjectArray(type, this); @@ -393,7 +408,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public > C readCollection(MessageCollectionType type) { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; C col = stream.readCollection(type, this); @@ -404,14 +419,14 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public > M readMap(MessageMapType type, boolean compress) { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; M map; if (compress) map = readCompressedMessageAndDeserialize( stream, - tmpReader -> tmpReader.state.item().stream.readMap(type, tmpReader) + r -> r.state.item().stream.readMap(type, r) ); else { map = stream.readMap(type, this); @@ -424,7 +439,7 @@ public ByteBuffer getBuffer() { /** {@inheritDoc} */ @Override public IgniteProductVersion readIgniteProductVersion() { - DirectByteBufferStream stream = state.item().stream; + DirectByteBufferStream stream = curStream; IgniteProductVersion v = stream.readIgniteProductVersion(); @@ -457,12 +472,16 @@ public ByteBuffer getBuffer() { @Override public void beforeNestedRead() { state.forward(); - state.item().stream.setBuffer(buf); + curStream = state.item().stream; + + curStream.setBuffer(buf); } /** {@inheritDoc} */ @Override public void afterNestedRead(boolean finished) { state.backward(finished); + + curStream = state.item().stream; } /** {@inheritDoc} */ @@ -491,24 +510,49 @@ private T readCompressedMessageAndDeserialize(DirectByteBufferStream stream, if (msg0.dataSize() == 0) return null; - byte[] uncompressed = msg0.uncompressed(); + // Reuse the temp reader across fields/messages instead of allocating a fresh state stack each time. + if (tmpReader == null) + tmpReader = new DirectMessageReader(msgFactory, cacheObjProc); + else + tmpReader.reset(); - ByteBuffer tmpBuf = ByteBuffer.allocateDirect(uncompressed.length); + tmpReader.setBuffer(ByteBuffer.wrap(msg0.uncompressed())); - tmpBuf.put(uncompressed); - tmpBuf.flip(); + T res; - DirectMessageReader tmpReader = new DirectMessageReader(msgFactory, cacheObjProc); + boolean ok = false; - tmpReader.setBuffer(tmpBuf); + try { + res = fun.apply(tmpReader); - T res = fun.apply(tmpReader); + // The payload buffer is always complete, so a partial read means a corrupted stream, not a lack of data. + if (!tmpReader.state.item().stream.lastFinished()) + throw new IgniteException("Failed to deserialize compressed payload: uncompressed data ended " + + "unexpectedly [dataSize=" + msg0.dataSize() + ']'); - lastRead = tmpReader.state.item().stream.lastFinished(); + ok = true; + } + finally { + // Don't reuse a reader with half-read state. + if (!ok) + tmpReader = null; + } + + // Don't pin the payload: the reader lives as long as the connection. + tmpReader.releasePayload(); + + lastRead = true; return res; } + /** Releases the payload buffer reference from all streams of this reader. */ + private void releasePayload() { + buf = EMPTY_BUF; + + state.forEachItem(item -> item.stream.setBuffer(EMPTY_BUF)); + } + /** */ private static class StateItem implements DirectMessageStateItem { diff --git a/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java b/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java index b92f6000fdd7a..67bd768342552 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java +++ b/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java @@ -19,6 +19,7 @@ import java.lang.reflect.Array; import java.util.Arrays; +import java.util.function.Consumer; import org.apache.ignite.internal.util.tostring.GridToStringExclude; import org.apache.ignite.internal.util.typedef.internal.S; import org.apache.ignite.lang.IgniteOutClosure; @@ -98,6 +99,19 @@ public void reset() { stack[0].reset(); } + /** + * @param c Closure to apply to every created item. + */ + public void forEachItem(Consumer c) { + // Items are created contiguously by forward(). + for (T item : stack) { + if (item == null) + break; + + c.accept(item); + } + } + /** {@inheritDoc} */ @Override public String toString() { return S.toString(DirectMessageState.class, this, "stack", Arrays.toString(stack)); diff --git a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java index e3e46c43ff35c..ca39336175af1 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java +++ b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java @@ -17,14 +17,13 @@ package org.apache.ignite.internal.managers.communication; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.IOException; import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.zip.DataFormatException; import java.util.zip.Deflater; -import java.util.zip.DeflaterOutputStream; import java.util.zip.Inflater; -import java.util.zip.InflaterInputStream; import org.apache.ignite.IgniteException; import org.apache.ignite.internal.util.typedef.internal.S; import org.apache.ignite.plugin.extensions.communication.Message; @@ -38,18 +37,15 @@ public class CompressedMessage implements Message { /** Chunk size. */ static final int CHUNK_SIZE = 10 * 1024; - /** Reader buffer capacity. */ - static final int BUFFER_CAPACITY = 10 * CHUNK_SIZE; + /** Compressed data chunks: filled by {@link #compress(ByteBuffer)} on send, by the serializer on receive. */ + List chunks; - /** Temporary buffer for compressed data received over the network. */ - ByteBuffer tmpBuf; + /** Index of the next chunk to send. */ + private int chunkIdx; /** Raw data size. */ int dataSize; - /** Chunked byte reader. */ - ChunkedByteReader chunkedReader; - /** Chunk. */ byte[] chunk; @@ -73,7 +69,7 @@ public CompressedMessage(ByteBuffer buf, int compressionLvl) { this.compressionLvl = compressionLvl; if (dataSize > 0) - chunkedReader = new ChunkedByteReader(compress(buf)); + compress(buf); } /** @return Raw data size. */ @@ -90,99 +86,106 @@ public byte[] uncompressed() { /** @return Next chunk of data or null. */ public byte[] nextChunk() { - return chunkedReader.nextChunk(); + return chunkIdx < chunks.size() ? chunks.get(chunkIdx++) : null; } - /** - * @param buf Buffer. - * @return Compressed data. - */ - private byte[] compress(ByteBuffer buf) { - byte[] data = new byte[dataSize]; + /** @param buf Buffer. */ + private void compress(ByteBuffer buf) { + Deflater deflater = new Deflater(compressionLvl, true); - buf.get(data); + try { + deflater.setInput(buf); + deflater.finish(); - ByteArrayOutputStream baos = new ByteArrayOutputStream(data.length); - Deflater deflater = new Deflater(compressionLvl, true); + chunks = new ArrayList<>(dataSize / CHUNK_SIZE + 1); - try (DeflaterOutputStream dos = new DeflaterOutputStream(baos, deflater)) { - dos.write(data); - dos.finish(); - } - catch (IOException ex) { - throw new IgniteException(ex); + // Incompressible data may expand: raw-deflate worst case below CHUNK_SIZE is one stored block (~11 bytes overhead). + byte[] chunk0 = new byte[dataSize >= CHUNK_SIZE ? CHUNK_SIZE : dataSize + 16]; + + int len = 0; + + while (!deflater.finished()) { + len += deflater.deflate(chunk0, len, chunk0.length - len); + + if (len == chunk0.length && !deflater.finished()) { + chunks.add(chunk0); + + chunk0 = new byte[CHUNK_SIZE]; + len = 0; + } + } + + if (len > 0) + chunks.add(len == chunk0.length ? chunk0 : Arrays.copyOf(chunk0, len)); } finally { deflater.end(); } - - return baos.toByteArray(); } /** @return Uncompressed data. */ private byte[] uncompress() { - if (tmpBuf == null) - return null; + if (chunks == null) + throw new IgniteException("Compressed stream is truncated [expected=" + dataSize + ", inflated=0]"); + + long compressedTotal = 0; - byte[] uncompressedData; + for (int i = 0; i < chunks.size(); i++) + compressedTotal += chunks.get(i).length; + + // Raw deflate cannot expand beyond ~1032:1, a larger claim means a corrupted size header; also bounds the + // upfront allocation by ~1032x of the actually received bytes. + if (dataSize > compressedTotal * 1032L + 64) + throw new IgniteException("Invalid compressed message data size [dataSize=" + dataSize + + ", compressedBytes=" + compressedTotal + ']'); + + byte[] data = new byte[dataSize]; Inflater inflater = new Inflater(true); - byte[] bytes = new byte[tmpBuf.position()]; + try { + int off = 0; + int i = 0; + + for (; i < chunks.size() && off < dataSize; i++) { + inflater.setInput(chunks.get(i)); - tmpBuf.flip(); - tmpBuf.get(bytes); + int n; - try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(bytes), inflater)) { - uncompressedData = iis.readAllBytes(); + while (off < dataSize && (n = inflater.inflate(data, off, dataSize - off)) > 0) + off += n; + } + + if (off != dataSize) + throw new IgniteException("Compressed stream is truncated [expected=" + dataSize + ", inflated=" + off + ']'); + + // Any extra inflatable byte means the size header is understated. + byte[] probe = new byte[1]; + + while (true) { + if (inflater.inflate(probe, 0, 1) > 0) + throw new IgniteException("Compressed stream is longer than expected [expected=" + dataSize + ']'); + + if (inflater.finished() || i == chunks.size()) + break; + + inflater.setInput(chunks.get(i++)); + } } - catch (IOException ex) { - throw new IgniteException(ex); + catch (DataFormatException e) { + throw new IgniteException(e); } finally { inflater.end(); } - assert uncompressedData != null; - assert uncompressedData.length == dataSize : "Expected=" + dataSize + ", actual=" + uncompressedData.length; + chunks = null; - tmpBuf = null; - - return uncompressedData; + return data; } /** {@inheritDoc} */ @Override public String toString() { return S.toString(CompressedMessage.class, this); } - - /** Byte reader returning data chunks of predefined size. */ - private static class ChunkedByteReader { - /** Input data. */ - private final byte[] inputData; - - /** Current position. */ - private int pos; - - /** Constructor. */ - ChunkedByteReader(byte[] inputData) { - this.inputData = inputData; - } - - /** @return Next chunk of bytes or null. */ - byte[] nextChunk() { - if (pos >= inputData.length) - return null; - - int curChunkSize = Math.min(inputData.length - pos, CHUNK_SIZE); - - byte[] chunk = new byte[curChunkSize]; - - System.arraycopy(inputData, pos, chunk, 0, curChunkSize); - - pos += curChunkSize; - - return chunk; - } - } } diff --git a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java index 594aab336e02d..c1a42085bc95e 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java +++ b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java @@ -17,12 +17,12 @@ package org.apache.ignite.internal.managers.communication; -import java.nio.ByteBuffer; +import java.util.ArrayList; +import org.apache.ignite.IgniteException; import org.apache.ignite.plugin.extensions.communication.MessageReader; import org.apache.ignite.plugin.extensions.communication.MessageSerializer; import org.apache.ignite.plugin.extensions.communication.MessageWriter; -import static org.apache.ignite.internal.managers.communication.CompressedMessage.BUFFER_CAPACITY; import static org.apache.ignite.internal.managers.communication.CompressedMessage.CHUNK_SIZE; /** Message serializer for compressed message. */ @@ -37,7 +37,7 @@ public class CompressedMessageSerializer implements MessageSerializer(Math.min(msg.dataSize / CHUNK_SIZE + 1, 1024)); - msg.tmpBuf = newTmpBuf; - } + msg.chunks.add(msg.chunk); - msg.tmpBuf.put(msg.chunk); + reader.decrementState(); - reader.decrementState(); - - msg.chunk = null; - } + msg.chunk = null; } } } diff --git a/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java b/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java index 53fa342460b41..98f1478c2b3a4 100644 --- a/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java +++ b/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java @@ -19,9 +19,13 @@ import java.nio.ByteBuffer; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.Random; import java.util.Set; import java.util.UUID; +import java.util.zip.Deflater; +import org.apache.ignite.IgniteException; import org.apache.ignite.internal.CoreMessagesProvider; import org.apache.ignite.internal.direct.DirectMessageReader; import org.apache.ignite.internal.direct.DirectMessageWriter; @@ -34,6 +38,7 @@ import org.apache.ignite.plugin.extensions.communication.Message; import org.apache.ignite.plugin.extensions.communication.MessageFactory; import org.apache.ignite.plugin.extensions.communication.MessageFactoryProvider; +import org.apache.ignite.testframework.GridTestUtils; import org.junit.Test; import static org.apache.ignite.marshaller.Marshallers.jdk; @@ -42,13 +47,14 @@ /** Test for {@link CompressedMessage}. */ public class CompressedMessageTest { + /** */ + private static final MessageFactory MSG_FACTORY = new IgniteMessageFactoryImpl(new MessageFactoryProvider[]{ + new CoreMessagesProvider(jdk(), jdk(), U.gridClassLoader())}); + /** */ @Test public void testWriteReadHugeMessage() { - MessageFactory msgFactory = new IgniteMessageFactoryImpl(new MessageFactoryProvider[]{ - new CoreMessagesProvider(jdk(), jdk(), U.gridClassLoader())}); - - DirectMessageWriter writer = new DirectMessageWriter(msgFactory); + DirectMessageWriter writer = new DirectMessageWriter(MSG_FACTORY); ByteBuffer tmpBuf = ByteBuffer.allocate(4096); @@ -74,9 +80,9 @@ public void testWriteReadHugeMessage() { assertTrue(compressedMsg.dataSize() > 0); - byte[] compressedData = U.field((Object)U.field(compressedMsg, "chunkedReader"), "inputData"); + List chunks = U.field(compressedMsg, "chunks"); - assertTrue(compressedData.length > CompressedMessage.CHUNK_SIZE * 2); + assertTrue(chunks.size() > 2); checkChunkCnt = false; } @@ -94,7 +100,7 @@ public void testWriteReadHugeMessage() { msgBuf.flip(); - DirectMessageReader reader = new DirectMessageReader(msgFactory, null); + DirectMessageReader reader = new DirectMessageReader(MSG_FACTORY, null); reader.setBuffer(msgBuf); @@ -105,6 +111,137 @@ public void testWriteReadHugeMessage() { assertEqualsFullMsg(fullMsg, (GridDhtPartitionsFullMessage)readMsg); } + /** Read must fail with an exception on a null chunk from the wire instead of looping forever. */ + @Test + public void testReadFailsOnNullChunk() { + DirectMessageWriter writer = new DirectMessageWriter(MSG_FACTORY); + + ByteBuffer buf = ByteBuffer.allocate(16); + + writer.setBuffer(buf); + + // Emulate a corrupted stream or an incompatible peer: dataSize > 0, non-final chunk, then the null-array + // marker (-1), which CompressedMessageSerializer.writeTo() never produces at the chunk position. + writer.writeInt(100); + writer.writeBoolean(false); + writer.writeByteArray(null); + + buf.flip(); + + DirectMessageReader reader = new DirectMessageReader(MSG_FACTORY, null); + + reader.setBuffer(buf); + + GridTestUtils.assertThrows(null, + () -> new CompressedMessageSerializer().readFrom(new CompressedMessage(), reader), + IgniteException.class, + "unexpected null chunk"); + } + + /** Read must fail fast on a negative data size from the wire. */ + @Test + public void testReadFailsOnNegativeDataSize() { + DirectMessageWriter writer = new DirectMessageWriter(MSG_FACTORY); + + ByteBuffer buf = ByteBuffer.allocate(16); + + writer.setBuffer(buf); + + writer.writeInt(-5); + + buf.flip(); + + DirectMessageReader reader = new DirectMessageReader(MSG_FACTORY, null); + + reader.setBuffer(buf); + + GridTestUtils.assertThrows(null, + () -> new CompressedMessageSerializer().readFrom(new CompressedMessage(), reader), + IgniteException.class, + "Invalid compressed message data size"); + } + + /** Uncompress must fail when dataSize > 0 but no chunks were received. */ + @Test + public void testUncompressFailsWithoutChunks() { + CompressedMessage rcvd = new CompressedMessage(); + + rcvd.dataSize = 100; + rcvd.finalChunk = true; + + GridTestUtils.assertThrows(null, rcvd::uncompressed, IgniteException.class, "truncated"); + } + + /** Uncompress must fail when the stream inflates to more bytes than the size header claims. */ + @Test + public void testUncompressFailsOnUnderstatedDataSize() { + byte[] data = new byte[1000]; + + CompressedMessage sent = new CompressedMessage(ByteBuffer.wrap(data), Deflater.BEST_SPEED); + + CompressedMessage rcvd = new CompressedMessage(); + + rcvd.dataSize = data.length - 1; + rcvd.chunks = sent.chunks; + rcvd.finalChunk = true; + + GridTestUtils.assertThrows(null, rcvd::uncompressed, IgniteException.class, "longer than expected"); + } + + /** Same as {@link #testUncompressFailsOnUnderstatedDataSize()}, but with a multi-chunk compressed stream. */ + @Test + public void testUncompressFailsOnUnderstatedDataSizeMultiChunk() { + byte[] data = new byte[CompressedMessage.CHUNK_SIZE * 3]; + + new Random(42).nextBytes(data); + + CompressedMessage sent = new CompressedMessage(ByteBuffer.wrap(data), Deflater.BEST_SPEED); + + assertTrue(sent.chunks.size() > 1); + + CompressedMessage rcvd = new CompressedMessage(); + + rcvd.dataSize = CompressedMessage.CHUNK_SIZE / 2; + rcvd.chunks = sent.chunks; + rcvd.finalChunk = true; + + GridTestUtils.assertThrows(null, rcvd::uncompressed, IgniteException.class, "longer than expected"); + } + + /** A complete envelope whose payload doesn't deserialize fully must fail instead of hanging as a partial read. */ + @Test + public void testReadFailsOnTruncatedPayload() { + DirectMessageWriter writer = new DirectMessageWriter(MSG_FACTORY); + + ByteBuffer tmpBuf = ByteBuffer.allocate(1 << 20); + + writer.setBuffer(tmpBuf); + + assertTrue(writer.writeMessage(fullMessage(), false)); + + tmpBuf.flip(); + + tmpBuf.limit(tmpBuf.limit() - 5); // Truncate the serialized message. + + CompressedMessage compressedMsg = new CompressedMessage(tmpBuf, Deflater.BEST_SPEED); + + DirectMessageWriter wireWriter = new DirectMessageWriter(MSG_FACTORY); + + ByteBuffer wire = ByteBuffer.allocate(1 << 20); + + wireWriter.setBuffer(wire); + + assertTrue(wireWriter.writeMessage(compressedMsg, false)); + + wire.flip(); + + DirectMessageReader reader = new DirectMessageReader(MSG_FACTORY, null); + + reader.setBuffer(wire); + + GridTestUtils.assertThrows(null, () -> reader.readMessage(true), IgniteException.class, "ended unexpectedly"); + } + /** */ private GridDhtPartitionsFullMessage fullMessage() { Map> partHistSuppliers = new HashMap<>();