Fix incorrect device status updates

components/device-mgt/io.entgra.device.mgt.core.device.mgt.core/src/main/java/io/entgra/device/mgt/core/device/mgt/core/status/task/impl/DeviceStatusMonitoringTask.java

@@ -37,10 +37,15 @@ import io.entgra.device.mgt.core.device.mgt.core.dao.DeviceManagementDAOFactory;
 import io.entgra.device.mgt.core.device.mgt.core.status.task.DeviceStatusTaskException;
 import io.entgra.device.mgt.core.device.mgt.core.task.impl.DynamicPartitionedScheduleTask;
 import org.wso2.carbon.context.CarbonContext;
+import org.wso2.carbon.context.PrivilegedCarbonContext;
+import org.wso2.carbon.user.api.UserStoreException;
+import org.wso2.carbon.user.core.service.RealmService;
 
 import java.sql.SQLException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * This implements the Task service which monitors the device activity periodically & update the device-status if
@@ -92,8 +97,35 @@ public class DeviceStatusMonitoringTask extends DynamicPartitionedScheduleTask {
         try {
             List<EnrolmentInfo> enrolmentInfoTobeUpdated = new ArrayList<>();
             List<DeviceMonitoringData> allDevicesForMonitoring = getAllDevicesForMonitoring();
+            Map<Integer, List<DeviceMonitoringData>> tenantDevicesMap = new HashMap<>();
+            List<DeviceMonitoringData> tenantMonitoringData = null;
+            //Delegate the devices in each tenant to a separate list to be updated the statuses.
+            //This improvement has been done since the tenants maintain a separate caches and the task is running
+            //in the super-tenant space. Hence, the device status updates are not reflected in the tenant caches.
+            //Refer to https://roadmap.entgra.net/issues/11386 for more information.
+            for (DeviceMonitoringData deviceMonitoringData : allDevicesForMonitoring) {
+                tenantMonitoringData = tenantDevicesMap.get(deviceMonitoringData.getTenantId());
+                if (tenantMonitoringData == null) {
+                    tenantMonitoringData = new ArrayList<>();
+                }
+                tenantMonitoringData.add(deviceMonitoringData);
+                tenantDevicesMap.put(deviceMonitoringData.getTenantId(), tenantMonitoringData);
+            }
+
+            List<DeviceMonitoringData> monitoringDevices = null;
             long timeMillis = System.currentTimeMillis();
-            for (DeviceMonitoringData monitoringData : allDevicesForMonitoring) {
+            //Retrieving the devices belongs for each tenants and updating the status of the devices.
+            for (Map.Entry<Integer, List<DeviceMonitoringData>> entry : tenantDevicesMap.entrySet()) {
+                Integer tenantId = entry.getKey();
+                RealmService realmService = DeviceManagementDataHolder.getInstance().getRealmService();
+                if (realmService != null) {
+                    String domain = realmService.getTenantManager().getDomain(tenantId);
+                    if (domain != null) {
+                        try {
+                            PrivilegedCarbonContext.startTenantFlow();
+                            PrivilegedCarbonContext.getThreadLocalCarbonContext().setTenantDomain(domain, true);
+                            monitoringDevices = entry.getValue();
+                            for (DeviceMonitoringData monitoringData : monitoringDevices) {
                                 long lastUpdatedTime = (timeMillis - monitoringData
                                         .getLastUpdatedTime()) / 1000;
 
@@ -118,7 +150,17 @@ public class DeviceStatusMonitoringTask extends DynamicPartitionedScheduleTask {
                                             monitoringData.getDevice(), monitoringData.getTenantId());
                                 }
                             }
-
+                        } finally {
+                            PrivilegedCarbonContext.endTenantFlow();
+                        }
+                    } else {
+                        log.error("Failed while running the device status update task. Failed while " +
+                                "extracting tenant domain of the tenant id : " + tenantId);
+                    }
+                } else {
+                    log.error("Failed while running the device status update task. RealmService is not initiated");
+                }
+            }
             if (!enrolmentInfoTobeUpdated.isEmpty()) {
                 try {
                     this.updateDeviceStatus(enrolmentInfoTobeUpdated);
@@ -127,11 +169,12 @@ public class DeviceStatusMonitoringTask extends DynamicPartitionedScheduleTask {
                             "device-status of devices of type '" + deviceType + "'", e);
                 }
             }
-
-
         } catch (DeviceManagementException e) {
             String msg = "Error occurred while retrieving devices list for monitoring.";
             log.error(msg, e);
+        } catch (UserStoreException e) {
+            String msg = "Error occurred while retrieving RealmService instance for updating device status.";
+            log.error(msg, e);
         }
     }