Closing connection to return connection back to the connection pool i…

…n case of any exceptions during read/write
Netflix · Jan 28, 2019 · 094fb50 · 094fb50
1 parent 63882aa
commit 094fb50
Show file tree

Hide file tree

Showing 3 changed files with 233 additions and 127 deletions.
diff --git a/...va/com/netflix/ndbench/plugin/cockroachdb/operations/CockroachDBSecondaryIndexPlugin.java b/...va/com/netflix/ndbench/plugin/cockroachdb/operations/CockroachDBSecondaryIndexPlugin.java
@@ -45,59 +45,95 @@ public CockroachDBSecondaryIndexPlugin(CockroachDBConfiguration cockroachDBConfi
     @Override
     public String readSingle(String key) throws Exception
     {
-        Connection connection = ds.getConnection();
-        ResultSet rs = connection.createStatement().executeQuery(readFromMainQuery + "'" + key + "'");
-        int rsSize = 0;
-        while (rs.next())
+        Connection connection = null;
+        try
         {
-            rsSize++;
-        }
+            connection = ds.getConnection();
+            ResultSet rs = connection.createStatement().executeQuery(readFromMainQuery + "'" + key + "'");
+            int rsSize = 0;
+            while (rs.next())
+            {
+                rsSize++;
+            }
+
+            if (rsSize == 0)
+            {
+                connection.close();
+                return CacheMiss;
+            }
+
+            if (rsSize > 1)
+            {
+                connection.close();
+                throw new Exception("Expecting only 1 row with a given key: " + key);
+            }
 
-        if (rsSize == 0)
-        {
             connection.close();
-            return CacheMiss;
+            return ResultOK;
         }
-
-        if (rsSize > 1)
+        catch (Exception ex)
         {
-            connection.close();
-            throw new Exception("Expecting only 1 row with a given key: " + key);
+            if (connection != null)
+            {
+                connection.close();
+            }
+            throw ex;
         }
-
-        connection.close();
-        return ResultOK;
     }
 
     @Override
     public String writeSingle(String key) throws Exception
     {
-        String columns = getNDelimitedStrings(config.getColsPerRow());
-        Connection connection = ds.getConnection();
-        connection
-        .createStatement()
-        .executeUpdate(writeToMainQuery + "('" + key + "', " + columns + ")");
-        connection.close();
-        return ResultOK;
+        Connection connection = null;
+        try
+        {
+            String columns = getNDelimitedStrings(config.getColsPerRow());
+            connection = ds.getConnection();
+            connection
+            .createStatement()
+            .executeUpdate(writeToMainQuery + "('" + key + "', " + columns + ")");
+            connection.close();
+            return ResultOK;
+        }
+        catch (Exception ex)
+        {
+            if (connection != null)
+            {
+                connection.close();
+            }
+            throw ex;
+        }
     }
 
     public void createTables() throws Exception
     {
-        Connection connection = ds.getConnection();
-        String columns = IntStream.range(0, config.getColsPerRow()).mapToObj(i -> "column" + i + " STRING").collect(Collectors.joining(", "));
-        connection
-        .createStatement()
-        .execute(String.format("CREATE TABLE IF NOT EXISTS %s.%s (key STRING PRIMARY KEY, %s)", config.getDBName(), config.getTableName(), columns));
-
-        // create secondary indices
-        for (int i = 0; i < config.getColsPerRow(); i++)
+        Connection connection = null;
+        try
         {
+            connection = ds.getConnection();
+            String columns = IntStream.range(0, config.getColsPerRow()).mapToObj(i -> "column" + i + " STRING").collect(Collectors.joining(", "));
             connection
             .createStatement()
-            .execute(String.format("CREATE INDEX IF NOT EXISTS %s_column%d_index on %s (column%d)", config.getTableName(), i, config.getTableName(), i));
-        }
+            .execute(String.format("CREATE TABLE IF NOT EXISTS %s.%s (key STRING PRIMARY KEY, %s)", config.getDBName(), config.getTableName(), columns));
 
-        connection.close();
+            // create secondary indices
+            for (int i = 0; i < config.getColsPerRow(); i++)
+            {
+                connection
+                .createStatement()
+                .execute(String.format("CREATE INDEX IF NOT EXISTS %s_column%d_index on %s (column%d)", config.getTableName(), i, config.getTableName(), i));
+            }
+
+            connection.close();
+        }
+        catch (Exception ex)
+        {
+            if (connection != null)
+            {
+                connection.close();
+            }
+            throw ex;
+        }
     }
 
     public void prepareStatements()

