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
+ */
+public 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;
@@ -3155,6 +3156,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 {
@@ -6545,6 +6555,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)
@@ -6640,6 +6655,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,
@@ -7344,6 +7364,23 @@ final boolean readingResponse() {
 
     protected ArrayList<byte[]> enclaveCEKs;
 
+    //Counter reference, so maxResultBuffer property can by acknowledged
+    private ICounter counter;
+
+    public ICounter getCounter() {
+        return counter;
+    }
+
+    public 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

@@ -1021,4 +1021,18 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      */
     void setSendTemporalDataTypesAsStringForBulkCopy(boolean sendTemporalDataTypesAsStringForBulkCopy);
 
+    /**
+     * 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,50 @@
+/*
+ * 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) {
+            logger.log(Level.WARNING, "MaxResultBuffer exceeded: {0}. Property was set to {1}.",
+                    new Object[] {number, maxResultBuffer});
+            throwExceededMaxResultBufferException(counter, maxResultBuffer);
+        }
+    }
+
+    private void throwExceededMaxResultBufferException(Object... arguments) throws SQLServerException {
+        MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_maxResultBufferPropertyDescription"));
+        throw new SQLServerException(form.format(arguments), null);
+    }
+}

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

@@ -0,0 +1,123 @@
+/*
+ * 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
+        if (StringUtils.isEmpty(input)) {
+            return number;
+        }
+        // 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) {
+                    logger.log(Level.INFO, ERROR_MESSAGE, new Object[] {input});
+                    throwNewInvalidMaxResultBufferParameterException(e, numberString);
+                }
+                return adjustMemory(number);
+            }
+        }
+        // check if only number was supplied
+        long multiplier = 1;
+        if (StringUtils.isNumeric(input)) {
+            number = Long.parseLong(input);
+            return adjustMemory(number, multiplier);
+        }
+        // check if prefix was supplied
+        multiplier = getMultiplier(input);
+
+        numberString = input.substring(0, input.length() - 1);
+
+        try {
+            number = Long.parseLong(numberString);
+        } catch (NumberFormatException e) {
+            logger.log(Level.INFO, ERROR_MESSAGE, new Object[] {input});
+            throwNewInvalidMaxResultBufferParameterException(e, numberString);
+        }
+        return adjustMemory(number, multiplier);
+    }
+
+    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:
+                logger.log(Level.INFO, ERROR_MESSAGE, new Object[] {input});
+                throwNewInvalidMaxResultBufferParameterException(null, input);
+        }
+        return multiplier;
+    }
+
+    private static long adjustMemory(long percentage) {
+        if (percentage > 90)
+            return (long) (0.9 * getMaxMemory());
+        else
+            return (long) ((percentage) / 100.0 * getMaxMemory());
+    }
+
+    private static long adjustMemory(long size, long multiplier) {
+        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"));
+        Object[] msgArgs = {arguments};
+        throw new SQLServerException(form.format(msgArgs), cause);
+    }
+
+}

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

@@ -2161,6 +2161,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) {
@@ -3179,13 +3184,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) {
                     /*
@@ -3197,10 +3209,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

@@ -1032,6 +1032,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.
      * 

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

@@ -366,7 +366,8 @@ enum SQLServerDriverStringProperty {
     KEY_STORE_PRINCIPAL_ID("keyStorePrincipalId", ""),
     CLIENT_CERTIFICATE("clientCertificate", ""),
     CLIENT_KEY("clientKey", ""),
-    CLIENT_KEY_PASSWORD("clientKeyPassword", "");
+    CLIENT_KEY_PASSWORD("clientKeyPassword", ""),
+    MAX_RESULT_BUFFER("maxResultBuffer", "");
 
     private final String name;
     private final String defaultValue;
@@ -634,7 +635,9 @@ public final class SQLServerDriver implements java.sql.Driver {
                     SQLServerDriverBooleanProperty.SEND_TEMPORAL_DATATYPES_AS_STRING_FOR_BULK_COPY.toString(),
                     Boolean.toString(SQLServerDriverBooleanProperty.SEND_TEMPORAL_DATATYPES_AS_STRING_FOR_BULK_COPY
                             .getDefaultValue()),
-                    false, TRUE_FALSE),};
+                    false, TRUE_FALSE),
+            new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.MAX_RESULT_BUFFER.toString(),
+                    SQLServerDriverStringProperty.MAX_RESULT_BUFFER.getDefaultValue(), false, null),};
 
     /**
      * Properties that can only be set by using Properties. Cannot set in connection string

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

@@ -649,5 +649,8 @@ protected Object[][] getContents() {
             {"R_unassignableError", "The class specified by the {0} property must be assignable to {1}."},
             {"R_InvalidCSVQuotes",
                     "Failed to parse the CSV file, verify that the fields are correctly enclosed in double quotes."},
-            {"R_TokenRequireUrl", "Token credentials require a URL using the HTTPS protocol scheme."},};
-};
+            {"R_TokenRequireUrl", "Token credentials require a URL using the HTTPS protocol scheme."},
+            {"R_maxResultBufferPropertyDescription",
+                    "MaxResultBuffer property exceeded: {0}. MaxResultBuffer was set to: {1}."},
+            {"R_maxResultBufferInvalidSyntax", "Invalid syntax: {0} in maxResultBuffer parameter"},};
+}