From 172cf5953f3b0b5581cd5e0e3ebeb502cfa0d4cd Mon Sep 17 00:00:00 2001 From: Ananya Garg Date: Wed, 4 Jun 2025 13:32:37 +0530 Subject: [PATCH 1/5] Optimized memory usage for TDS packets using ByteBufferManager and BufferPool. --- .../microsoft/sqlserver/jdbc/BufferPool.java | 24 +++++++++++ .../sqlserver/jdbc/ByteBufferManager.java | 41 +++++++++++++++++++ .../microsoft/sqlserver/jdbc/IOBuffer.java | 8 +++- 3 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java create mode 100644 src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java b/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java new file mode 100644 index 000000000..f160771f9 --- /dev/null +++ b/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java @@ -0,0 +1,24 @@ +package com.microsoft.sqlserver.jdbc; + +import java.util.concurrent.ConcurrentLinkedQueue; + +public class BufferPool { + private final int bufferSize; + private final ConcurrentLinkedQueue pool = new ConcurrentLinkedQueue<>(); + + BufferPool(int bufferSize) { + this.bufferSize = bufferSize; + } + + public byte[] rent() { + byte[] buffer = pool.poll(); + return buffer == null ? new byte[bufferSize] : buffer; + } + + public void release(byte[] buffer) { + if (buffer.length == bufferSize) { + pool.offer(buffer); + } + // If the buffer size does not match, we do not store it in the pool + } +} diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java b/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java new file mode 100644 index 000000000..5496768ea --- /dev/null +++ b/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java @@ -0,0 +1,41 @@ +package com.microsoft.sqlserver.jdbc; + +import java.util.concurrent.ConcurrentHashMap; + +public class ByteBufferManager { + + private static int[] poolSizes = { 1024, 2048, 4096, 8192, 16384, 32768, 65536 }; + private final static ConcurrentHashMap pools = new ConcurrentHashMap<>(); + + static { + for (int size : poolSizes) { + pools.put(size, new BufferPool(size)); + } + } + + public static byte[] rentBytes(int size) { + int poolSize = getPoolSize(size); + + if (poolSize == -1) { + return new byte[size]; + } + return pools.get(poolSize).rent(); + } + + private static int getPoolSize(int size) { + if (size <= 0) return 1024; + + int nextPower = Integer.highestOneBit(size - 1) << 1; + return nextPower <= 65536 ? nextPower : -1; + } + + // Return a previously rented byte array + public static void release(byte[] array) { + int size = array.length; + BufferPool pool = pools.get(size); + if (pool != null) { + pool.release(array); + } + // else discard — no pooling for this size + } +} \ No newline at end of file diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java index f51a5430c..67a0dffaa 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java @@ -6675,11 +6675,16 @@ final public String toString() { } TDSPacket(int size) { - payload = new byte[size]; + // payload = new byte[size]; + payload = ByteBufferManager.rentBytes(size); payloadLength = 0; next = null; } + void releasePayload() { + ByteBufferManager.release(payload); + } + final boolean isEOM() { return TDS.STATUS_BIT_EOM == (header[TDS.PACKET_HEADER_MESSAGE_STATUS] & TDS.STATUS_BIT_EOM); } @@ -6977,6 +6982,7 @@ final boolean readPacket() throws SQLServerException { // interrupts. If an interrupt happened prior to disabling, then expect // to read the attention ack packet as well. if (newPacket.isEOM()) { + newPacket.releasePayload(); ++tdsChannel.numMsgsRcvd; // Notify the command (if any) that we've reached the end of the response. From a53abd19b57d9dbb3d4523c2a4e45cd2d098fe89 Mon Sep 17 00:00:00 2001 From: Ananya Garg Date: Wed, 4 Jun 2025 16:02:08 +0530 Subject: [PATCH 2/5] refactor: instantiated ByteBufferManager in TDSReader and injected into TDSPacket. --- .../microsoft/sqlserver/jdbc/ByteBufferManager.java | 6 +++--- .../java/com/microsoft/sqlserver/jdbc/IOBuffer.java | 13 ++++++++----- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java b/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java index 5496768ea..ef658c0be 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java @@ -7,13 +7,13 @@ public class ByteBufferManager { private static int[] poolSizes = { 1024, 2048, 4096, 8192, 16384, 32768, 65536 }; private final static ConcurrentHashMap pools = new ConcurrentHashMap<>(); - static { + public ByteBufferManager() { for (int size : poolSizes) { pools.put(size, new BufferPool(size)); } } - public static byte[] rentBytes(int size) { + public byte[] rentBytes(int size) { int poolSize = getPoolSize(size); if (poolSize == -1) { @@ -30,7 +30,7 @@ private static int getPoolSize(int size) { } // Return a previously rented byte array - public static void release(byte[] array) { + public void release(byte[] array) { int size = array.length; BufferPool pool = pools.get(size); if (pool != null) { diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java index 67a0dffaa..fff5205e6 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java @@ -6668,21 +6668,23 @@ final class TDSPacket { final byte[] payload; int payloadLength; volatile TDSPacket next; + private final ByteBufferManager byteBufferManager; final public String toString() { return "TDSPacket(SPID:" + Util.readUnsignedShortBigEndian(header, TDS.PACKET_HEADER_SPID) + " Seq:" + header[TDS.PACKET_HEADER_SEQUENCE_NUM] + ")"; } - TDSPacket(int size) { + TDSPacket(int size, ByteBufferManager byteBufferManager) { + this.byteBufferManager = byteBufferManager; // payload = new byte[size]; - payload = ByteBufferManager.rentBytes(size); + payload = byteBufferManager.rentBytes(size); payloadLength = 0; next = null; } void releasePayload() { - ByteBufferManager.release(payload); + byteBufferManager.release(payload); } final boolean isEOM() { @@ -6743,7 +6745,8 @@ final SQLServerConnection getConnection() { return con; } - private transient TDSPacket currentPacket = new TDSPacket(0); + private transient ByteBufferManager byteBufferManager = new ByteBufferManager(); + private transient TDSPacket currentPacket = new TDSPacket(0, byteBufferManager); private transient TDSPacket lastPacket = currentPacket; private int payloadOffset = 0; private int packetNum = 0; @@ -6884,7 +6887,7 @@ final boolean readPacket() throws SQLServerException { assert tdsChannel.numMsgsRcvd < tdsChannel.numMsgsSent : "numMsgsRcvd:" + tdsChannel.numMsgsRcvd + " should be less than numMsgsSent:" + tdsChannel.numMsgsSent; - TDSPacket newPacket = new TDSPacket(con.getTDSPacketSize()); + TDSPacket newPacket = new TDSPacket(con.getTDSPacketSize(), byteBufferManager); if ((null != command) && // if cancelQueryTimeout is set, we should wait for the total amount of // queryTimeout + cancelQueryTimeout to From 23fe04f98182a91b0c18b6d75ed213143da1fa63 Mon Sep 17 00:00:00 2001 From: Ananya Garg Date: Wed, 18 Jun 2025 16:52:36 +0530 Subject: [PATCH 3/5] fixed failure --- .../java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java b/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java index ef658c0be..924c9e0c6 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/ByteBufferManager.java @@ -19,7 +19,8 @@ public byte[] rentBytes(int size) { if (poolSize == -1) { return new byte[size]; } - return pools.get(poolSize).rent(); + BufferPool pool = pools.computeIfAbsent(poolSize, BufferPool::new); + return pool.rent(); } private static int getPoolSize(int size) { From 79b4b72ebdd50fa063b31281a253e5c56f618bae Mon Sep 17 00:00:00 2001 From: Ananya Garg Date: Wed, 18 Jun 2025 18:38:23 +0530 Subject: [PATCH 4/5] fixed pipeline failure --- src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java index fff5205e6..aab93d93c 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java @@ -6862,6 +6862,7 @@ private boolean nextPacket() throws SQLServerException { if (logger.isLoggable(Level.FINEST)) logger.finest(toString() + " Moving to next packet -- unlinking consumed packet"); + consumedPacket.releasePayload(); consumedPacket.next = null; } currentPacket = nextPacket; @@ -6985,7 +6986,6 @@ final boolean readPacket() throws SQLServerException { // interrupts. If an interrupt happened prior to disabling, then expect // to read the attention ack packet as well. if (newPacket.isEOM()) { - newPacket.releasePayload(); ++tdsChannel.numMsgsRcvd; // Notify the command (if any) that we've reached the end of the response. From 30dd55ef9e1f24fb4cccef65773a1d7e33f0ad0f Mon Sep 17 00:00:00 2001 From: Ananya Garg Date: Mon, 23 Jun 2025 12:32:21 +0530 Subject: [PATCH 5/5] updated code as per review comment --- src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java b/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java index f160771f9..17ceaff50 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/BufferPool.java @@ -1,5 +1,6 @@ package com.microsoft.sqlserver.jdbc; +import java.util.Arrays; import java.util.concurrent.ConcurrentLinkedQueue; public class BufferPool { @@ -18,6 +19,7 @@ public byte[] rent() { public void release(byte[] buffer) { if (buffer.length == bufferSize) { pool.offer(buffer); + Arrays.fill(buffer, (byte) 0); // Clear the array for re-use } // If the buffer size does not match, we do not store it in the pool }