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.
This commit is contained in:
Gian Merlino 2022-10-25 09:31:56 -07:00 committed by GitHub
parent 1e39bc65cc
commit 2b0d873c7e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 18 additions and 6 deletions

View File

@ -210,7 +210,8 @@ public class DirectStatement extends AbstractStatement implements Cancelable
// Context keys for authorization. Use the user-provided keys, // Context keys for authorization. Use the user-provided keys,
// NOT the keys from the query context which, by this point, // NOT the keys from the query context which, by this point,
// will have been extended with internally-defined values. // will have been extended with internally-defined values.
queryPlus.context().keySet())) { queryPlus.context().keySet()
)) {
validate(planner); validate(planner);
authorize(planner, authorizer()); authorize(planner, authorizer());
@ -314,4 +315,11 @@ public class DirectStatement extends AbstractStatement implements Cancelable
state = State.CLOSED; state = State.CLOSED;
} }
} }
@Override
public void closeQuietly()
{
sqlToolbox.sqlLifecycleManager.remove(sqlQueryId(), this);
super.closeQuietly();
}
} }

View File

@ -411,7 +411,8 @@ public class DruidJdbcResultSet implements Closeable
throw new RuntimeException(t); throw new RuntimeException(t);
} }
finally { 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(); stmt.close();
} }
} }

View File

@ -55,6 +55,7 @@ import org.apache.druid.sql.SqlLifecycleManager.Cancelable;
import org.apache.druid.sql.SqlPlanningException; import org.apache.druid.sql.SqlPlanningException;
import org.apache.druid.sql.SqlRowTransformer; import org.apache.druid.sql.SqlRowTransformer;
import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlStatementFactory;
import org.apache.druid.utils.CloseableUtils;
import javax.annotation.Nullable; import javax.annotation.Nullable;
import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletRequest;
@ -174,8 +175,12 @@ public class SqlResource
throw new RuntimeException(ex); throw new RuntimeException(ex);
} }
finally { finally {
yielder.close(); final Exception finalE = e;
endLifecycle(stmt, e, os.getCount());
CloseableUtils.closeAll(
yielder,
() -> endLifecycle(stmt, finalE, os.getCount())
);
} }
} }
} }
@ -245,7 +250,6 @@ public class SqlResource
HttpStatement stmt HttpStatement stmt
) )
{ {
sqlLifecycleManager.remove(stmt.sqlQueryId(), stmt);
stmt.closeQuietly(); stmt.closeQuietly();
} }
@ -261,7 +265,6 @@ public class SqlResource
} else { } else {
reporter.failed(e); reporter.failed(e);
} }
sqlLifecycleManager.remove(stmt.sqlQueryId(), stmt);
stmt.close(); stmt.close();
} }