From 66f33f484569a199f3b522292fa8a19dfa0aec7e Mon Sep 17 00:00:00 2001
From: Amalka Subasinghe <amalka.subasinghe>
Date: Mon, 24 Jul 2023 13:17:46 +0530
Subject: [PATCH] notify cluster formation changed nug fixes

---
 .../internal/HeartBeatBeaconComponent.java    | 22 ++++++++++---------
 .../HeartBeatManagementServiceImpl.java       |  9 ++++++++
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/internal/HeartBeatBeaconComponent.java b/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/internal/HeartBeatBeaconComponent.java
index 5f07342a41..a526e216fb 100644
--- a/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/internal/HeartBeatBeaconComponent.java
+++ b/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/internal/HeartBeatBeaconComponent.java
@@ -61,20 +61,22 @@ public class HeartBeatBeaconComponent {
                 DataSourceConfig dsConfig = HeartBeatBeaconConfig.getInstance().getDataSourceConfig();
                 HeartBeatBeaconDAOFactory.init(dsConfig);
 
+                ClusterFormationChangedNotifierRepository clusterFormationChangedNotifierRepository
+                        = new ClusterFormationChangedNotifierRepository();
+                List<String> notifiers = HeartBeatBeaconConfig.getInstance().getNotifiers();
+                if (notifiers != null && notifiers.size() > 0) {
+                    for (String notifier : notifiers) {
+                        clusterFormationChangedNotifierRepository.addNotifier(notifier);
+                    }
+                }
+                HeartBeatBeaconDataHolder.getInstance().setClusterFormationChangedNotifierRepository(
+                        clusterFormationChangedNotifierRepository);
+
                 //Setting up executors to notify heart beat status */
                 HeartBeatExecutor.setUpNotifiers(HeartBeatBeaconUtils.getServerDetails());
             }
 
-            ClusterFormationChangedNotifierRepository clusterFormationChangedNotifierRepository
-                    = new ClusterFormationChangedNotifierRepository();
-            List<String> notifiers = HeartBeatBeaconConfig.getInstance().getNotifiers();
-            if (notifiers != null && notifiers.size() > 0) {
-                for (String notifier : notifiers) {
-                    clusterFormationChangedNotifierRepository.addNotifier(notifier);
-                }
-            }
-            HeartBeatBeaconDataHolder.getInstance().setClusterFormationChangedNotifierRepository(
-                    clusterFormationChangedNotifierRepository);
+
 
             if (log.isDebugEnabled()) {
                 log.debug("Heart Beat Notifier bundle has been successfully initialized");
diff --git a/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/service/HeartBeatManagementServiceImpl.java b/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/service/HeartBeatManagementServiceImpl.java
index 7aa362376a..7a31a6ebcc 100644
--- a/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/service/HeartBeatManagementServiceImpl.java
+++ b/components/heartbeat-management/io.entgra.device.mgt.core.server.bootup.heartbeat.beacon/src/main/java/io/entgra/device/mgt/core/server/bootup/heartbeat/beacon/service/HeartBeatManagementServiceImpl.java
@@ -261,7 +261,9 @@ public class HeartBeatManagementServiceImpl implements HeartBeatManagementServic
     public void notifyClusterFormationChanged(int elapsedTimeInSeconds) throws HeartBeatManagementException {
         if (HeartBeatBeaconConfig.getInstance().isEnabled()) {
             try {
+                HeartBeatBeaconDAOFactory.beginTransaction();
                 Map<String, ServerContext> servers = heartBeatDAO.getActiveServerDetails(elapsedTimeInSeconds);
+                HeartBeatBeaconDAOFactory.commitTransaction();
                 if (servers != null && !servers.isEmpty()) {
                     String serverUUID = HeartBeatBeaconDataHolder.getInstance().getLocalServerUUID();
                     ServerContext serverContext = servers.get(serverUUID);
@@ -289,6 +291,13 @@ public class HeartBeatManagementServiceImpl implements HeartBeatManagementServic
                 String msg = "Error occurred while notifyClusterFormationChanged.";
                 log.error(msg, e);
                 throw new HeartBeatManagementException(msg, e);
+            } catch (TransactionManagementException e) {
+                HeartBeatBeaconDAOFactory.rollbackTransaction();
+                String msg = "Error occurred while electing candidate for dynamic task execution. Issue in opening a connection to the underlying data source";
+                log.error(msg, e);
+                throw new HeartBeatManagementException(msg, e);
+            } finally {
+                HeartBeatBeaconDAOFactory.closeConnection();
             }
         } else {
             String msg = "Heart Beat Configuration Disabled. Error while notifyClusterFormationChanged.";