fix for one managment server shutdown

Signed-off-by: Abhishek Kumar <abhishek.mrt22@gmail.com>

Author: shwstppr

server/src/main/java/com/cloud/ha/HighAvailabilityManagerImpl.java

@@ -1122,6 +1122,8 @@ public boolean stop() {
 
         _executor.shutdown();
 
+        _haDao.markServerPendingWorksAsInvestigating(_msServer.getId());
+
         return true;
     }
 

server/src/main/java/com/cloud/ha/dao/HighAvailabilityDao.java

@@ -87,4 +87,5 @@ public interface HighAvailabilityDao extends GenericDao<HaWorkVO, Long> {
     List<HaWorkVO> listPendingMigrationsForVm(long vmId);
     int expungeByVmList(List<Long> vmIds, Long batchSize);
     void markPendingWorksAsInvestigating();
+    void markServerPendingWorksAsInvestigating(long managementServerId);
 }

server/src/main/java/com/cloud/ha/dao/HighAvailabilityDaoImpl.java

@@ -272,14 +272,30 @@ public int expungeByVmList(List<Long> vmIds, Long batchSize) {
         return batchExpunge(sc, batchSize);
     }
 
+    protected void updatePendingWorkToInvestigating(SearchCriteria<HaWorkVO> sc) {
+        HaWorkVO haWorkVO = createForUpdate();
+        haWorkVO.setStep(Step.Investigating);
+        UpdateBuilder updateBuilder = getUpdateBuilder(haWorkVO);
+        update(updateBuilder, sc, null);
+    }
+
     @Override
     public void markPendingWorksAsInvestigating() {
         final SearchCriteria<HaWorkVO> sc = TBASearch.create();
         sc.setParameters("time", System.currentTimeMillis() >> 10);
         sc.setParameters("step", Step.Done, Step.Cancelled);
-        HaWorkVO haWorkVO = createForUpdate();
-        haWorkVO.setStep(Step.Investigating);
-        UpdateBuilder updateBuilder = getUpdateBuilder(haWorkVO);
-        update(updateBuilder, sc, null);
+        updatePendingWorkToInvestigating(sc);
+    }
+
+    @Override
+    public void markServerPendingWorksAsInvestigating(long managementServerId) {
+        SearchBuilder<HaWorkVO> sb = createSearchBuilder();
+        sb.and("server", sb.entity().getServerId(), Op.EQ);
+        sb.and("step", sb.entity().getStep(), Op.NIN);
+        sb.done();
+        SearchCriteria<HaWorkVO> sc = sb.create();
+        sc.setParameters("server", managementServerId);
+        sc.setParameters("step", Step.Done, Step.Cancelled);
+        updatePendingWorkToInvestigating(sc);
     }
 }

server/src/test/java/com/cloud/ha/dao/HighAvailabilityDaoImplTest.java

@@ -71,7 +71,6 @@ public void testExpungeByVmList() {
                 .batchExpunge(sc, batchSize);
     }
 
-
     @Test
     public void testMarkPendingWorksAsInvestigating() throws Exception {
         SearchBuilder<HaWorkVO> mockTBASearch = Mockito.mock(SearchBuilder.class);
@@ -96,4 +95,31 @@ public void testMarkPendingWorksAsInvestigating() throws Exception {
             Mockito.verify(highAvailabilityDaoImpl).update(Mockito.eq(mockUpdateBuilder), Mockito.eq(mockSearchCriteria), Mockito.isNull());
         }
     }
+
+    @Test
+    public void testMarkServerPendingWorksAsInvestigating() {
+        SearchBuilder<HaWorkVO> mockSearch = Mockito.mock(SearchBuilder.class);
+        Mockito.doReturn(Mockito.mock(HaWorkVO.class)).when(mockSearch).entity();
+        Mockito.doReturn(mockSearch).when(highAvailabilityDaoImpl).createSearchBuilder();
+        SearchCriteria<HaWorkVO> mockSearchCriteria = Mockito.mock(SearchCriteria.class);
+        UpdateBuilder mockUpdateBuilder = Mockito.mock(UpdateBuilder.class);
+        Mockito.when(mockSearch.create()).thenReturn(mockSearchCriteria);
+        Mockito.doNothing().when(mockSearchCriteria).setParameters(Mockito.eq("server"), Mockito.eq(1L));
+        Mockito.doNothing().when(mockSearchCriteria).setParameters(Mockito.eq("step"), Mockito.eq(HighAvailabilityManager.Step.Done), Mockito.eq(HighAvailabilityManager.Step.Cancelled));
+        HaWorkVO haWorkVO = new HaWorkVO(1L, VirtualMachine.Type.User, null,
+                null, 1L, null, 0, 0,
+                HighAvailabilityManager.ReasonType.HostMaintenance);
+        Mockito.when(highAvailabilityDaoImpl.createForUpdate()).thenReturn(haWorkVO);
+        Mockito.when(highAvailabilityDaoImpl.createForUpdate()).thenReturn(haWorkVO);
+        try(MockedStatic<GenericDaoBase> genericDaoBaseMockedStatic = Mockito.mockStatic(GenericDaoBase.class)) {
+            genericDaoBaseMockedStatic.when(() -> GenericDaoBase.getUpdateBuilder(Mockito.any())).thenReturn(mockUpdateBuilder);
+            Mockito.doReturn(5).when(highAvailabilityDaoImpl).update(Mockito.any(UpdateBuilder.class), Mockito.any(), Mockito.nullable(Integer.class));
+            highAvailabilityDaoImpl.markServerPendingWorksAsInvestigating(1L);
+            Mockito.verify(mockSearch).create();
+            Mockito.verify(mockSearchCriteria).setParameters(Mockito.eq("server"), Mockito.eq(1L));
+            Mockito.verify(mockSearchCriteria).setParameters(Mockito.eq("step"), Mockito.eq(HighAvailabilityManager.Step.Done), Mockito.eq(HighAvailabilityManager.Step.Cancelled));
+            Assert.assertEquals(HighAvailabilityManager.Step.Investigating, haWorkVO.getStep()); // Ensure the step is set correctly
+            Mockito.verify(highAvailabilityDaoImpl).update(Mockito.eq(mockUpdateBuilder), Mockito.eq(mockSearchCriteria), Mockito.isNull());
+        }
+    }
 }