Replace usage of NeverInline in jdbc-pgsql

extensions/jdbc/jdbc-postgresql/runtime/src/main/java/io/quarkus/jdbc/postgresql/runtime/graal/DomHelper.java

@@ -1,8 +1,8 @@
 package io.quarkus.jdbc.postgresql.runtime.graal;
 
 import java.io.StringWriter;
-import java.lang.reflect.InvocationTargetException;
 import java.sql.SQLException;
+import java.util.function.BiFunction;
 
 import javax.xml.transform.Transformer;
 import javax.xml.transform.TransformerException;
@@ -11,50 +11,52 @@
 import javax.xml.transform.dom.DOMSource;
 import javax.xml.transform.stream.StreamResult;
 
+import org.graalvm.nativeimage.ImageSingletons;
 import org.postgresql.core.BaseConnection;
 import org.postgresql.util.GT;
 import org.postgresql.util.PSQLException;
 import org.postgresql.util.PSQLState;
 import org.postgresql.xml.DefaultPGXmlFactoryFactory;
 import org.postgresql.xml.PGXmlFactoryFactory;
 
-import com.oracle.svm.core.annotate.NeverInline;
-
 /**
  * Used by PgSQLXML: easier to keep the actual code separated from the substitutions.
  */
-public final class DomHelper {
+final class DomHelper {
 
-    //This is the actual code reaching to the JDK XML types; no other code except
-    //the reflective call in the previous method should trigger this inclusion.
-    public static String processDomResult(DOMResult domResult, BaseConnection conn) throws SQLException {
-        return maybeProcessAsDomResult(domResult, conn);
-    }
+    // Stays null unless XML result processing becomes reachable.
+    private BiFunction<DOMResult, BaseConnection, String> processDomResult;
 
     // This is only ever invoked if field domResult got initialized; which in turn
     // is only possible if setResult(Class) is reachable.
-    public static String maybeProcessAsDomResult(DOMResult domResult, BaseConnection conn) throws SQLException {
-        try {
-            return (String) DomHelper.class.getMethod(obfuscatedMethodName(), DOMResult.class, BaseConnection.class)
-                    .invoke(null, domResult, conn);
-        } catch (NoSuchMethodException | IllegalAccessException e) {
-            throw new RuntimeException("Unexpected failure in reflective call - please report", e);
-        } catch (InvocationTargetException e) {
-            final Throwable cause = e.getCause();
-            if (cause instanceof SQLException) {
-                //Propagate normal SQLException(s):
-                throw (SQLException) cause;
-            } else {
-                throw new RuntimeException("Unexpected failure in reflective call - please report", e);
+    public static String processDomResult(DOMResult domResult, BaseConnection conn) throws SQLException {
+        BiFunction<DOMResult, BaseConnection, String> func = ImageSingletons.lookup(DomHelper.class).processDomResult;
+        if (func == null) {
+            return null;
+        } else {
+            try {
+                return func.apply(domResult, conn);
+            } catch (UncheckedSQLException e) {
+                throw (SQLException) e.getCause();
             }
         }
     }
 
-    @NeverInline("Prevent GraalVM from figuring out the target method to be invoked reflectively, so that it can't automatically register it for reflection")
-    private static String obfuscatedMethodName() {
-        return "reallyProcessDomResult";
+    // Called by SQLXMLFeature when setResult(Class) becomes reachable
+    static void enableXmlProcessing() {
+        ImageSingletons.lookup(DomHelper.class).processDomResult = DomHelper::doProcessDomResult;
     }
 
+    public static String doProcessDomResult(DOMResult domResult, BaseConnection conn) {
+        try {
+            return reallyProcessDomResult(domResult, conn);
+        } catch (SQLException e) {
+            throw new UncheckedSQLException(e);
+        }
+    }
+
+    // This XML processing code should only be reachable (via processDomResult function) iff it's
+    // actually used. I.e. setResult(Class) is reachable.
     public static String reallyProcessDomResult(DOMResult domResult, BaseConnection conn) throws SQLException {
         TransformerFactory factory = getXmlFactoryFactory(conn).newTransformerFactory();
         try {
@@ -77,3 +79,11 @@ private static PGXmlFactoryFactory getXmlFactoryFactory(BaseConnection conn) thr
         return DefaultPGXmlFactoryFactory.INSTANCE;
     }
 }
+
+@SuppressWarnings("serial")
+final class UncheckedSQLException extends RuntimeException {
+
+    UncheckedSQLException(SQLException cause) {
+        super(cause);
+    }
+}

extensions/jdbc/jdbc-postgresql/runtime/src/main/java/io/quarkus/jdbc/postgresql/runtime/graal/SQLXMLFeature.java

@@ -3,11 +3,8 @@
 import java.lang.reflect.Method;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-import javax.xml.transform.dom.DOMResult;
-
+import org.graalvm.nativeimage.ImageSingletons;
 import org.graalvm.nativeimage.hosted.Feature;
-import org.graalvm.nativeimage.hosted.RuntimeReflection;
-import org.postgresql.core.BaseConnection;
 
 public final class SQLXMLFeature implements Feature {
 
@@ -18,6 +15,11 @@ public final class SQLXMLFeature implements Feature {
      */
     private static final boolean log = Boolean.getBoolean("io.quarkus.jdbc.postgresql.graalvm.diagnostics");
 
+    @Override
+    public void afterRegistration(AfterRegistrationAccess access) {
+        ImageSingletons.add(DomHelper.class, new DomHelper());
+    }
+
     @Override
     public void beforeAnalysis(BeforeAnalysisAccess access) {
         Class<?> pgSQLXMLClass = access.findClassByName("io.quarkus.jdbc.postgresql.runtime.graal.PgSQLXML");
@@ -39,17 +41,7 @@ private void identifiedXMLProcessingInDriver(DuringAnalysisAccess duringAnalysis
                 System.out.println(
                         "Quarkus' automatic feature for GraalVM native images: enabling support for XML processing in the PostgreSQL JDBC driver");
             }
-            enableDomXMLProcessingInDriver(duringAnalysisAccess);
-        }
-    }
-
-    private void enableDomXMLProcessingInDriver(DuringAnalysisAccess duringAnalysisAccess) {
-        final Class<?> classByName = duringAnalysisAccess.findClassByName("io.quarkus.jdbc.postgresql.runtime.graal.DomHelper");
-        try {
-            final Method method = classByName.getMethod("reallyProcessDomResult", DOMResult.class, BaseConnection.class);
-            RuntimeReflection.register(method);
-        } catch (NoSuchMethodException e) {
-            throw new RuntimeException(e);
+            DomHelper.enableXmlProcessing();
         }
     }