diff --git a/.../main/java/com/netflix/ndbench/plugin/cockroachdb/operations/CockroachDBSimplePlugin.java b/.../main/java/com/netflix/ndbench/plugin/cockroachdb/operations/CockroachDBSimplePlugin.java
@@ -44,55 +44,91 @@ public CockroachDBSimplePlugin(CockroachDBConfiguration cockroachDBConfiguration
     @Override
     public String readSingle(String key) throws Exception
     {
-        Connection connection = ds.getConnection();
-        ResultSet rs = connection.createStatement().executeQuery(readQuery + "'" + key + "'");
-        int rsSize = 0;
-        while (rs.next())
+        Connection connection = null;
+        try
         {
-            rsSize++;
-        }
+            connection = ds.getConnection();
+            ResultSet rs = connection.createStatement().executeQuery(readQuery + "'" + key + "'");
+            int rsSize = 0;
+            while (rs.next())
+            {
+                rsSize++;
+            }
+
+            if (rsSize == 0)
+            {
+                connection.close();
+                return CacheMiss;
+            }
+
+            if (rsSize > 1)
+            {
+                connection.close();
+                throw new Exception("Expecting only 1 row with a given key: " + key);
+            }
 
-        if (rsSize == 0)
-        {
             connection.close();
-            return CacheMiss;
+            return ResultOK;
         }
-
-        if (rsSize > 1)
+        catch (Exception ex)
         {
-            connection.close();
-            throw new Exception("Expecting only 1 row with a given key: " + key);
+            if (connection != null)
+            {
+                connection.close();
+            }
+            throw ex;
         }
-
-        connection.close();
-        return ResultOK;
     }
 
     @Override
     public String writeSingle(String key) throws Exception
     {
-        String values = getNDelimitedStrings(config.getColsPerRow());
-        Connection connection = ds.getConnection();
+        Connection connection = null;
+        try
+        {
+            String values = getNDelimitedStrings(config.getColsPerRow());
+            connection = ds.getConnection();
 
-        connection
-        .createStatement()
-        .executeUpdate(writeQuery + "('" + key + "', 1 ," + values + ")");
-        connection.close();
+            connection
+            .createStatement()
+            .executeUpdate(writeQuery + "('" + key + "', 1 ," + values + ")");
+            connection.close();
 
-        return ResultOK;
+            return ResultOK;
+        }
+        catch (Exception ex)
+        {
+            if (connection != null)
+            {
+                connection.close();
+            }
+            throw ex;
+        }
     }
 
     public void createTables() throws Exception
     {
-        String values = IntStream.range(0, config.getColsPerRow()).mapToObj(i -> "value" + i + " STRING").collect(Collectors.joining(", "));
+        Connection connection = null;
+        try
+        {
+            String values = IntStream.range(0, config.getColsPerRow()).mapToObj(i -> "value" + i + " STRING").collect(Collectors.joining(", "));
 
-        Connection connection = ds.getConnection();
+            connection = ds.getConnection();
 
-        connection
-        .createStatement()
-        .execute(String.format("CREATE TABLE IF NOT EXISTS %s.%s (key STRING PRIMARY KEY, column1 INT, %s)", config.getDBName(), config.getTableName(), values));
+            connection
+            .createStatement()
+            .execute(String.format("CREATE TABLE IF NOT EXISTS %s.%s (key STRING PRIMARY KEY, column1 INT, %s)", config.getDBName(), config.getTableName(), values));
 
-        connection.close();
+            connection.close();
+        }
+        catch (Exception ex)
+        {
+            if (connection != null)
+            {
+                connection.close();
+            }
+            throw ex;
+        }
     }
 
     public void prepareStatements()