Add maxResultBuffer property

src/main/java/com/microsoft/sqlserver/jdbc/ICounter.java

@@ -0,0 +1,27 @@
+/*
+ * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
+ * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+
+package com.microsoft.sqlserver.jdbc;
+
+/**
+ * Interface for MaxResultBufferCounter
+ */
+interface ICounter {
+
+    /**
+     * Increases the state of Counter
+     *
+     * @param bytes
+     *        Number of bytes to increase state
+     * @throws SQLServerException
+     *         Exception is thrown, when limit of Counter is exceeded
+     */
+    void increaseCounter(long bytes) throws SQLServerException;
+
+    /**
+     * Resets the state of Counter
+     */
+    void resetCounter();
+}

src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java

@@ -51,6 +51,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Properties;
 import java.util.Set;
 import java.util.SimpleTimeZone;
 import java.util.TimeZone;
@@ -3156,6 +3157,15 @@ boolean isEOMSent() {
         traceID = "TDSWriter@" + Integer.toHexString(hashCode()) + " (" + con.toString() + ")";
     }
 
+    /**
+     * Checks If tdsMessageType is RPC or QUERY
+     * 
+     * @return boolean
+     */
+    boolean checkIfTdsMessageTypeIsBatchOrRPC() {
+        return tdsMessageType == TDS.PKT_QUERY || tdsMessageType == TDS.PKT_RPC;
+    }
+
     // TDS message start/end operations
 
     void preparePacket() throws SQLServerException {
@@ -6546,6 +6556,11 @@ private boolean nextPacket() throws SQLServerException {
         // This action must be synchronized against against another thread calling
         // readAllPackets() to read in ALL of the remaining packets of the current response.
         if (null == consumedPacket.next) {
+            // if the read comes from getNext() and responseBuffering is Adaptive (in this place is), then reset Counter
+            // State
+            if (command.getTDSWriter().checkIfTdsMessageTypeIsBatchOrRPC()) {
+                command.getCounter().resetCounter();
+            }
             readPacket();
 
             if (null == consumedPacket.next)
@@ -6641,6 +6656,11 @@ synchronized final boolean readPacket() throws SQLServerException {
             System.arraycopy(newPacket.header, 0, logBuffer, 0, TDS.PACKET_HEADER_SIZE);
         }
 
+        // if messageType is RPC or QUERY, then increment Counter's state
+        if (tdsChannel.getWriter().checkIfTdsMessageTypeIsBatchOrRPC()) {
+            command.getCounter().increaseCounter(packetLength);
+        }
+
         // Now for the payload...
         for (int payloadBytesRead = 0; payloadBytesRead < newPacket.payloadLength;) {
             int bytesRead = tdsChannel.read(newPacket.payload, payloadBytesRead,
@@ -7345,6 +7365,23 @@ final boolean readingResponse() {
 
     protected ArrayList<byte[]> enclaveCEKs;
 
+    // Counter reference, so maxResultBuffer property can by acknowledged
+    private ICounter counter;
+
+    ICounter getCounter() {
+        return counter;
+    }
+
+    void createCounter(ICounter previousCounter, Properties activeConnectionProperties) {
+        if (null == previousCounter) {
+            String maxResultBuffer = activeConnectionProperties
+                    .getProperty(SQLServerDriverStringProperty.MAX_RESULT_BUFFER.toString());
+            counter = new MaxResultBufferCounter(Long.parseLong(maxResultBuffer));
+        } else {
+            counter = previousCounter;
+        }
+    }
+
     /**
      * Creates this command with an optional timeout.
      *

src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java

@@ -1051,4 +1051,19 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      *        Active Directory Service Principal secret.
      */
     void setAADSecurePrincipalSecret(String AADSecurePrincipalSecret);
+
+    /**
+     * Returns value of 'maxResultBuffer' from Connection String.
+     *
+     * @return 'maxResultBuffer' property.
+     */
+    String getMaxResultBuffer();
+
+    /**
+     * Specifies value for 'maxResultBuffer' property
+     *
+     * @param maxResultBuffer
+     *        String value for 'maxResultBuffer'
+     */
+    void setMaxResultBuffer(String maxResultBuffer);
 }

src/main/java/com/microsoft/sqlserver/jdbc/MaxResultBufferCounter.java

@@ -0,0 +1,52 @@
+/*
+ * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
+ * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+
+package com.microsoft.sqlserver.jdbc;
+
+import java.text.MessageFormat;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+
+/**
+ * Implementation of ICounter for 'maxResultBuffer' property.
+ */
+public class MaxResultBufferCounter implements ICounter {
+
+    private final Logger logger = Logger.getLogger("com.microsoft.sqlserver.jdbc.MaxResultBufferCounter");
+
+    private long counter = 0;
+    private final long maxResultBuffer;
+
+    public MaxResultBufferCounter(long maxResultBuffer) {
+        this.maxResultBuffer = maxResultBuffer;
+    }
+
+    public void increaseCounter(long bytes) throws SQLServerException {
+        if (maxResultBuffer > 0) {
+            counter += bytes;
+            checkForMaxResultBufferOverflow(counter);
+        }
+    }
+
+    public void resetCounter() {
+        counter = 0;
+    }
+
+    private void checkForMaxResultBufferOverflow(long number) throws SQLServerException {
+        if (number > maxResultBuffer) {
+            if (logger.isLoggable(Level.SEVERE)) {
+                logger.log(Level.SEVERE, SQLServerException.getErrString("R_maxResultBufferPropertyExceeded"),
+                        new Object[] {number, maxResultBuffer});
+            }
+            throwExceededMaxResultBufferException(counter, maxResultBuffer);
+        }
+    }
+
+    private void throwExceededMaxResultBufferException(Object... arguments) throws SQLServerException {
+        MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_maxResultBufferPropertyExceeded"));
+        throw new SQLServerException(form.format(arguments), null);
+    }
+}

src/main/java/com/microsoft/sqlserver/jdbc/MaxResultBufferParser.java

@@ -0,0 +1,151 @@
+/*
+ * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
+ * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+
+package com.microsoft.sqlserver.jdbc;
+
+import java.lang.management.ManagementFactory;
+import java.text.MessageFormat;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+
+/**
+ * Parser created to parse String value from Connection String to equivalent number of bytes for JDBC Driver to work on.
+ */
+public class MaxResultBufferParser {
+
+    private static final Logger logger = Logger.getLogger("com.microsoft.sqlserver.jdbc.MaxResultBufferParser");
+    private static final String[] PERCENT_PHRASES = {"percent", "pct", "p"};
+    private static final String ERROR_MESSAGE = "MaxResultBuffer property is badly formatted: {0}.";
+
+    private MaxResultBufferParser() {}
+
+    /**
+     *
+     * Returns number of bytes for maxResultBuffer property
+     *
+     * @param input
+     *        String value for maxResultProperty provided in Connection String
+     * @return 'maxResultBuffer' property as number of bytes
+     * @throws SQLServerException
+     *         Is Thrown when maxResultProperty's syntax is wrong
+     */
+    public static long validateMaxResultBuffer(String input) throws SQLServerException {
+        String numberString;
+        long number = -1;
+
+        // check for null values and empty String "", if so return -1 (default value)
+        if (StringUtils.isEmpty(input) || input.equals("-1")) {
+            return number;
+        }
+
+        // check if input is number
+        if (!StringUtils.isEmpty(input) && input.matches("-?\\d+(\\.\\d+)?")) {
+            try {
+                number = Long.parseLong(input);
+            } catch (NumberFormatException e) {
+                if (logger.isLoggable(Level.SEVERE)) {
+                    logger.log(Level.SEVERE, ERROR_MESSAGE, new Object[] {input});
+                }
+                throwNewInvalidMaxResultBufferParameterException(e, input);
+            }
+            return adjustMemory(number, 1);
+        } else {
+            // check PERCENT_PHRASES
+            for (String percentPhrase : PERCENT_PHRASES) {
+                if (input.endsWith(percentPhrase)) {
+                    numberString = input.substring(0, input.length() - percentPhrase.length());
+                    try {
+                        number = Long.parseLong(numberString);
+                    } catch (NumberFormatException e) {
+                        if (logger.isLoggable(Level.SEVERE)) {
+                            logger.log(Level.SEVERE, ERROR_MESSAGE, new Object[] {input});
+                        }
+                        throwNewInvalidMaxResultBufferParameterException(e, numberString);
+                    }
+                    return adjustMemoryPercentage(number);
+                }
+            }
+
+            // check if prefix was supplied
+            long multiplier = getMultiplier(input);
+            numberString = input.substring(0, input.length() - 1);
+
+            try {
+                number = Long.parseLong(numberString);
+            } catch (NumberFormatException e) {
+                if (logger.isLoggable(Level.SEVERE)) {
+                    logger.log(Level.SEVERE, ERROR_MESSAGE, new Object[] {input});
+                }
+                throwNewInvalidMaxResultBufferParameterException(e, numberString);
+            }
+            return adjustMemory(number, multiplier);
+        }
+    }
+
+    private static void checkForNegativeValue(long value) throws SQLServerException {
+        if (value <= 0) {
+            Object[] objectToThrow = new Object[] {value};
+            MessageFormat form = new MessageFormat(
+                    SQLServerException.getErrString("R_maxResultBufferNegativeParameterValue"));
+            if (logger.isLoggable(Level.SEVERE)) {
+                logger.log(Level.SEVERE, SQLServerException.getErrString("R_maxResultBufferNegativeParameterValue"),
+                        objectToThrow);
+            }
+            throw new SQLServerException(form.format(objectToThrow), new Throwable());
+        }
+    }
+
+    private static long getMultiplier(String input) throws SQLServerException {
+        long multiplier = 1;
+        switch (Character.toUpperCase(input.charAt(input.length() - 1))) {
+            case 'K':
+                multiplier = 1_000L;
+                break;
+            case 'M':
+                multiplier = 1_000_000L;
+                break;
+            case 'G':
+                multiplier = 1_000_000_000L;
+                break;
+            case 'T':
+                multiplier = 1_000_000_000_000L;
+                break;
+            default:
+                if (logger.isLoggable(Level.SEVERE)) {
+                    logger.log(Level.SEVERE, ERROR_MESSAGE, new Object[] {input});
+                }
+                throwNewInvalidMaxResultBufferParameterException(null, input);
+        }
+        return multiplier;
+    }
+
+    private static long adjustMemoryPercentage(long percentage) throws SQLServerException {
+        checkForNegativeValue(percentage);
+        if (percentage > 90)
+            return (long) (0.9 * getMaxMemory());
+        else
+            return (long) ((percentage) / 100.0 * getMaxMemory());
+    }
+
+    private static long adjustMemory(long size, long multiplier) throws SQLServerException {
+        checkForNegativeValue(size);
+        if (size * multiplier > 0.9 * getMaxMemory())
+            return (long) (0.9 * getMaxMemory());
+        else
+            return size * multiplier;
+    }
+
+    private static long getMaxMemory() {
+        return ManagementFactory.getMemoryMXBean().getHeapMemoryUsage().getMax();
+    }
+
+    private static void throwNewInvalidMaxResultBufferParameterException(Throwable cause,
+            Object... arguments) throws SQLServerException {
+        MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_maxResultBufferInvalidSyntax"));
+        throw new SQLServerException(form.format(arguments), cause);
+    }
+
+}

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java

@@ -2196,6 +2196,11 @@ else if (0 == requestedPacketSize)
                 sendTemporalDataTypesAsStringForBulkCopy = isBooleanPropertyOn(sPropKey, sPropValue);
             }
 
+            sPropKey = SQLServerDriverStringProperty.MAX_RESULT_BUFFER.toString();
+            sPropValue = activeConnectionProperties.getProperty(sPropKey);
+            activeConnectionProperties.setProperty(sPropKey,
+                    String.valueOf(MaxResultBufferParser.validateMaxResultBuffer(sPropValue)));
+
             sPropKey = SQLServerDriverBooleanProperty.DELAY_LOADING_LOBS.toString();
             sPropValue = activeConnectionProperties.getProperty(sPropKey);
             if (null == sPropValue) {
@@ -3216,13 +3221,20 @@ final void terminate(int driverErrorCode, String message, Throwable throwable) t
      */
     boolean executeCommand(TDSCommand newCommand) throws SQLServerException {
         synchronized (schedulerLock) {
+            ICounter previousCounter = null;
             /*
              * Detach (buffer) the response from any previously executing command so that we can execute the new
              * command. Note that detaching the response does not process it. Detaching just buffers the response off of
              * the wire to clear the TDS channel.
              */
             if (null != currentCommand) {
                 try {
+
+                    /**
+                     * If currentCommand needs to be detached, reset Counter to acknowledge number of Bytes in remaining
+                     * packets
+                     */
+                    currentCommand.getCounter().resetCounter();
                     currentCommand.detach();
                 } catch (SQLServerException e) {
                     /*
@@ -3234,10 +3246,14 @@ boolean executeCommand(TDSCommand newCommand) throws SQLServerException {
                         connectionlogger.fine("Failed to detach current command : " + e.getMessage());
                     }
                 } finally {
+                    previousCounter = currentCommand.getCounter();
                     currentCommand = null;
                 }
             }
-
+            /**
+             * Add Counter reference to newCommand
+             */
+            newCommand.createCounter(previousCounter, activeConnectionProperties);
             /*
              * The implementation of this scheduler is pretty simple... Since only one command at a time may use a
              * connection (to avoid TDS protocol errors), just synchronize to serialize command execution.

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java

@@ -1056,6 +1056,17 @@ public void setSendTemporalDataTypesAsStringForBulkCopy(boolean sendTemporalData
                 sendTemporalDataTypesAsStringForBulkCopy);
     }
 
+    @Override
+    public String getMaxResultBuffer() {
+        return getStringProperty(connectionProps, SQLServerDriverStringProperty.MAX_RESULT_BUFFER.toString(),
+                SQLServerDriverStringProperty.MAX_RESULT_BUFFER.getDefaultValue());
+    }
+
+    @Override
+    public void setMaxResultBuffer(String maxResultBuffer) {
+        setStringProperty(connectionProps, SQLServerDriverStringProperty.MAX_RESULT_BUFFER.toString(), maxResultBuffer);
+    }
+
     /**
      * Sets a property string value.
      *