Configurable Retry Logic II - Connection Retry

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

@@ -39,6 +39,10 @@
     private static final String FORWARD_SLASH = "/";
     private static final String EQUALS_SIGN = "=";
     private static final String RETRY_EXEC = "retryExec";
+    private static final String RETRY_CONN = "retryConn";
+    private static final String STATEMENT = "statement";
+    private static final String CONNECTION = "connection";
+    private static boolean replaceFlag = false; // Are we replacing the list of transient errors?
     /**
      * The time the properties file was last modified.
      */
@@ -52,14 +56,23 @@
      */
     private static final AtomicReference<String> lastQuery = new AtomicReference<>("");
     /**
-     * The previously read rules from the connection string.
+     * The previously read statement rules from the connection string.
      */
-    private static final AtomicReference<String> prevRulesFromConnectionString = new AtomicReference<>("");
+    private static final AtomicReference<String> prevStmtRulesFromConnString = new AtomicReference<>("");
+    /**
+     * The previously read connection rules from the connection string.
+     */
+    private static final AtomicReference<String> prevConnRulesFromConnString = new AtomicReference<>("");
     /**
      * The list of statement retry rules.
      */
     private static final AtomicReference<HashMap<Integer, ConfigurableRetryRule>> stmtRules = new AtomicReference<>(
             new HashMap<>());
+    /**
+     * The list of connection retry rules.
+     */
+    private static final AtomicReference<HashMap<Integer, ConfigurableRetryRule>> connRules = new AtomicReference<>(
+            new HashMap<>());
     private static ConfigurableRetryLogic singleInstance;
 
     /**
@@ -70,7 +83,8 @@
      */
     private ConfigurableRetryLogic() throws SQLServerException {
         timeLastRead.compareAndSet(0, new Date().getTime());
-        setUpRules(null);
+        setUpRules(null, STATEMENT);
+        setUpRules(null, CONNECTION);
     }
 
     /**
@@ -102,7 +116,8 @@
 
     /**
      * If it has been INTERVAL_BETWEEN_READS_IN_MS (30 secs) since last read, see if we last did a file read, if so
-     * only reread if the file has been modified. If no file read, set up rules using the prev. connection string rules.
+     * only reread if the file has been modified. If no file read, set up rules using the previous connection
+     * string (statement and connection) rules
      *
      * @throws SQLServerException
      *         when an exception occurs
@@ -116,10 +131,12 @@
                 // If timeLastModified is set, we previously read from file, so we setUpRules also reading from file
                 File f = new File(getCurrentClassPath());
                 if (f.lastModified() != timeLastModified.get()) {
-                    setUpRules(null);
+                    setUpRules(null, STATEMENT);
+                    setUpRules(null, CONNECTION);
                 }
             } else {
-                setUpRules(prevRulesFromConnectionString.get());
+                setUpRules(prevStmtRulesFromConnString.get(), STATEMENT);
+                setUpRules(prevConnRulesFromConnString.get(), CONNECTION);
             }
         }
     }
@@ -133,8 +150,10 @@
      *         when an exception occurs
      */
     void setFromConnectionString(String newRules) throws SQLServerException {
-        prevRulesFromConnectionString.set(newRules);
-        setUpRules(prevRulesFromConnectionString.get());
+        prevStmtRulesFromConnString.set(newRules);
+        prevConnRulesFromConnString.set(newRules);
+        setUpRules(prevStmtRulesFromConnString.get(), STATEMENT);
+        setUpRules(prevConnRulesFromConnString.get(), CONNECTION);
     }
 
     /**
@@ -160,34 +179,42 @@
      * Sets up rules based on either connection string option or file read.
      *
      * @param cxnStrRules
-     *        if null, rules are constructed from file, else, this parameter is used to construct rules
+     *          if null, rules are constructed from file, else, this parameter is used to construct rules
+     * @param ruleType
+     *          either "statement" or "connection" for statement or connection rules respectively
      * @throws SQLServerException
-     *         if an exception occurs
+     *          if an exception occurs
      */
