From 8b774e542f9b32872b8162dd06519bf08976a1ac Mon Sep 17 00:00:00 2001
From: Zheng Feng <zh.feng@gmail.com>
Date: Wed, 31 May 2023 19:42:00 +0800
Subject: [PATCH] Defer register XAResourceRecovery after recovery manager
 service is created

---
 .../jta/runtime/QuarkusRecoveryService.java   | 41 ++++++++++++++++++-
 .../jta/JdbcObjectStoreTestProfile.java       |  5 +++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/extensions/narayana-jta/runtime/src/main/java/io/quarkus/narayana/jta/runtime/QuarkusRecoveryService.java b/extensions/narayana-jta/runtime/src/main/java/io/quarkus/narayana/jta/runtime/QuarkusRecoveryService.java
index 18c5e80438c10..55ff7ae05fe22 100644
--- a/extensions/narayana-jta/runtime/src/main/java/io/quarkus/narayana/jta/runtime/QuarkusRecoveryService.java
+++ b/extensions/narayana-jta/runtime/src/main/java/io/quarkus/narayana/jta/runtime/QuarkusRecoveryService.java
@@ -1,17 +1,54 @@
 package io.quarkus.narayana.jta.runtime;
 
+import java.util.ArrayList;
+import java.util.List;
+
+import org.jboss.tm.XAResourceRecovery;
+
 import com.arjuna.ats.jbossatx.jta.RecoveryManagerService;
 
-public class QuarkusRecoveryService {
+public class QuarkusRecoveryService extends RecoveryManagerService {
     private static RecoveryManagerService recoveryManagerService;
+    private List<XAResourceRecovery> xaResources;
+    private boolean isCreated;
 
     public static RecoveryManagerService getInstance() {
         if (recoveryManagerService == null) {
-            recoveryManagerService = new RecoveryManagerService();
+            recoveryManagerService = new QuarkusRecoveryService();
         }
         return recoveryManagerService;
     }
 
     private QuarkusRecoveryService() {
+        xaResources = new ArrayList<>();
+        isCreated = false;
+    }
+
+    @Override
+    public void addXAResourceRecovery(XAResourceRecovery xares) {
+        if (isCreated) {
+            super.addXAResourceRecovery(xares);
+        } else {
+            xaResources.add(xares);
+        }
+    }
+
+    @Override
+    public void removeXAResourceRecovery(XAResourceRecovery xares) {
+        if (isCreated) {
+            super.removeXAResourceRecovery(xares);
+        } else {
+            xaResources.remove(xares);
+        }
+    }
+
+    @Override
+    public void create() {
+        super.create();
+        isCreated = true;
+        for (XAResourceRecovery xares : xaResources) {
+            super.addXAResourceRecovery(xares);
+        }
+        xaResources.clear();
     }
 }
diff --git a/integration-tests/narayana-jta/src/test/java/io/quarkus/narayana/jta/JdbcObjectStoreTestProfile.java b/integration-tests/narayana-jta/src/test/java/io/quarkus/narayana/jta/JdbcObjectStoreTestProfile.java
index 88fa090fa3a3b..f4aecafd2583c 100644
--- a/integration-tests/narayana-jta/src/test/java/io/quarkus/narayana/jta/JdbcObjectStoreTestProfile.java
+++ b/integration-tests/narayana-jta/src/test/java/io/quarkus/narayana/jta/JdbcObjectStoreTestProfile.java
@@ -12,6 +12,11 @@ public Map<String, String> getConfigOverrides() {
         props.put("quarkus.transaction-manager.object-store.type", "jdbc");
         props.put("quarkus.transaction-manager.object-store.create-table", "true");
         props.put("quarkus.transaction-manager.enable-recovery", "true");
+
+        props.put("quarkus.datasource.test.db-kind", "h2");
+        props.put("quarkus.datasource.test.jdbc.url", "jdbc:h2:mem:default");
+        props.put("quarkus.datasource.test.jdbc.transactions", "xa");
+
         return props;
     }
 }