From 3f4b422c5573b8ac7c12a2f1a5740e6b5b857e53 Mon Sep 17 00:00:00 2001 From: Tian Jiang Date: Thu, 5 Nov 2026 18:09:06 +0800 Subject: [PATCH] Add configurable Thrift buffer reuse thresholds to Java clients --- .../apache/iotdb/rpc/AutoResizingBuffer.java | 7 + .../rpc/AutoScalingBufferReadTransport.java | 34 +- .../rpc/AutoScalingBufferWriteTransport.java | 51 ++- .../rpc/DeepCopyRpcTransportFactory.java | 18 +- .../TCompressedElasticFramedTransport.java | 30 +- .../iotdb/rpc/TElasticFramedTransport.java | 19 +- .../rpc/TSnappyElasticFramedTransport.java | 30 +- ...TimeoutChangeableTFastFramedTransport.java | 29 +- ...meoutChangeableTSnappyFramedTransport.java | 29 +- .../rpc/DeepCopyRpcTransportFactoryTest.java | 41 +++ ...TCompressedElasticFramedTransportTest.java | 295 ++++++++++++++++++ .../iotdb/session/AbstractSessionBuilder.java | 1 + .../apache/iotdb/session/NodesSupplier.java | 57 +++- .../org/apache/iotdb/session/Session.java | 16 +- .../iotdb/session/SessionConnection.java | 4 +- .../iotdb/session/TableSessionBuilder.java | 14 + .../iotdb/session/ThriftConnection.java | 19 +- .../iotdb/session/pool/SessionPool.java | 22 +- .../session/pool/TableSessionPoolBuilder.java | 15 + .../session/SessionSeedNodeFallbackTest.java | 137 ++++++++ 20 files changed, 832 insertions(+), 36 deletions(-) create mode 100644 iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransportTest.java diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java index d9772985dc614..f2a57cb5ca862 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java @@ -86,6 +86,13 @@ public byte[] array() { return this.array; } + void resetToInitialCapacity() throws IOException { + reserveCurrentCapacityIfReleased(); + resize(initialCapacity, initialCapacity); + bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; + lastShrinkTime = System.currentTimeMillis(); + } + @Override public void close() { if (!closed) { diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java index 035f8753de1ab..87a0e93b7d819 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java @@ -28,20 +28,28 @@ public class AutoScalingBufferReadTransport extends NonOpenTransport { private final AutoResizingBuffer buf; + private final int maxReusableBufferSize; + private byte[] temporaryBuffer; private int pos = 0; private int limit = 0; public AutoScalingBufferReadTransport(int initialCapacity) throws IOException { + this(initialCapacity, Integer.MAX_VALUE); + } + + AutoScalingBufferReadTransport(int initialCapacity, int maxReusableBufferSize) + throws IOException { this.buf = new AutoResizingBuffer(initialCapacity); + this.maxReusableBufferSize = maxReusableBufferSize; } public void fill(TTransport inTrans, int length) throws TTransportException { try { - buf.resizeIfNecessary(length); + resizeIfNecessary(length); } catch (IOException e) { throw new TTransportException(e); } - inTrans.readAll(buf.array(), 0, length); + inTrans.readAll(getBuffer(), 0, length); pos = 0; limit = length; } @@ -49,7 +57,7 @@ public void fill(TTransport inTrans, int length) throws TTransportException { @Override public final int read(byte[] target, int off, int len) { int amtToRead = Math.min(len, getBytesRemainingInBuffer()); - System.arraycopy(buf.array(), pos, target, off, amtToRead); + System.arraycopy(getBuffer(), pos, target, off, amtToRead); consumeBuffer(amtToRead); return amtToRead; } @@ -62,6 +70,12 @@ public void write(byte[] buf, int off, int len) { @Override public final void consumeBuffer(int len) { pos += len; + if (pos == limit && temporaryBuffer != null) { + // Protocols may keep their own reference to a consumed binary value. Drop only our reference. + temporaryBuffer = null; + pos = 0; + limit = 0; + } } @Override @@ -82,7 +96,7 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException { @Override public final byte[] getBuffer() { - return buf.array(); + return temporaryBuffer == null ? buf.array() : temporaryBuffer; } @Override @@ -96,12 +110,22 @@ public final int getBytesRemainingInBuffer() { } public void resizeIfNecessary(int size) throws IOException { - buf.resizeIfNecessary(size); + if (size > maxReusableBufferSize) { + byte[] newBuffer = new byte[size]; + buf.resetToInitialCapacity(); + temporaryBuffer = newBuffer; + } else { + buf.resizeIfNecessary(size); + temporaryBuffer = null; + } } @Override public void close() { super.close(); + temporaryBuffer = null; + pos = 0; + limit = 0; buf.close(); } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java index eab45f3da48f7..9a71a1fa8865d 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java @@ -23,6 +23,7 @@ import org.apache.thrift.transport.TTransportException; import java.io.IOException; +import java.util.Arrays; /** * Note that this class is mainly copied from class {@link @@ -32,10 +33,18 @@ public class AutoScalingBufferWriteTransport extends NonOpenTransport { private final AutoResizingBuffer buf; + private final int maxReusableBufferSize; + private byte[] temporaryBuffer; private int pos; public AutoScalingBufferWriteTransport(int initialCapacity) throws IOException { + this(initialCapacity, Integer.MAX_VALUE); + } + + AutoScalingBufferWriteTransport(int initialCapacity, int maxReusableBufferSize) + throws IOException { this.buf = new AutoResizingBuffer(initialCapacity); + this.maxReusableBufferSize = maxReusableBufferSize; this.pos = 0; } @@ -49,11 +58,11 @@ public void write(byte[] toWrite, int off, int len) throws TTransportException { try { // A small prefix during serialization does not indicate a small frame. Only the complete // frame size, supplied after flushing, may trigger shrinking. - buf.ensureCapacity(pos + len); + ensureCapacity(pos + len); } catch (IOException e) { throw new TTransportException(e); } - System.arraycopy(toWrite, off, buf.array(), pos, len); + System.arraycopy(toWrite, off, getBuffer(), pos, len); pos += len; } @@ -63,21 +72,55 @@ public int getPos() { public void reset() { pos = 0; + temporaryBuffer = null; + } + + void resetAndResize() throws IOException { + int length = pos; + reset(); + // A completed large frame must not recreate the temporary array just released by reset(). + if (length <= maxReusableBufferSize) { + buf.resizeIfNecessary(length); + } } @Override public void close() { super.close(); + reset(); buf.close(); } public void resizeIfNecessary(int size) throws IOException { - buf.resizeIfNecessary(size); + if (size > maxReusableBufferSize) { + ensureCapacity(size); + } else { + buf.resizeIfNecessary(size); + temporaryBuffer = null; + } + } + + private void ensureCapacity(int size) throws IOException { + if (size > maxReusableBufferSize) { + if (temporaryBuffer == null || temporaryBuffer.length < size) { + // Preserve an unfinished frame when incremental serialization crosses the threshold. + // Keep amortized growth for subsequent small writes within the same temporary frame. + int capacity = + temporaryBuffer == null + ? size + : Math.max(size, temporaryBuffer.length + (temporaryBuffer.length >> 1)); + byte[] newBuffer = Arrays.copyOf(getBuffer(), capacity); + buf.resetToInitialCapacity(); + temporaryBuffer = newBuffer; + } + } else { + buf.ensureCapacity(size); + } } @Override public byte[] getBuffer() { - return buf.array(); + return temporaryBuffer == null ? buf.array() : temporaryBuffer; } @Override diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactory.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactory.java index 7c2c37ebf3f4a..1699aa975a1ab 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactory.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactory.java @@ -42,14 +42,28 @@ public static DeepCopyRpcTransportFactory getInstance( return create(USE_SNAPPY, thriftDefaultBufferSize, thriftMaxFrameSize); } + /** Creates a client factory whose reusable-buffer threshold is independent of other clients. */ + public static DeepCopyRpcTransportFactory getInstance( + int thriftDefaultBufferSize, int thriftMaxFrameSize, int maxReusableBufferSize) { + return create(USE_SNAPPY, thriftDefaultBufferSize, thriftMaxFrameSize, maxReusableBufferSize); + } + private static DeepCopyRpcTransportFactory create( boolean useSnappy, int thriftDefaultBufferSize, int thriftMaxFrameSize) { + return create(useSnappy, thriftDefaultBufferSize, thriftMaxFrameSize, Integer.MAX_VALUE); + } + + private static DeepCopyRpcTransportFactory create( + boolean useSnappy, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + int maxReusableBufferSize) { return useSnappy ? new DeepCopyRpcTransportFactory( new TimeoutChangeableTSnappyFramedTransport.Factory( - thriftDefaultBufferSize, thriftMaxFrameSize, true)) + thriftDefaultBufferSize, thriftMaxFrameSize, true, maxReusableBufferSize)) : new DeepCopyRpcTransportFactory( new TimeoutChangeableTFastFramedTransport.Factory( - thriftDefaultBufferSize, thriftMaxFrameSize, true)); + thriftDefaultBufferSize, thriftMaxFrameSize, true, maxReusableBufferSize)); } } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java index 36cca36b9b1e5..ce4440a9e1042 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java @@ -36,10 +36,28 @@ protected TCompressedElasticFramedTransport( int thriftMaxFrameSize, boolean copyBinary) throws TTransportException { - super(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + this(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + /** + * @param maxReusableBufferSize requests larger than this number of bytes use temporary arrays in + * all four buffers and reset their reusable storage to {@code thriftDefaultBufferSize}; + * {@link Integer#MAX_VALUE} preserves the original resizing policy + */ + protected TCompressedElasticFramedTransport( + TTransport underlying, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) + throws TTransportException { + super( + underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); try { - writeCompressBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize); - readCompressBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize); + writeCompressBuffer = + new AutoScalingBufferWriteTransport(thriftDefaultBufferSize, maxReusableBufferSize); + readCompressBuffer = + new AutoScalingBufferReadTransport(thriftDefaultBufferSize, maxReusableBufferSize); } catch (IOException e) { closeAllocatedBuffers(); throw new TTransportException(e); @@ -97,11 +115,13 @@ public void flush() throws TTransportException { underlying.write(writeCompressBuffer.getBuffer(), 0, compressedLength); } catch (IOException e) { throw new TTransportException(e); + } finally { + // The compressed output has been handed to the underlying transport, so do not retain it. + writeCompressBuffer.reset(); } - writeBuffer.reset(); try { - writeBuffer.resizeIfNecessary(length); + writeBuffer.resetAndResize(); } catch (IOException e) { throw new TTransportException(e); } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java index e15c9f594248e..8428ef3c94a4e 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java @@ -87,14 +87,26 @@ public TElasticFramedTransport( int thriftMaxFrameSize, boolean copyBinary) throws TTransportException { + this(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + protected TElasticFramedTransport( + TTransport underlying, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) + throws TTransportException { this.underlying = underlying; this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.copyBinary = copyBinary; alignUnderlyingConfiguration(); try { - readBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize); - writeBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize); + readBuffer = + new AutoScalingBufferReadTransport(thriftDefaultBufferSize, maxReusableBufferSize); + writeBuffer = + new AutoScalingBufferWriteTransport(thriftDefaultBufferSize, maxReusableBufferSize); } catch (IOException e) { closeAllocatedBuffers(); throw new TTransportException(e); @@ -316,9 +328,8 @@ public void flush() throws TTransportException { TFramedTransport.encodeFrameSize(length, i32buf); underlying.write(i32buf, 0, 4); underlying.write(writeBuffer.getBuffer(), 0, length); - writeBuffer.reset(); try { - writeBuffer.resizeIfNecessary(length); + writeBuffer.resetAndResize(); } catch (IOException e) { throw new TTransportException(e); } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSnappyElasticFramedTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSnappyElasticFramedTransport.java index fd9d4ffbba359..6d98290b11083 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSnappyElasticFramedTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSnappyElasticFramedTransport.java @@ -29,6 +29,8 @@ public class TSnappyElasticFramedTransport extends TCompressedElasticFramedTrans public static class Factory extends TElasticFramedTransport.Factory { + private final int maxReusableBufferSize; + public Factory() { this(RpcUtils.THRIFT_DEFAULT_BUF_CAPACITY, RpcUtils.THRIFT_FRAME_MAX_SIZE, true); } @@ -38,13 +40,25 @@ public Factory(int thriftDefaultBufferSize, boolean copyBinary) { } public Factory(int thriftDefaultBufferSize, int thriftMaxFrameSize, boolean copyBinary) { + this(thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + /** + * @param maxReusableBufferSize buffer requests above this size in bytes use temporary arrays + */ + public Factory( + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) { super(thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + this.maxReusableBufferSize = maxReusableBufferSize; } @Override public TTransport getTransport(TTransport trans) throws TTransportException { return new TSnappyElasticFramedTransport( - trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); } } @@ -61,6 +75,20 @@ public TSnappyElasticFramedTransport( super(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); } + /** + * @param maxReusableBufferSize buffer requests above this size in bytes use temporary arrays + */ + public TSnappyElasticFramedTransport( + TTransport underlying, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) + throws TTransportException { + super( + underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); + } + @Override protected int uncompressedLength(byte[] buf, int off, int len) throws IOException { return Snappy.uncompressedLength(buf, off, len); diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTFastFramedTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTFastFramedTransport.java index 505e1a99405a8..bd6f38eb76242 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTFastFramedTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTFastFramedTransport.java @@ -34,7 +34,18 @@ public class TimeoutChangeableTFastFramedTransport extends TElasticFramedTranspo public TimeoutChangeableTFastFramedTransport( TSocket underlying, int thriftDefaultBufferSize, int thriftMaxFrameSize, boolean copyBinary) throws TTransportException { - super(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + this(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + public TimeoutChangeableTFastFramedTransport( + TSocket underlying, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) + throws TTransportException { + super( + underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); this.underlyingSocket = underlying; } @@ -59,21 +70,31 @@ public static class Factory extends TTransportFactory { private final int thriftDefaultBufferSize; protected final int thriftMaxFrameSize; protected final boolean copyBinary; + private final int maxReusableBufferSize; public Factory(int thriftDefaultBufferSize, int thriftMaxFrameSize, boolean copyBinary) { + this(thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + public Factory( + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) { this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.copyBinary = copyBinary; + this.maxReusableBufferSize = maxReusableBufferSize; } @Override public TTransport getTransport(TTransport trans) throws TTransportException { - if (trans instanceof TSocket) { + if (trans instanceof TSocket socket) { return new TimeoutChangeableTFastFramedTransport( - (TSocket) trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + socket, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); } else { return new TElasticFramedTransport( - trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); } } } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTSnappyFramedTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTSnappyFramedTransport.java index ecd807e38d65a..d8c9cb975bdfc 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTSnappyFramedTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TimeoutChangeableTSnappyFramedTransport.java @@ -34,7 +34,18 @@ public class TimeoutChangeableTSnappyFramedTransport extends TSnappyElasticFrame public TimeoutChangeableTSnappyFramedTransport( TSocket underlying, int thriftDefaultBufferSize, int thriftMaxFrameSize, boolean copyBinary) throws TTransportException { - super(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + this(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + public TimeoutChangeableTSnappyFramedTransport( + TSocket underlying, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) + throws TTransportException { + super( + underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); this.underlyingSocket = underlying; } @@ -59,21 +70,31 @@ public static class Factory extends TTransportFactory { private final int thriftDefaultBufferSize; protected final int thriftMaxFrameSize; protected final boolean copyBinary; + private final int maxReusableBufferSize; public Factory(int thriftDefaultBufferSize, int thriftMaxFrameSize, boolean copyBinary) { + this(thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE); + } + + public Factory( + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + boolean copyBinary, + int maxReusableBufferSize) { this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.copyBinary = copyBinary; + this.maxReusableBufferSize = maxReusableBufferSize; } @Override public TTransport getTransport(TTransport trans) throws TTransportException { - if (trans instanceof TSocket) { + if (trans instanceof TSocket socket) { return new TimeoutChangeableTSnappyFramedTransport( - (TSocket) trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + socket, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); } else { return new TSnappyElasticFramedTransport( - trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary); + trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize); } } } diff --git a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactoryTest.java b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactoryTest.java index 21c7d2945a32e..7ca73ee4c810b 100644 --- a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactoryTest.java +++ b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/DeepCopyRpcTransportFactoryTest.java @@ -19,19 +19,60 @@ package org.apache.iotdb.rpc; +import org.apache.thrift.transport.TMemoryBuffer; import org.apache.thrift.transport.TMemoryInputTransport; import org.apache.thrift.transport.TTransport; import org.apache.thrift.transport.TTransportException; import org.junit.Test; import java.nio.ByteBuffer; +import java.util.Random; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; public class DeepCopyRpcTransportFactoryTest { + // Two client factories use different thresholds in the same JVM. Both plain and Snappy frames + // must round-trip, while only the configured client releases its large reusable buffers. + @Test + public void testIndependentReusableBufferThresholds() throws Exception { + boolean previousUseSnappy = BaseRpcTransportFactory.USE_SNAPPY; + try { + for (boolean snappy : new boolean[] {false, true}) { + BaseRpcTransportFactory.setUseSnappy(snappy); + try (TElasticFramedTransport temporary = + (TElasticFramedTransport) + DeepCopyRpcTransportFactory.getInstance(32, 4096, 256) + .getTransport(new TMemoryBuffer(0)); + TElasticFramedTransport reusable = + (TElasticFramedTransport) + DeepCopyRpcTransportFactory.getInstance(32, 4096) + .getTransport(new TMemoryBuffer(0))) { + byte[] payload = new byte[1024]; + new Random(0).nextBytes(payload); + for (TElasticFramedTransport transport : + new TElasticFramedTransport[] {temporary, reusable}) { + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + } + assertEquals(32, temporary.writeBuffer.getBuffer().length); + assertEquals(32, temporary.readBuffer.getBuffer().length); + assertTrue(reusable.writeBuffer.getBuffer().length >= payload.length); + assertTrue(reusable.readBuffer.getBuffer().length >= payload.length); + } + } + } finally { + BaseRpcTransportFactory.setUseSnappy(previousUseSnappy); + } + } + @Test public void testGetInstanceReturnsIndependentClientFactories() { assertNotSame( diff --git a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransportTest.java b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransportTest.java new file mode 100644 index 0000000000000..829dce8a27b80 --- /dev/null +++ b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransportTest.java @@ -0,0 +1,295 @@ +/* + * 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.iotdb.rpc; + +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.protocol.TCompactProtocol; +import org.apache.thrift.protocol.TProtocol; +import org.apache.thrift.transport.TMemoryBuffer; +import org.apache.thrift.transport.TTransportException; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Random; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +public class TCompressedElasticFramedTransportTest { + + private static final int INITIAL_CAPACITY = 32; + private static final int BUFFER_THRESHOLD = 256; + private static final int MAX_FRAME_SIZE = 4096; + + private final AtomicLong usedMemory = new AtomicLong(); + + @Before + public void setUp() { + AutoResizingBufferMemoryManager.setMemoryControl( + new AutoResizingBufferMemoryControl() { + @Override + public boolean allocate(long sizeInBytes) { + usedMemory.addAndGet(sizeInBytes); + return true; + } + + @Override + public void release(long sizeInBytes) { + usedMemory.addAndGet(-sizeInBytes); + } + }); + } + + @After + public void tearDown() { + AutoResizingBufferMemoryManager.resetMemoryControl(); + } + + // Grow all four reusable buffers, then send an incompressible large frame in pieces. The large + // frame must round-trip while all reusable buffers return to their initial, accounted capacities. + @Test + public void testLargeFrameResetsAllBuffersAndSmallFramesResumeReuse() throws Exception { + TSnappyElasticFramedTransport.Factory factory = + new TSnappyElasticFramedTransport.Factory( + INITIAL_CAPACITY, MAX_FRAME_SIZE, false, BUFFER_THRESHOLD); + try (TSnappyElasticFramedTransport transport = + (TSnappyElasticFramedTransport) factory.getTransport(new TMemoryBuffer(0))) { + roundTrip(transport, randomBytes(128)); + for (AutoResizingBuffer buffer : reusableBuffers(transport)) { + assertTrue(buffer.array().length > INITIAL_CAPACITY); + } + + byte[] payload = randomBytes(1024); + transport.write(payload, 0, 200); + transport.write(payload, 200, 100); + transport.write(payload, 300, payload.length - 300); + assertEquals(INITIAL_CAPACITY, reusableBuffer(transport.writeBuffer).array().length); + transport.flush(); + assertEquals(INITIAL_CAPACITY, transport.writeBuffer.getBuffer().length); + assertEquals( + INITIAL_CAPACITY, compressionBuffer(transport, "writeCompressBuffer").array().length); + + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, 100); + assertEquals(payload.length - 100, transport.getBytesRemainingInBuffer()); + assertEquals(payload.length, transport.getBuffer().length); + for (AutoResizingBuffer buffer : reusableBuffers(transport)) { + assertEquals(INITIAL_CAPACITY, buffer.array().length); + } + assertEquals(4L * INITIAL_CAPACITY, usedMemory.get()); + transport.readAll(actual, 100, actual.length - 100); + assertArrayEquals(payload, actual); + assertEquals(INITIAL_CAPACITY, transport.getBuffer().length); + + roundTrip(transport, randomBytes(64)); + for (AutoResizingBuffer buffer : reusableBuffers(transport)) { + assertTrue(buffer.array().length > INITIAL_CAPACITY); + } + } + assertEquals(0, usedMemory.get()); + } + + // Exactly the threshold still uses the reusable read buffer. One extra byte selects temporary + // storage, preserves partial reads, and releases that storage when the last byte is consumed. + @Test + public void testReadThresholdBoundaryAndTemporaryBufferLifetime() throws Exception { + try (AutoScalingBufferReadTransport input = + new AutoScalingBufferReadTransport(INITIAL_CAPACITY, BUFFER_THRESHOLD)) { + TMemoryBuffer wire = new TMemoryBuffer(0); + byte[] payload = randomBytes(BUFFER_THRESHOLD); + wire.write(payload); + input.fill(wire, payload.length); + assertSame(reusableBuffer(input).array(), input.getBuffer()); + assertEquals(BUFFER_THRESHOLD, usedMemory.get()); + byte[] actual = new byte[payload.length]; + input.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + + payload = randomBytes(BUFFER_THRESHOLD + 1); + wire.write(payload); + input.fill(wire, payload.length); + assertEquals(payload.length, input.getBuffer().length); + assertEquals(INITIAL_CAPACITY, reusableBuffer(input).array().length); + assertEquals(INITIAL_CAPACITY, usedMemory.get()); + actual = new byte[payload.length]; + input.readAll(actual, 0, 100); + input.readAll(actual, 100, actual.length - 100); + assertArrayEquals(payload, actual); + assertEquals(0, input.getBytesRemainingInBuffer()); + assertEquals(INITIAL_CAPACITY, input.getBuffer().length); + } + assertEquals(0, usedMemory.get()); + } + + // A frame crosses the threshold after an exact-boundary prefix, then grows through more small + // writes. All bytes must survive both transitions, and reset must release the temporary buffer. + @Test + public void testIncrementalWritesPreservePrefixAcrossThreshold() throws Exception { + try (AutoScalingBufferWriteTransport output = + new AutoScalingBufferWriteTransport(INITIAL_CAPACITY, BUFFER_THRESHOLD)) { + byte[] payload = randomBytes(1024); + output.write(payload, 0, BUFFER_THRESHOLD); + assertSame(reusableBuffer(output).array(), output.getBuffer()); + assertEquals(BUFFER_THRESHOLD, usedMemory.get()); + for (int i = BUFFER_THRESHOLD; i < payload.length; i++) { + output.write(payload, i, 1); + } + assertEquals(payload.length, output.getPos()); + assertArrayEquals(payload, Arrays.copyOf(output.getBuffer(), output.getPos())); + assertEquals(INITIAL_CAPACITY, reusableBuffer(output).array().length); + assertEquals(INITIAL_CAPACITY, usedMemory.get()); + + output.reset(); + assertEquals(0, output.getPos()); + assertEquals(INITIAL_CAPACITY, output.getBuffer().length); + output.write(payload, 0, 64); + assertSame(reusableBuffer(output).array(), output.getBuffer()); + assertArrayEquals(Arrays.copyOf(payload, 64), Arrays.copyOf(output.getBuffer(), 64)); + } + assertEquals(0, usedMemory.get()); + } + + // A highly compressible payload has a small wire frame but a large uncompressed frame. Both + // sizing paths must honor the threshold; zero and thresholds below the initial size also work. + @Test + public void testCompressibleFramesAndSmallThresholds() throws Exception { + for (int threshold : new int[] {0, INITIAL_CAPACITY / 2, BUFFER_THRESHOLD}) { + try (TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport( + new TMemoryBuffer(0), INITIAL_CAPACITY, MAX_FRAME_SIZE, true, threshold)) { + roundTrip(transport, new byte[1024]); + for (AutoResizingBuffer buffer : reusableBuffers(transport)) { + assertEquals(INITIAL_CAPACITY, buffer.array().length); + } + assertEquals(4L * INITIAL_CAPACITY, usedMemory.get()); + roundTrip(transport, new byte[] {42}); + } + assertEquals(0, usedMemory.get()); + } + } + + // Binary and compact protocols, with and without binary copying, read successive large/small + // frames. A previously returned binary value must remain valid after temporary storage is freed. + @Test + public void testProtocolsRetainBinaryValuesAcrossFrames() throws Exception { + for (boolean compact : new boolean[] {false, true}) { + for (boolean copyBinary : new boolean[] {false, true}) { + try (TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport( + new TMemoryBuffer(0), + INITIAL_CAPACITY, + MAX_FRAME_SIZE, + copyBinary, + BUFFER_THRESHOLD)) { + TProtocol protocol = + compact ? new TCompactProtocol(transport) : new TBinaryProtocol(transport); + byte[] payload = randomBytes(1024); + protocol.writeBinary(ByteBuffer.wrap(payload)); + transport.flush(); + ByteBuffer firstValue = protocol.readBinary(); + protocol.writeBinary(ByteBuffer.wrap(new byte[] {42})); + transport.flush(); + assertEquals(ByteBuffer.wrap(new byte[] {42}), protocol.readBinary()); + protocol.writeBinary(ByteBuffer.wrap(randomBytes(2048))); + transport.flush(); + protocol.readBinary(); + assertEquals(ByteBuffer.wrap(payload), firstValue); + } + assertEquals(0, usedMemory.get()); + } + } + } + + // Temporary storage must not bypass the maximum wire frame size. A rejected flush closes the + // transport and releases all four reusable allocations, including a pending temporary write. + @Test + public void testTemporaryFramesStillEnforceFrameLimit() throws Exception { + try (TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport(new TMemoryBuffer(0), INITIAL_CAPACITY, 128, true, 64)) { + transport.write(randomBytes(1024)); + TTransportException exception = assertThrows(TTransportException.class, transport::flush); + assertEquals(TTransportException.CORRUPTED_DATA, exception.getType()); + assertEquals(0, usedMemory.get()); + assertEquals(INITIAL_CAPACITY, transport.writeBuffer.getBuffer().length); + } + assertEquals(0, usedMemory.get()); + } + + // The original constructor keeps its reuse policy: even a large frame leaves grown reusable + // buffers in place until the existing shrink policy runs, and close releases their accounting. + @Test + public void testOriginalConstructorPreservesReusePolicy() throws Exception { + try (TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport( + new TMemoryBuffer(0), INITIAL_CAPACITY, MAX_FRAME_SIZE, true)) { + roundTrip(transport, randomBytes(1024)); + for (AutoResizingBuffer buffer : reusableBuffers(transport)) { + assertTrue(buffer.array().length > BUFFER_THRESHOLD); + } + } + assertEquals(0, usedMemory.get()); + } + + private void roundTrip(TSnappyElasticFramedTransport transport, byte[] payload) + throws TTransportException { + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + } + + private byte[] randomBytes(int size) { + byte[] payload = new byte[size]; + new Random(0).nextBytes(payload); + return payload; + } + + private AutoResizingBuffer[] reusableBuffers(TSnappyElasticFramedTransport transport) + throws Exception { + return new AutoResizingBuffer[] { + reusableBuffer(transport.writeBuffer), + reusableBuffer(transport.readBuffer), + compressionBuffer(transport, "writeCompressBuffer"), + compressionBuffer(transport, "readCompressBuffer") + }; + } + + private AutoResizingBuffer compressionBuffer( + TCompressedElasticFramedTransport transport, String name) throws Exception { + Field field = TCompressedElasticFramedTransport.class.getDeclaredField(name); + field.setAccessible(true); + return reusableBuffer(field.get(transport)); + } + + private AutoResizingBuffer reusableBuffer(Object transport) throws Exception { + Field field = transport.getClass().getDeclaredField("buf"); + field.setAccessible(true); + return (AutoResizingBuffer) field.get(transport); + } +} diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java index c41185e1bec30..af89ec6bb73ca 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/AbstractSessionBuilder.java @@ -42,6 +42,7 @@ public abstract class AbstractSessionBuilder { public ZoneId zoneId = null; public int thriftDefaultBufferSize = SessionConfig.DEFAULT_INITIAL_BUFFER_CAPACITY; public int thriftMaxFrameSize = SessionConfig.DEFAULT_MAX_FRAME_SIZE; + public int thriftMaxReusableBufferSize = Integer.MAX_VALUE; public int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; // this field only take effect in write request, nothing to do with any other type requests, // like query, load and so on. diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java index fe91277a27dc8..62593dfcd495d 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java @@ -87,6 +87,7 @@ public class NodesSupplier implements INodeSupplier, Runnable { private final int thriftDefaultBufferSize; private final int thriftMaxFrameSize; + private final int thriftMaxReusableBufferSize; private final int connectionTimeoutInMs; @@ -163,6 +164,49 @@ public static NodesSupplier createNodeSupplier( boolean enableRPCCompression, String version) { + return createNodeSupplier( + endPointList, + seedNodes, + executorService, + userName, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompression, + version, + Integer.MAX_VALUE); + } + + /** Bootstrap discovery using the owning Session's buffer reuse threshold. */ + @SuppressWarnings("unsafeThreadSchedule") + public static NodesSupplier createNodeSupplier( + List endPointList, + List seedNodes, + ScheduledExecutorService executorService, + String userName, + String password, + ZoneId zoneId, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + int connectionTimeoutInMs, + boolean useSSL, + String trustStore, + String trustStorePwd, + String keyStore, + String keyStorePwd, + String sslProtocol, + boolean enableRPCCompression, + String version, + int thriftMaxReusableBufferSize) { + NodesSupplier nodesSupplier = new NodesSupplier( endPointList, @@ -181,7 +225,8 @@ public static NodesSupplier createNodeSupplier( keyStorePwd, sslProtocol, enableRPCCompression, - version); + version, + thriftMaxReusableBufferSize); // call executorService.scheduleAtFixedRate here in a separate line // instead of the constructor to prevent leaking the "this" reference to @@ -208,7 +253,8 @@ private NodesSupplier( String keyStorePwd, String sslProtocol, boolean enableRPCCompression, - String version) { + String version, + int thriftMaxReusableBufferSize) { this.seedNodes = Collections.unmodifiableList( seedNodes.stream().distinct().map(TEndPoint::deepCopy).collect(Collectors.toList())); @@ -228,6 +274,7 @@ private NodesSupplier( this.zoneId = zoneId == null ? ZoneId.systemDefault() : zoneId; this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; + this.thriftMaxReusableBufferSize = thriftMaxReusableBufferSize; this.connectionTimeoutInMs = connectionTimeoutInMs; this.version = version; } @@ -351,7 +398,11 @@ private boolean refreshFrom(ThriftConnection current) { private ThriftConnection createConnection(TEndPoint endPoint) { ThriftConnection current = new ThriftConnection( - endPoint, thriftDefaultBufferSize, thriftMaxFrameSize, connectionTimeoutInMs); + endPoint, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + thriftMaxReusableBufferSize); synchronized (clientLock) { if (closed) { return null; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 9e3d41ca65c0f..89b3e348b163a 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -154,6 +154,7 @@ public class Session implements ISession { protected ZoneId zoneId; protected int thriftDefaultBufferSize; protected int thriftMaxFrameSize; + protected int thriftMaxReusableBufferSize = Integer.MAX_VALUE; protected TEndPoint defaultEndPoint; protected SessionConnection defaultSessionConnection; private boolean isClosed = true; @@ -487,6 +488,7 @@ public Session(AbstractSessionBuilder builder) { this.zoneId = builder.zoneId; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; + this.thriftMaxReusableBufferSize = builder.thriftMaxReusableBufferSize; this.connectionTimeoutInMs = builder.connectionTimeoutInMs; this.version = builder.version; this.useSSL = builder.useSSL; @@ -570,7 +572,8 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout keyStorePwd, sslProtocol, enableRPCCompaction, - version.toString()); + version.toString(), + thriftMaxReusableBufferSize); } else { // A static supplier has no separate seed tier. Keep every configured address after reopen. Set staticNodes = new LinkedHashSet<>(getNodeUrls()); @@ -4534,6 +4537,17 @@ public Builder thriftMaxFrameSize(int thriftMaxFrameSize) { return this; } + /** + * Buffer requests larger than this number of bytes use temporary arrays and reset reusable + * buffers to {@code thriftDefaultBufferSize}. Applies to compressed and uncompressed + * transports. The default {@link Integer#MAX_VALUE} keeps the original reuse policy; zero makes + * every nonempty buffer request use temporary storage. + */ + public Builder thriftMaxReusableBufferSize(int thriftMaxReusableBufferSize) { + this.thriftMaxReusableBufferSize = thriftMaxReusableBufferSize; + return this; + } + public Builder enableRedirection(boolean enableRedirection) { this.enableRedirection = enableRedirection; return this; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java index 49ecfef949440..f2158ef0ffacb 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java @@ -202,7 +202,9 @@ private void init( throws IoTDBConnectionException, StatementExecutionException { DeepCopyRpcTransportFactory transportFactory = DeepCopyRpcTransportFactory.getInstance( - session.thriftDefaultBufferSize, session.thriftMaxFrameSize); + session.thriftDefaultBufferSize, + session.thriftMaxFrameSize, + session.thriftMaxReusableBufferSize); try { if (transport != null && transport.isOpen()) { close(); diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index a39af4c847a95..b4be67b11b62c 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -158,6 +158,20 @@ public TableSessionBuilder thriftMaxFrameSize(int thriftMaxFrameSize) { return this; } + /** + * Buffer requests larger than this threshold use temporary arrays and reset reusable buffers to + * {@code thriftDefaultBufferSize}. Applies to compressed and uncompressed transports. Zero makes + * every nonempty buffer request use temporary storage. + * + * @param thriftMaxReusableBufferSize the buffer reuse threshold in bytes. + * @return the current {@link TableSessionBuilder} instance. + * @defaultValue Integer.MAX_VALUE (keeps the original reuse policy) + */ + public TableSessionBuilder thriftMaxReusableBufferSize(int thriftMaxReusableBufferSize) { + this.thriftMaxReusableBufferSize = thriftMaxReusableBufferSize; + return this; + } + /** * Enables or disables redirection for cluster nodes. * diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java index 71cc63526d711..e15320c1a1519 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java @@ -56,6 +56,7 @@ public class ThriftConnection { protected final int thriftDefaultBufferSize; protected final int thriftMaxFrameSize; + private final int thriftMaxReusableBufferSize; protected final int connectionTimeoutInMs; @@ -75,9 +76,24 @@ public ThriftConnection( int thriftDefaultBufferSize, int thriftMaxFrameSize, int connectionTimeoutInMs) { + this( + endPoint, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + Integer.MAX_VALUE); + } + + ThriftConnection( + TEndPoint endPoint, + int thriftDefaultBufferSize, + int thriftMaxFrameSize, + int connectionTimeoutInMs, + int thriftMaxReusableBufferSize) { this.endPoint = endPoint; this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; + this.thriftMaxReusableBufferSize = thriftMaxReusableBufferSize; this.connectionTimeoutInMs = connectionTimeoutInMs; } @@ -207,7 +223,8 @@ protected TTransport createTransport( String sslProtocol) throws TTransportException { DeepCopyRpcTransportFactory factory = - DeepCopyRpcTransportFactory.getInstance(thriftDefaultBufferSize, thriftMaxFrameSize); + DeepCopyRpcTransportFactory.getInstance( + thriftDefaultBufferSize, thriftMaxFrameSize, thriftMaxReusableBufferSize); return useSSL ? factory.getTransport( endPoint.getIp(), diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index 16b79c6415c4b..a3e5b6d473b7f 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java @@ -143,6 +143,7 @@ public class SessionPool implements ISessionPool { private int thriftDefaultBufferSize; private int thriftMaxFrameSize; + private int thriftMaxReusableBufferSize = Integer.MAX_VALUE; /** * Timeout of query can be set by users. A negative number means using the default configuration @@ -544,6 +545,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; + this.thriftMaxReusableBufferSize = builder.thriftMaxReusableBufferSize; this.enableAutoFetch = builder.enableAutoFetch; this.useSSL = builder.useSSL; this.trustStore = builder.trustStore; @@ -603,6 +605,7 @@ private Session constructNewSession() { .zoneId(zoneId) .thriftDefaultBufferSize(thriftDefaultBufferSize) .thriftMaxFrameSize(thriftMaxFrameSize) + .thriftMaxReusableBufferSize(thriftMaxReusableBufferSize) .enableRedirection(enableRedirection) .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) .version(version) @@ -632,6 +635,7 @@ private Session constructNewSession() { .zoneId(zoneId) .thriftDefaultBufferSize(thriftDefaultBufferSize) .thriftMaxFrameSize(thriftMaxFrameSize) + .thriftMaxReusableBufferSize(thriftMaxReusableBufferSize) .enableRedirection(enableRedirection) .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) .version(version) @@ -675,6 +679,7 @@ private void initThreadPool() { private void initAvailableNodes(List endPointList) { this.availableNodes = NodesSupplier.createNodeSupplier( + endPointList, endPointList, executorService, user, @@ -690,7 +695,8 @@ private void initAvailableNodes(List endPointList) { keyStorePwd, sslProtocol, enableThriftCompression, - version.toString()); + version.toString(), + thriftMaxReusableBufferSize); } // if this method throws an exception, either the server is broken, or the ip/port/user/password @@ -3747,6 +3753,20 @@ public Builder thriftMaxFrameSize(int thriftMaxFrameSize) { return this; } + /** + * Buffer requests larger than this threshold use temporary arrays and reset reusable buffers to + * {@code thriftDefaultBufferSize}. Applies to compressed and uncompressed transports for all + * sessions and background node discovery. The default {@link Integer#MAX_VALUE} keeps the + * original reuse policy; zero makes every nonempty buffer request use temporary storage. + * + * @param thriftMaxReusableBufferSize the buffer reuse threshold in bytes + * @return this builder + */ + public Builder thriftMaxReusableBufferSize(int thriftMaxReusableBufferSize) { + this.thriftMaxReusableBufferSize = thriftMaxReusableBufferSize; + return this; + } + public Builder enableThriftRpcCompaction(boolean enableCompaction) { this.isThriftRpcCompactionEnabled = enableCompaction; return this; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index f19c5fc8878a6..8104a774306fc 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -171,6 +171,21 @@ public TableSessionPoolBuilder thriftMaxFrameSize(int thriftMaxFrameSize) { return this; } + /** + * Buffer requests larger than this threshold use temporary arrays and reset reusable buffers to + * {@code thriftDefaultBufferSize}. Applies to compressed and uncompressed transports for all + * sessions and background node discovery. Zero makes every nonempty buffer request use temporary + * storage. + * + * @param thriftMaxReusableBufferSize the buffer reuse threshold in bytes. + * @return the current {@link TableSessionPoolBuilder} instance. + * @defaultValue Integer.MAX_VALUE (keeps the original reuse policy) + */ + public TableSessionPoolBuilder thriftMaxReusableBufferSize(int thriftMaxReusableBufferSize) { + this.thriftMaxReusableBufferSize = thriftMaxReusableBufferSize; + return this; + } + /** * Enables or disables rpc compression for the connection. * diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java index 24f7e6ab7cb27..ad5767a07c609 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionSeedNodeFallbackTest.java @@ -22,9 +22,12 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.ITableSessionPool; import org.apache.iotdb.isession.pool.SessionDataSetWrapper; +import org.apache.iotdb.rpc.AutoScalingBufferWriteTransport; import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.rpc.UrlUtils; @@ -33,6 +36,7 @@ import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion; import org.apache.iotdb.session.pool.SessionPool; +import org.apache.iotdb.session.pool.TableSessionPoolBuilder; import org.apache.thrift.TApplicationException; import org.apache.thrift.TException; @@ -71,6 +75,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -83,6 +88,8 @@ * Exercise discovery and reconnection over real Thrift connections with controlled availability. */ public class SessionSeedNodeFallbackTest { + private static final String LARGE_QUERY = "SHOW AVAILABLE URLS" + " ".repeat(4096); + private final List nodes = new ArrayList<>(); private final List attempts = new CopyOnWriteArrayList<>(); private final List queries = new CopyOnWriteArrayList<>(); @@ -457,6 +464,136 @@ public void testPoolListConstructorCopiesUrlsForFutureSessions() throws Exceptio } } + // TableSession must pass the default, finite and zero thresholds into its opened transport; + // a large request retains its write buffer only with the default reuse policy. + @Test + public void testTableSessionBufferReuseThreshold() throws Exception { + for (int threshold : new int[] {Integer.MAX_VALUE, 128, 0}) { + TableSessionBuilder builder = + new TableSessionBuilder() + .nodeUrls(Collections.singletonList(url(0))) + .thriftDefaultBufferSize(32) + .enableAutoFetch(false); + if (threshold != Integer.MAX_VALUE) { + builder.thriftMaxReusableBufferSize(threshold); + } + try (ITableSession tableSession = builder.build()) { + Session session = Whitebox.getInternalState(tableSession, "session"); + try (SessionDataSet result = tableSession.executeQueryStatement(LARGE_QUERY)) { + assertTrue(result.hasNext()); + assertLargeRequestBuffer(session, threshold); + } + assertTrue(reconnect(session.defaultSessionConnection)); + try (SessionDataSet result = tableSession.executeQueryStatement(LARGE_QUERY)) { + assertTrue(result.hasNext()); + assertLargeRequestBuffer(session, threshold); + } + } + } + } + + // Both host/port and seed-list pools must configure every lazily opened Session and their + // background discovery connection, preserving the default policy when no threshold is supplied. + @Test + public void testSessionPoolBufferReuseThreshold() throws Exception { + // Keep queries on the default connection whose retained buffer is checked below. + advertised = endpoints(0); + for (boolean singleHost : new boolean[] {true, false}) { + for (int threshold : new int[] {Integer.MAX_VALUE, 128, 0}) { + SessionPool.Builder builder = + new SessionPool.Builder().maxSize(2).thriftDefaultBufferSize(32); + if (singleHost) { + builder.host(nodes.get(0).endpoint.getIp()).port(nodes.get(0).endpoint.getPort()); + } else { + builder.nodeUrls(Collections.singletonList(url(0))); + } + if (threshold != Integer.MAX_VALUE) { + builder.thriftMaxReusableBufferSize(threshold); + } + SessionPool pool = builder.build(); + try { + try (SessionDataSetWrapper first = pool.executeQueryStatement(LARGE_QUERY); + SessionDataSetWrapper second = pool.executeQueryStatement(LARGE_QUERY)) { + assertTrue(first.hasNext()); + assertTrue(second.hasNext()); + assertNotSame(first.getSession(), second.getSession()); + assertLargeRequestBuffer((Session) first.getSession(), threshold); + assertLargeRequestBuffer((Session) second.getSession(), threshold); + } + assertDiscoveryBufferThreshold(pool, threshold); + } finally { + pool.close(); + } + } + } + } + + // Holding two table sessions forces the pool to open two connections. Both connections and + // discovery must inherit the default, finite or zero threshold from TableSessionPoolBuilder. + @Test + public void testTableSessionPoolBufferReuseThreshold() throws Exception { + advertised = endpoints(0); + for (int threshold : new int[] {Integer.MAX_VALUE, 128, 0}) { + TableSessionPoolBuilder builder = + new TableSessionPoolBuilder() + .nodeUrls(Collections.singletonList(url(0))) + .maxSize(2) + .thriftDefaultBufferSize(32); + if (threshold != Integer.MAX_VALUE) { + builder.thriftMaxReusableBufferSize(threshold); + } + try (ITableSessionPool pool = builder.build()) { + try (ITableSession first = pool.getSession(); + ITableSession second = pool.getSession()) { + Session firstSession = Whitebox.getInternalState(first, "session"); + Session secondSession = Whitebox.getInternalState(second, "session"); + assertNotSame(firstSession, secondSession); + try (SessionDataSet result = first.executeQueryStatement(LARGE_QUERY)) { + assertTrue(result.hasNext()); + assertLargeRequestBuffer(firstSession, threshold); + } + try (SessionDataSet result = second.executeQueryStatement(LARGE_QUERY)) { + assertTrue(result.hasNext()); + assertLargeRequestBuffer(secondSession, threshold); + } + } + assertDiscoveryBufferThreshold(Whitebox.getInternalState(pool, "sessionPool"), threshold); + } + } + } + + private void assertLargeRequestBuffer(Session session, int threshold) { + // Check before closing the result: closeOperation sends a small frame that can grow reuse + // again. + AutoScalingBufferWriteTransport output = + assertTransportBufferThreshold(session.defaultSessionConnection, threshold); + if (threshold == Integer.MAX_VALUE) { + assertTrue(output.getBuffer().length >= LARGE_QUERY.length()); + } else { + assertEquals(32, output.getBuffer().length); + } + } + + private void assertDiscoveryBufferThreshold(SessionPool pool, int threshold) { + NodesSupplier current = Whitebox.getInternalState(pool, "availableNodes"); + current.run(); + ThriftConnection connection = Whitebox.getInternalState(current, "client"); + AutoScalingBufferWriteTransport output = assertTransportBufferThreshold(connection, threshold); + if (threshold == 0) { + assertEquals(32, output.getBuffer().length); + } + } + + private AutoScalingBufferWriteTransport assertTransportBufferThreshold( + Object connection, int threshold) { + TTransport transport = Whitebox.getInternalState(connection, "transport"); + AutoScalingBufferWriteTransport output = Whitebox.getInternalState(transport, "writeBuffer"); + Object input = Whitebox.getInternalState(transport, "readBuffer"); + assertEquals(threshold, (int) Whitebox.getInternalState(output, "maxReusableBufferSize")); + assertEquals(threshold, (int) Whitebox.getInternalState(input, "maxReusableBufferSize")); + return output; + } + @Test public void testCloseAbortsBlockedSessionOpenAndDoesNotTryRemainingCandidates() throws Exception { supplier.close();