-    private static void setUpRules(String cxnStrRules) throws SQLServerException {
+    private static void setUpRules(String cxnStrRules, String ruleType) throws SQLServerException {
         LinkedList<String> temp;
-
-        stmtRules.set(new HashMap<>());
         lastQuery.set("");
 
         if (cxnStrRules == null || cxnStrRules.isEmpty()) {
-            temp = readFromFile();
+            if (ruleType.equals(STATEMENT)) {
+                temp = readFromFile(RETRY_EXEC);
+            } else {
+                temp = readFromFile(RETRY_CONN);
+            }
+
         } else {
             temp = new LinkedList<>();
             Collections.addAll(temp, cxnStrRules.split(SEMI_COLON));
         }
-        createRules(temp);
+        createRules(temp, ruleType);
     }
 
     /**
      * Creates and stores rules based on the inputted list of rules.
      *
      * @param listOfRules
      *        the list of rules, as a String LinkedList
+     * @param ruleType
+     *          the type of rule; either "statement" or "connection
      * @throws SQLServerException
      *         if unable to create rules from the inputted list
      */
-    private static void createRules(LinkedList<String> listOfRules) throws SQLServerException {
+    private static void createRules(LinkedList<String> listOfRules, String ruleType) throws SQLServerException {
+        connRules.set(new HashMap<>());
         stmtRules.set(new HashMap<>());
 
         for (String potentialRule : listOfRules) {
@@ -198,10 +225,29 @@
 
                 for (String retryError : arr) {
                     ConfigurableRetryRule splitRule = new ConfigurableRetryRule(retryError, rule);
-                    stmtRules.get().put(Integer.parseInt(splitRule.getError()), splitRule);
+                    if (rule.isConnection) {
+//                        if (rule.replaceExisting) {
+//                            if (!replaceFlag) {
+//                                connRules.set(new HashMap<>());
+//                            }
+//                            replaceFlag = true;
+//                        }
+                        connRules.get().put(Integer.parseInt(splitRule.getError()), splitRule);
+                    } else {
+                        stmtRules.get().put(Integer.parseInt(splitRule.getError()), splitRule);
+                    }
+
                 }
             } else {
-                stmtRules.get().put(Integer.parseInt(rule.getError()), rule);
+                if (rule.isConnection) {
+//                    if (rule.replaceExisting) {
+//                        connRules.set(new HashMap<>());
+//                        replaceFlag = true;
+//                    }
+                    connRules.get().put(Integer.parseInt(rule.getError()), rule);
+                } else {
+                    stmtRules.get().put(Integer.parseInt(rule.getError()), rule);
+                }
             }
         }
     }
@@ -241,7 +287,7 @@
      * @throws SQLServerException
      *         if unable to read from the file
      */
-    private static LinkedList<String> readFromFile() throws SQLServerException {
+    private static LinkedList<String> readFromFile(String connectionStringProperty) throws SQLServerException {
         String filePath = getCurrentClassPath();
         LinkedList<String> list = new LinkedList<>();
 
@@ -250,7 +296,7 @@
             try (BufferedReader buffer = new BufferedReader(new FileReader(f))) {
                 String readLine;
                 while ((readLine = buffer.readLine()) != null) {
-                    if (readLine.startsWith(RETRY_EXEC)) {
+                    if (readLine.startsWith(connectionStringProperty)) { // Either "retryExec" or "retryConn"
                         String value = readLine.split(EQUALS_SIGN)[1];
                         Collections.addAll(list, value.split(SEMI_COLON));
                     }
@@ -280,13 +326,25 @@
      * @throws SQLServerException
      *         when an exception occurs
      */
-    ConfigurableRetryRule searchRuleSet(int ruleToSearchFor) throws SQLServerException {
+    ConfigurableRetryRule searchRuleSet(int ruleToSearchFor, String ruleSet) throws SQLServerException {
         refreshRuleSet();
-        for (Map.Entry<Integer, ConfigurableRetryRule> entry : stmtRules.get().entrySet()) {
-            if (entry.getKey() == ruleToSearchFor) {
-                return entry.getValue();
+        if (ruleSet.equals(STATEMENT)) {
+            for (Map.Entry<Integer, ConfigurableRetryRule> entry : stmtRules.get().entrySet()) {
+                if (entry.getKey() == ruleToSearchFor) {
+                    return entry.getValue();
+                }
+            }
+        } else {
+            for (Map.Entry<Integer, ConfigurableRetryRule> entry : connRules.get().entrySet()) {
+                if (entry.getKey() == ruleToSearchFor) {
+                    return entry.getValue();
+                }
             }
         }
         return null;
     }
+
+    boolean getReplaceFlag() {
+        return replaceFlag;
+    }
 }

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

@@ -25,6 +25,8 @@
     private int retryCount = 1;
     private String retryQueries = "";
     private String retryError;
+    boolean isConnection = false;
+    boolean replaceExisting = false;
 
     private ArrayList<Integer> waitTimes = new ArrayList<>();
 
@@ -70,6 +72,18 @@
         this.retryCount = baseRule.retryCount;
         this.retryQueries = baseRule.retryQueries;
         this.waitTimes = baseRule.waitTimes;
+        this.isConnection = baseRule.isConnection;
+    }
+
+    private String appendOrReplace(String retryError) {
+        if (retryError.charAt(0) == '+') {
+            replaceExisting = false;
+            StringUtils.isNumeric(retryError.substring(1));
+            return retryError.substring(1);
+        } else {
+            replaceExisting = true;
+            return retryError;
+        }
     }
 
     /**
@@ -152,7 +166,12 @@
      *         if a rule or parameter has invalid inputs
      */
     private void addElements(String[] rule) throws SQLServerException {
-        if (rule.length == 2 || rule.length == 3) {
+        if (rule.length == 1) {
+            String errorWithoutOptionalPrefix = appendOrReplace(rule[0]);
+            checkParameter(errorWithoutOptionalPrefix);
+            isConnection = true;
+            retryError = errorWithoutOptionalPrefix;
+        } else if (rule.length == 2 || rule.length == 3) {
             checkParameter(rule[0]);
             retryError = rule[0];
             String[] timings = rule[1].split(COMMA);

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

@@ -1363,6 +1363,22 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
      */
     String getRetryExec();
 
+    /**
+     * Returns value of 'retryConn' from Connection String.
+     *
+     * @param retryConn
+     *          Set of rules used for connection retry
+     */
+    void setRetryConn(String retryConn);
+
+    /**
+     * Sets the value for 'retryConn' property
+     *
+     * @return retryConn
+     *         String value
+     */
+    String getRetryConn();
+
     /**
      * useFlexibleCallableStatements is temporarily removed. This is meant as a no-op.
      *

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

@@ -1086,6 +1086,28 @@
         this.retryExec = retryExec;
     }
 
+    private String retryConn = SQLServerDriverStringProperty.RETRY_CONN.getDefaultValue();
+
+    /**
+     * Returns the set of configurable connection retry rules set in retryConn
+     *
+     * @return
+     *         A string containing statement retry rules.
+     */
+    public String getRetryConn() {
+        return retryConn;
+    }
+
+    /**
+     * Sets the list of configurable connection retry rules, for the given connection, in retryConn.
+     *
+     * @param retryConn
+     *        The list of retry rules to set, as a string.
+     */
+    public void setRetryConn(String retryConn) {
+        this.retryConn = retryConn;
+    }
+
     /** Session Recovery Object */
     private transient IdleConnectionResiliency sessionRecovery = new IdleConnectionResiliency(this);
 
@@ -2024,10 +2046,23 @@
                         }
                         throw e;
                     } else {
-                        // only retry if transient error
+                        // Only retry if matches configured CRL rules, or transient error (if CRL is not in use)
                         SQLServerError sqlServerError = e.getSQLServerError();
-                        if (!TransientError.isTransientError(sqlServerError)) {
+                        if (null == sqlServerError) {
                             throw e;
+                        } else {
+                            ConfigurableRetryRule rule = ConfigurableRetryLogic.getInstance()
+                                    .searchRuleSet(sqlServerError.getErrorNumber(), "connection");
+
+                            if (null == rule) {
+                                if (ConfigurableRetryLogic.getInstance().getReplaceFlag()) {
+                                    throw e;
+                                } else {
+                                    if (!TransientError.isTransientError(sqlServerError)) {
+                                        throw e;
+                                    }
+                                }
+                            }
                         }
 
                         // check if there's time to retry, no point to wait if no time left

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

@@ -1383,16 +1383,48 @@ public boolean getCalcBigDecimalPrecision() {
                 SQLServerDriverBooleanProperty.CALC_BIG_DECIMAL_PRECISION.getDefaultValue());
     }
 
+    /**
+     * Sets the 'retryExec' setting.
+     *
+     * @param retryExec
+     *        String property giving the custom statement retry rules to use for configurable retry logic
+     */
     @Override
     public void setRetryExec(String retryExec) {
         setStringProperty(connectionProps, SQLServerDriverStringProperty.RETRY_EXEC.toString(), retryExec);
     }
 
+    /**
+     * Returns the value for 'retryExec'.
+     *
+     * @return retryExec String value
+     */
     @Override
     public String getRetryExec() {
         return getStringProperty(connectionProps, SQLServerDriverStringProperty.RETRY_EXEC.toString(), null);
     }
 
+    /**
+     * Sets the 'retryConn' setting.
+     *
+     * @param retryConn
+     *        String property giving the custom connection retry rules to use for configurable retry logic
+     */
+    @Override
+    public void setRetryConn(String retryConn) {
+        setStringProperty(connectionProps, SQLServerDriverStringProperty.RETRY_CONN.toString(), retryConn);
+    }
+
+    /**
+     * Returns the value for 'retryConn'.
+     *
+     * @return retryConn String value
+     */
+    @Override
+    public String getRetryConn() {
+        return getStringProperty(connectionProps, SQLServerDriverStringProperty.RETRY_CONN.toString(), null);
+    }
+
     /**
      * Sets a property string value.
      *