feat: Set memory size limit in multi-sql statements with no ring buffer

Author: luis-garza-dev

src/main/java/com/amazon/redshift/core/v3/QueryExecutorImpl.java

@@ -7,9 +7,7 @@
 package com.amazon.redshift.core.v3;
 
 import com.amazon.redshift.RedshiftProperty;
-import com.amazon.redshift.copy.CopyIn;
 import com.amazon.redshift.copy.CopyOperation;
-import com.amazon.redshift.copy.CopyOut;
 import com.amazon.redshift.core.CommandCompleteParser;
 import com.amazon.redshift.core.Encoding;
 import com.amazon.redshift.core.EncodingPredictor;
@@ -71,7 +69,6 @@
 import java.util.Set;
 import java.util.TimeZone;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.regex.Pattern;
@@ -158,6 +155,8 @@ private static boolean looksLikePrepare(String sql) {
   // Query or some execution on a socket in process
   private final Lock m_executingLock = new ReentrantLock();
 
+  private static final long INVALID_TUPLE_SIZE = -1L;
+
   /**
    * {@code CommandComplete(B)} messages are quite common, so we reuse instance to parse those
    */
@@ -1900,6 +1899,7 @@ private void processResultsOnThread(ResultHandler handler,
     												&& (!bothRowsAndStatus); // RETURNING clause 
 
     List<Tuple> tuples = null;
+    long totalTupleSize = 0;
 
     int c;
     boolean endQuery = false;
@@ -2269,8 +2269,15 @@ private void processResultsOnThread(ResultHandler handler,
                                 tuples = new ArrayList<Tuple>();
                             }
 
-                            if(!skipRow)
+                            if (!skipRow) {
+                                if (enableFetchRingBuffer) {
+                                    totalTupleSize = checkAndUpdateTupleSize(tuple, totalTupleSize, handler);
+                                    if (handler.getException() != null || totalTupleSize == INVALID_TUPLE_SIZE) {
+                                        return;
+                                    }
+                                }
                                 tuples.add(tuple);
+                            }
                         }
                     }
 
@@ -2492,6 +2499,40 @@ private void processResultsOnThread(ResultHandler handler,
 
   }
 
+  /**
+   * Checks and updates the tuple size ensuring it doesn't exceed buffer limits.
+   *
+   * @param tuple The tuple to check
+   * @param totalTupleSize Current total size of tuples
+   * @param handler Error handler for reporting issues
+   * @return Updated total size or `INVALID_TUPLE_SIZE` if an error occurred
+  */
+  private long checkAndUpdateTupleSize(Tuple tuple, long totalTupleSize, ResultHandler handler) {
+    int nodeSize = RedshiftMemoryUtils.calculateNodeSize(tuple);
+
+    // Check for potential overflow before addition
+    if (Long.MAX_VALUE - totalTupleSize < nodeSize) {
+        handler.handleError(
+                new RedshiftException(
+                        GT.tr("Buffer size calculation overflow. Current size: {0}, Adding: {1}",
+                                totalTupleSize, nodeSize),
+                        RedshiftState.OUT_OF_MEMORY));
+        return INVALID_TUPLE_SIZE;
+    }
+
+    long newTotalSize = totalTupleSize + nodeSize;
+    if (newTotalSize > this.fetchRingBufferSize) {
+        handler.handleError(
+                new RedshiftException(
+                        GT.tr("Total result set size of {0} bytes would exceed maximum allowed size of {1} bytes",
+                                newTotalSize, this.fetchRingBufferSize),
+                        RedshiftState.OUT_OF_MEMORY));
+        return INVALID_TUPLE_SIZE;
+    }
+
+    return newTotalSize;
+  }
+
   /**
    * Ignore the response message by reading the message length and skipping over those bytes in the
    * communication stream.

src/main/java/com/amazon/redshift/core/v3/RedshiftMemoryUtils.java

@@ -0,0 +1,70 @@
+package com.amazon.redshift.core.v3;
+
+import com.amazon.redshift.core.Tuple;
+import com.amazon.redshift.jdbc.RedshiftConnectionImpl;
+
+/**
+ * Utility class for memory-related calculations in Redshift JDBC driver.
+ * Provides methods to estimate memory usage of result set data structures.
+ */
+public final class RedshiftMemoryUtils {
+
+    /**
+     * Scaling factor to account for JVM memory overhead and alignment.
+     * The factor 1.2 adds a 20% safety margin to the calculated size.
+     */
+    static final double MEMORY_ESTIMATE_SCALING_FACTOR = 1.2;
+
+    /**
+     * Memory overhead for a node in a 64-bit JVM:
+     * - 8 bytes: Object reference for Tuple
+     * - 8 bytes: Reference for next pointer
+     * - 16 bytes: Object header overhead
+     */
+    private static final int NODE_OVERHEAD_64BIT = 32;
+
+    /**
+     * Memory overhead for a node in a 32-bit JVM:
+     * - 4 bytes: Object reference for Tuple
+     * - 4 bytes: Reference for next pointer
+     * - 8 bytes: Object header overhead
+     */
+    private static final int NODE_OVERHEAD_32BIT = 16;
+
+    /**
+     * JVM-specific memory overhead for node structures.
+     */
+    private static final int NODE_OVERHEAD = RedshiftConnectionImpl.IS_64_BIT_JVM
+            ? NODE_OVERHEAD_64BIT
+            : NODE_OVERHEAD_32BIT;
+
+    private RedshiftMemoryUtils() {
+        throw new AssertionError("Utility class should not be instantiated");
+    }
+
+    /**
+     * Calculates the estimated memory size in bytes for a tuple node, including JVM overhead.
+     * The calculation includes:
+     * <ul>
+     *   <li>The actual tuple data size</li>
+     *   <li>JVM object overhead (32 bytes for 64-bit JVM, 16 bytes for 32-bit JVM)</li>
+     *   <li>A scaling factor to account for additional JVM overhead</li>
+     * </ul>
+     *
+     * Note: We don't worry about overflow in this method because:
+     * - MEMORY_ESTIMATE_SCALING_FACTOR is 1.2, so even with scaling we're only increasing values by 20%
+     * - NODE_OVERHEAD is a modest constant value
+     * The main overflow risk comes from accumulating multiple tuples' sizes.
+     *
+     * @param row The tuple to calculate size for, can be null
+     * @return Estimated size in bytes, or 0 if the tuple is null
+     */
+    public static int calculateNodeSize(Tuple row) {
+        if (row == null) {
+            return 0;
+        }
+
+        int rawSize = row.getTupleSize() + NODE_OVERHEAD;
+        return (int) (rawSize * MEMORY_ESTIMATE_SCALING_FACTOR);
+    }
+}

