[druid] branch master updated: Fix two sources of SQL statement leaks. (#13259)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

gian pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git


The following commit(s) were added to refs/heads/master by this push:
     new 2b0d873c7e Fix two sources of SQL statement leaks. (#13259)
2b0d873c7e is described below

commit 2b0d873c7e4bf244fce12c45e381d00a61a277e2
Author: Gian Merlino <gi...@gmail.com>
AuthorDate: Tue Oct 25 09:31:56 2022 -0700

    Fix two sources of SQL statement leaks. (#13259)
    
    * Fix two sources of SQL statement leaks.
    
    1) SqlTaskResource and DruidJdbcResultSet leaked statements 100% of the
       time, since they call stmt.plan(), which adds statements to
       SqlLifecycleManager, and they do not explicitly remove them.
    
    2) SqlResource leaked statements if yielder.close() threw an exception.
       (And also would not emit metrics, since in that case it failed to
       call stmt.close as well.)
    
    * Only closeQuietly is needed.
---
 sql/src/main/java/org/apache/druid/sql/DirectStatement.java   | 10 +++++++++-
 .../java/org/apache/druid/sql/avatica/DruidJdbcResultSet.java |  3 ++-
 sql/src/main/java/org/apache/druid/sql/http/SqlResource.java  | 11 +++++++----
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
index 776dc30916..73b57c0719 100644
--- a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
+++ b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
@@ -210,7 +210,8 @@ public class DirectStatement extends AbstractStatement implements Cancelable
         // Context keys for authorization. Use the user-provided keys,
         // NOT the keys from the query context which, by this point,
         // will have been extended with internally-defined values.
-        queryPlus.context().keySet())) {
+        queryPlus.context().keySet()
+    )) {
       validate(planner);
       authorize(planner, authorizer());
 
@@ -314,4 +315,11 @@ public class DirectStatement extends AbstractStatement implements Cancelable
       state = State.CLOSED;
     }
   }
+
+  @Override
+  public void closeQuietly()
+  {
+    sqlToolbox.sqlLifecycleManager.remove(sqlQueryId(), this);
+    super.closeQuietly();
+  }
 }
diff --git a/sql/src/main/java/org/apache/druid/sql/avatica/DruidJdbcResultSet.java b/sql/src/main/java/org/apache/druid/sql/avatica/DruidJdbcResultSet.java
index 1eb0d1aa5e..70ec88ff1b 100644
--- a/sql/src/main/java/org/apache/druid/sql/avatica/DruidJdbcResultSet.java
+++ b/sql/src/main/java/org/apache/druid/sql/avatica/DruidJdbcResultSet.java
@@ -411,7 +411,8 @@ public class DruidJdbcResultSet implements Closeable
       throw new RuntimeException(t);
     }
     finally {
-      // Closing the statement cause logs and metrics to be emitted.
+      // Closing the statement causes logs and metrics to be emitted, and causes the statement to be removed
+      // from the SqlLifecycleManager.
       stmt.close();
     }
   }
diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java
index dc4bfbcfb6..968a722caa 100644
--- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java
+++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java
@@ -55,6 +55,7 @@ import org.apache.druid.sql.SqlLifecycleManager.Cancelable;
 import org.apache.druid.sql.SqlPlanningException;
 import org.apache.druid.sql.SqlRowTransformer;
 import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.utils.CloseableUtils;
 
 import javax.annotation.Nullable;
 import javax.servlet.http.HttpServletRequest;
@@ -174,8 +175,12 @@ public class SqlResource
                     throw new RuntimeException(ex);
                   }
                   finally {
-                    yielder.close();
-                    endLifecycle(stmt, e, os.getCount());
+                    final Exception finalE = e;
+
+                    CloseableUtils.closeAll(
+                        yielder,
+                        () -> endLifecycle(stmt, finalE, os.getCount())
+                    );
                   }
                 }
               }
@@ -245,7 +250,6 @@ public class SqlResource
       HttpStatement stmt
   )
   {
-    sqlLifecycleManager.remove(stmt.sqlQueryId(), stmt);
     stmt.closeQuietly();
   }
 
@@ -261,7 +265,6 @@ public class SqlResource
     } else {
       reporter.failed(e);
     }
-    sqlLifecycleManager.remove(stmt.sqlQueryId(), stmt);
     stmt.close();
   }
 


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org