Configurable Retry Logic II - Connection Retry

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

@@ -39,6 +39,9 @@ public class ConfigurableRetryLogic {
     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 boolean replaceFlag = false; // Are we replacing the list of transient errors?
     /**
      * The time the properties file was last modified.
      */
@@ -52,14 +55,23 @@ public class ConfigurableRetryLogic {
      */
     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 +82,8 @@ public class ConfigurableRetryLogic {
      */
     private ConfigurableRetryLogic() throws SQLServerException {
         timeLastRead.compareAndSet(0, new Date().getTime());
-        setUpRules(null);
+        setUpStatementRules(null);
+        setUpConnectionRules(null);
     }
 
     /**
@@ -102,7 +115,8 @@ public static ConfigurableRetryLogic getInstance() throws SQLServerException {
 
     /**
      * 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,25 +130,40 @@ private static void refreshRuleSet() throws SQLServerException {
                 // 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);
+                    setUpStatementRules(null);
+                    setUpConnectionRules(null);
                 }
             } else {
-                setUpRules(prevRulesFromConnectionString.get());
+                setUpStatementRules(prevStmtRulesFromConnString.get());
+                setUpConnectionRules(prevConnRulesFromConnString.get());
             }
         }
     }
 
     /**
-     * Sets rules given from connection string.
+     * Sets statement rules given from connection string.
+     *
+     * @param newRules
+     *        the new rules to use
+     * @throws SQLServerException
+     *         when an exception occurs
+     */
+    void setStatementRulesFromConnectionString(String newRules) throws SQLServerException {
+        prevStmtRulesFromConnString.set(newRules);
+        setUpStatementRules(prevStmtRulesFromConnString.get());
+    }
+
+    /**
+     * Sets connection rules given from connection string.
      *
      * @param newRules
      *        the new rules to use
      * @throws SQLServerException
      *         when an exception occurs
      */
-    void setFromConnectionString(String newRules) throws SQLServerException {
-        prevRulesFromConnectionString.set(newRules);
-        setUpRules(prevRulesFromConnectionString.get());
+    void setConnectionRulesFromConnectionString(String newRules) throws SQLServerException {
+        prevConnRulesFromConnString.set(newRules);
+        setUpConnectionRules(prevConnRulesFromConnString.get());
     }
 
     /**
@@ -164,19 +193,34 @@ String getLastQuery() {
      * @throws SQLServerException
      *         if an exception occurs
      */
-    private static void setUpRules(String cxnStrRules) throws SQLServerException {
+    private static void setUpStatementRules(String cxnStrRules) throws SQLServerException {
         LinkedList<String> temp;
 
         stmtRules.set(new HashMap<>());
         lastQuery.set("");
 
         if (cxnStrRules == null || cxnStrRules.isEmpty()) {
-            temp = readFromFile();
+            temp = readFromFile(RETRY_EXEC);
         } else {
             temp = new LinkedList<>();
             Collections.addAll(temp, cxnStrRules.split(SEMI_COLON));
         }
-        createRules(temp);
+        createStatementRules(temp);
+    }
+
+    private static void setUpConnectionRules(String cxnStrRules) throws SQLServerException {
+        LinkedList<String> temp;
+
+        connRules.set(new HashMap<>());
+        lastQuery.set("");
+
+        if (cxnStrRules == null || cxnStrRules.isEmpty()) {
+            temp = readFromFile(RETRY_CONN);
+        } else {
+            temp = new LinkedList<>();
+            Collections.addAll(temp, cxnStrRules.split(SEMI_COLON));
+        }
+        createConnectionRules(temp);
     }
 
     /**
@@ -187,7 +231,7 @@ private static void setUpRules(String cxnStrRules) throws SQLServerException {
      * @throws SQLServerException
      *         if unable to create rules from the inputted list
      */
-    private static void createRules(LinkedList<String> listOfRules) throws SQLServerException {
+    private static void createStatementRules(LinkedList<String> listOfRules) throws SQLServerException {
         stmtRules.set(new HashMap<>());
 
         for (String potentialRule : listOfRules) {
@@ -206,6 +250,29 @@ private static void createRules(LinkedList<String> listOfRules) throws SQLServer
         }
     }
 
+    private static void createConnectionRules(LinkedList<String> listOfRules) throws SQLServerException {
+        connRules.set(new HashMap<>());
+        replaceFlag = false;
+
+        for (String potentialRule : listOfRules) {
+            ConfigurableRetryRule rule = new ConfigurableRetryRule(potentialRule);
+            if (rule.replaceExisting) {
+                replaceFlag = true;
+            }
+
+            if (rule.getError().contains(COMMA)) {
+                String[] arr = rule.getError().split(COMMA);
+
+                for (String retryError : arr) {
+                    ConfigurableRetryRule splitRule = new ConfigurableRetryRule(retryError, rule);
+                    connRules.get().put(Integer.parseInt(splitRule.getError()), splitRule);
+                }
+            } else {
+                connRules.get().put(Integer.parseInt(rule.getError()), rule);
+            }
+        }
+    }
+
     /**
      * Gets the current class path (for use in file reading).
      *
@@ -241,7 +308,7 @@ private static String getCurrentClassPath() throws SQLServerException {
      * @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 +317,7 @@ private static LinkedList<String> readFromFile() throws SQLServerException {
             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 +347,25 @@ private static LinkedList<String> readFromFile() throws SQLServerException {
      * @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 @@ class ConfigurableRetryRule {
     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 @@ private void copyFromRule(ConfigurableRetryRule baseRule) {
         this.retryCount = baseRule.retryCount;
         this.retryQueries = baseRule.retryQueries;
         this.waitTimes = baseRule.waitTimes;
+        this.isConnection = baseRule.isConnection;
+    }
+
+    private String appendOrReplace(String retryError) {
+        if (retryError.startsWith(PLUS_SIGN)) {
+            replaceExisting = false;
+            StringUtils.isNumeric(retryError.substring(1));
+            return retryError.substring(1);
+        } else {
+            replaceExisting = true;
+            return retryError;
+        }
     }
 
     /**
@@ -152,7 +166,12 @@ private void checkParameter(String value) throws SQLServerException {
      *         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 @@ public void setRetryExec(String retryExec) {
         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 @@ Connection connect(Properties propsIn, SQLServerPooledConnection pooledConnectio
                         }
                         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
@@ -2367,7 +2402,16 @@ Connection connectInternal(Properties propsIn,
                     activeConnectionProperties.setProperty(sPropKey, sPropValue);
                 }
                 retryExec = sPropValue;
-                ConfigurableRetryLogic.getInstance().setFromConnectionString(sPropValue);
+                ConfigurableRetryLogic.getInstance().setStatementRulesFromConnectionString(sPropValue);
+
+                sPropKey = SQLServerDriverStringProperty.RETRY_CONN.toString();
+                sPropValue = activeConnectionProperties.getProperty(sPropKey);
+                if (null == sPropValue) {
+                    sPropValue = SQLServerDriverStringProperty.RETRY_CONN.getDefaultValue();
+                    activeConnectionProperties.setProperty(sPropKey, sPropValue);
+                }
+                retryConn = sPropValue;
+                ConfigurableRetryLogic.getInstance().setConnectionRulesFromConnectionString(sPropValue);
 
                 sPropKey = SQLServerDriverBooleanProperty.CALC_BIG_DECIMAL_PRECISION.toString();
                 sPropValue = activeConnectionProperties.getProperty(sPropKey);