src/main/java/com/amazon/redshift/core/v3/RedshiftRowsBlockingQueue.java

@@ -37,8 +37,6 @@ public class RedshiftRowsBlockingQueue<E> extends LinkedBlockingQueue<E> {
   private RedshiftLogger logger;  
 
   private Portal currentSuspendedPortal;
-
-  private final double MEMORY_ESTIMATE_SCALING_FACTOR = 1.2;
 
   public RedshiftRowsBlockingQueue(int capacity) {
 		super(capacity);
@@ -108,7 +106,7 @@ public void put(E e) throws InterruptedException {
 
 						super.put(e);
 
-						currentBufSize = totalFetchRingBufferSize.addAndGet(getNodeSize(row));
+						currentBufSize = totalFetchRingBufferSize.addAndGet(RedshiftMemoryUtils.calculateNodeSize(row));
 
 						if (currentBufSize < fetchRingBufferSizeCapacity)
 							notFull.signal();
@@ -118,7 +116,7 @@ public void put(E e) throws InterruptedException {
 				}
 				else {
 					super.put(e);
-					totalFetchRingBufferSize.addAndGet(getNodeSize((Tuple)e));
+					totalFetchRingBufferSize.addAndGet(RedshiftMemoryUtils.calculateNodeSize((Tuple)e));
 				}
 			}
 		} // By size
@@ -135,7 +133,7 @@ public E take() throws InterruptedException {
   		Tuple row = (Tuple)e;
   		long currentBufSize;
   		boolean bufWasFull = (totalFetchRingBufferSize.get() >= fetchRingBufferSizeCapacity);  		
-  		currentBufSize = totalFetchRingBufferSize.addAndGet(-getNodeSize(row));
+  		currentBufSize = totalFetchRingBufferSize.addAndGet(-RedshiftMemoryUtils.calculateNodeSize(row));
 
   		// Signal the waiters
   		if (bufWasFull) {
@@ -238,18 +236,4 @@ private void signalNotFull() {
           putLock.unlock();
       }
   }
-
-  /**
-   * Returns the size in bytes of an individual node of Ring buffer queue/linked list
-   */
-  private int getNodeSize(Tuple row) {
-	  /**
-	   * Node overheads are 32 bytes for 64-bit JVM and 16 bytes for 32-bit JVM
-	   * For 64-bit JVM: (8 + 8 + 16) => 8 byte reference for Tuple object + 8 byte reference for next
-	   *                 + 16 byte Node object header overhead
-	   * Each of these are reduced to half in case of 32-bit JVM.
-	   */
-	  int estimatedNodeSize = row.getTupleSize() + (RedshiftConnectionImpl.IS_64_BIT_JVM ? 32 : 16);
-	  return (int) (estimatedNodeSize * MEMORY_ESTIMATE_SCALING_FACTOR); // using a scaling factor for avoiding OOM errors
-  }
 }