From 349fce4e6321c2651456e87a11d1c8c4e3bc1fe6 Mon Sep 17 00:00:00 2001 From: "amalka.subasinghe" Date: Thu, 24 Aug 2023 16:29:22 +0530 Subject: [PATCH] improved push notofication task to get allocated operations when heart beat enabled --- .../task/PushNotificationSchedulerTask.java | 27 ++++++++++++++++--- .../PushNotificationSchedulerTaskTest.java | 20 +++++++++----- 2 files changed, 38 insertions(+), 9 deletions(-) diff --git a/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/main/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTask.java b/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/main/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTask.java index 2e859bcead..c4c9df0e56 100644 --- a/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/main/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTask.java +++ b/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/main/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTask.java @@ -17,6 +17,9 @@ */ package io.entgra.device.mgt.core.device.mgt.core.push.notification.mgt.task; +import io.entgra.device.mgt.core.device.mgt.common.ServerCtxInfo; +import io.entgra.device.mgt.core.server.bootup.heartbeat.beacon.exception.HeartBeatManagementException; +import io.entgra.device.mgt.core.server.bootup.heartbeat.beacon.service.HeartBeatManagementService; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.wso2.carbon.context.PrivilegedCarbonContext; @@ -63,9 +66,27 @@ public class PushNotificationSchedulerTask implements Runnable { try { //Get next available operation list per device batch OperationManagementDAOFactory.openConnection(); - operationMappingsTenantMap = operationDAO.getOperationMappingsByStatus(Operation.Status - .PENDING, Operation.PushNotificationStatus.SCHEDULED, DeviceConfigurationManager.getInstance() - .getDeviceManagementConfig().getPushNotificationConfiguration().getSchedulerBatchSize()); + try { + if (DeviceManagementDataHolder.getInstance().getHeartBeatService().isTaskPartitioningEnabled()) { + ServerCtxInfo serverCtxInfo = DeviceManagementDataHolder.getInstance().getHeartBeatService().getServerCtxInfo(); + if (serverCtxInfo != null) { + operationMappingsTenantMap = operationDAO.getAllocatedOperationMappingsByStatus(Operation.Status + .PENDING, Operation.PushNotificationStatus.SCHEDULED, DeviceConfigurationManager.getInstance() + .getDeviceManagementConfig().getPushNotificationConfiguration().getSchedulerBatchSize(), + serverCtxInfo.getActiveServerCount(), serverCtxInfo.getLocalServerHashIdx()); + } else { + if (log.isDebugEnabled()) { + log.debug("Active server information not recorded yet."); + } + } + } else { + operationMappingsTenantMap = operationDAO.getOperationMappingsByStatus(Operation.Status + .PENDING, Operation.PushNotificationStatus.SCHEDULED, DeviceConfigurationManager.getInstance() + .getDeviceManagementConfig().getPushNotificationConfiguration().getSchedulerBatchSize()); + } + } catch (HeartBeatManagementException e) { + throw new RuntimeException(e); + } } catch (OperationManagementDAOException e) { log.error("Unable to retrieve scheduled pending operations for task.", e); } finally { diff --git a/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/test/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTaskTest.java b/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/test/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTaskTest.java index 7d92b9faa1..df04c2fb88 100644 --- a/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/test/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTaskTest.java +++ b/components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/test/java/io/entgra/device/mgt/core/device/mgt/core/push/notification/mgt/task/PushNotificationSchedulerTaskTest.java @@ -17,12 +17,6 @@ */ package io.entgra.device.mgt.core.device.mgt.core.push.notification.mgt.task; -import org.apache.commons.logging.Log; -import org.apache.commons.logging.LogFactory; -import org.mockito.Mockito; -import org.testng.Assert; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.Test; import io.entgra.device.mgt.core.device.mgt.common.exceptions.DeviceManagementException; import io.entgra.device.mgt.core.device.mgt.common.operation.mgt.OperationManagementException; import io.entgra.device.mgt.core.device.mgt.core.common.BaseDeviceManagementTest; @@ -37,6 +31,14 @@ import io.entgra.device.mgt.core.device.mgt.core.operation.mgt.dao.OperationMana import io.entgra.device.mgt.core.device.mgt.core.operation.mgt.dao.OperationManagementDAOFactory; import io.entgra.device.mgt.core.device.mgt.core.service.DeviceManagementProviderService; import io.entgra.device.mgt.core.device.mgt.core.service.DeviceManagementProviderServiceImpl; +import io.entgra.device.mgt.core.server.bootup.heartbeat.beacon.exception.HeartBeatManagementException; +import io.entgra.device.mgt.core.server.bootup.heartbeat.beacon.service.HeartBeatManagementService; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.mockito.Mockito; +import org.testng.Assert; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; import org.wso2.carbon.registry.core.exceptions.RegistryException; import java.sql.SQLException; @@ -52,13 +54,16 @@ public class PushNotificationSchedulerTaskTest extends BaseDeviceManagementTest private PushNotificationSchedulerTask pushNotificationSchedulerTask; private OperationDAO operationDAO; + private HeartBeatManagementService heartBeatManagementService; @BeforeClass public void init() throws DeviceManagementException, RegistryException { DeviceConfigurationManager.getInstance().initConfig(); log.info("Initializing Push Notification Scheduler Test Class"); DeviceManagementServiceComponent.notifyStartupListeners(); this.deviceMgtProviderService = Mockito.mock(DeviceManagementProviderServiceImpl.class, Mockito.CALLS_REAL_METHODS); + this.heartBeatManagementService = Mockito.mock(HeartBeatManagementService.class, Mockito.CALLS_REAL_METHODS); DeviceManagementDataHolder.getInstance().setDeviceManagementProvider(this.deviceMgtProviderService); + DeviceManagementDataHolder.getInstance().setHeartBeatService(this.heartBeatManagementService); this.operationDAO = OperationManagementDAOFactory.getOperationDAO(); this.pushNotificationSchedulerTask = new PushNotificationSchedulerTask(); } @@ -69,6 +74,7 @@ public class PushNotificationSchedulerTaskTest extends BaseDeviceManagementTest OperationManagementDAOException { try { log.info("Attempting to execute push notification task scheduler"); + Mockito.when(this.heartBeatManagementService.isTaskPartitioningEnabled()).thenReturn(false); Mockito.doReturn(new TestNotificationStrategy()).when(this.deviceMgtProviderService) .getNotificationStrategyByDeviceType(Mockito.anyString()); Mockito.doReturn(new io.entgra.device.mgt.core.device.mgt.common.operation.mgt.Operation()) @@ -81,6 +87,8 @@ public class PushNotificationSchedulerTaskTest extends BaseDeviceManagementTest .getPushNotificationConfiguration().getSchedulerBatchSize()); Assert.assertEquals(operationMappingsTenantMap.size(), 0); log.info("Push notification task execution complete."); + } catch (HeartBeatManagementException e) { + throw new RuntimeException(e); } finally { OperationManagementDAOFactory.closeConnection(); }