[Fleet] Fix rolling upgrade CANCEL and UI fixes

x-pack/plugins/fleet/public/applications/fleet/sections/agents/agent_list_page/hooks/use_current_upgrades.tsx

@@ -12,9 +12,9 @@ import { sendGetCurrentUpgrades, sendPostCancelAction, useStartServices } from '
 
 import type { CurrentUpgrade } from '../../../../types';
 
-const POLL_INTERVAL = 30 * 1000;
+const POLL_INTERVAL = 2 * 60 * 1000; // 2 minutes
 
-export function useCurrentUpgrades() {
+export function useCurrentUpgrades(onAbortSuccess: () => void) {
   const [currentUpgrades, setCurrentUpgrades] = useState<CurrentUpgrade[]>([]);
   const currentTimeoutRef = useRef<NodeJS.Timeout>();
   const isCancelledRef = useRef<boolean>(false);
@@ -65,7 +65,7 @@ export function useCurrentUpgrades() {
           return;
         }
         await sendPostCancelAction(currentUpgrade.actionId);
-        await refreshUpgrades();
+        await Promise.all([refreshUpgrades(), onAbortSuccess()]);
       } catch (err) {
         notifications.toasts.addError(err, {
           title: i18n.translate('xpack.fleet.currentUpgrade.abortRequestError', {
@@ -74,7 +74,7 @@ export function useCurrentUpgrades() {
         });
       }
     },
-    [refreshUpgrades, notifications.toasts, overlays]
+    [refreshUpgrades, notifications.toasts, overlays, onAbortSuccess]
   );
 
   // Poll for upgrades

x-pack/plugins/fleet/public/applications/fleet/sections/agents/agent_list_page/index.tsx

@@ -338,7 +338,7 @@ export const AgentListPage: React.FunctionComponent<{}> = () => {
   }, [flyoutContext]);
 
   // Current upgrades
-  const { abortUpgrade, currentUpgrades, refreshUpgrades } = useCurrentUpgrades();
+  const { abortUpgrade, currentUpgrades, refreshUpgrades } = useCurrentUpgrades(fetchData);
 
   const columns = [
     {
@@ -545,7 +545,7 @@ export const AgentListPage: React.FunctionComponent<{}> = () => {
         selectionMode={selectionMode}
         currentQuery={kuery}
         selectedAgents={selectedAgents}
-        refreshAgents={() => fetchData()}
+        refreshAgents={() => Promise.all([fetchData(), refreshUpgrades()])}
       />
       <EuiSpacer size="m" />
       {/* Agent total, bulk actions and status bar */}

x-pack/plugins/fleet/server/services/agents/actions.test.ts

@@ -8,6 +8,11 @@
 import { elasticsearchServiceMock } from '@kbn/core/server/mocks';
 
 import { cancelAgentAction } from './actions';
+import { bulkUpdateAgents } from './crud';
+
+jest.mock('./crud');
+
+const mockedBulkUpdateAgents = bulkUpdateAgents as jest.Mock;
 
 describe('Agent actions', () => {
   describe('cancelAgentAction', () => {
@@ -67,5 +72,30 @@ describe('Agent actions', () => {
         })
       );
     });
+
+    it('should cancel UPGRADE action', async () => {
+      const esClient = elasticsearchServiceMock.createInternalClient();
+      esClient.search.mockResolvedValue({
+        hits: {
+          hits: [
+            {
+              _source: {
+                type: 'UPGRADE',
+                action_id: 'action1',
+                agents: ['agent1', 'agent2'],
+                expiration: '2022-05-12T18:16:18.019Z',
+              },
+            },
+          ],
+        },
+      } as any);
+      await cancelAgentAction(esClient, 'action1');
+
+      expect(mockedBulkUpdateAgents).toBeCalled();
+      expect(mockedBulkUpdateAgents).toBeCalledWith(expect.anything(), [
+        expect.objectContaining({ agentId: 'agent1' }),
+        expect.objectContaining({ agentId: 'agent2' }),
+      ]);
+    });
   });
 });

x-pack/plugins/fleet/server/services/agents/actions.ts

@@ -17,6 +17,8 @@ import type {
 import { AGENT_ACTIONS_INDEX, SO_SEARCH_LIMIT } from '../../../common/constants';
 import { AgentActionNotFoundError } from '../../errors';
 
+import { bulkUpdateAgents } from './crud';
+
 const ONE_MONTH_IN_MS = 2592000000;
 
 export async function createAgentAction(
@@ -131,6 +133,18 @@ export async function cancelAgentAction(esClient: ElasticsearchClient, actionId:
       created_at: now,
       expiration: hit._source.expiration,
     });
+    if (hit._source.type === 'UPGRADE') {
+      await bulkUpdateAgents(
+        esClient,
+        hit._source.agents.map((agentId) => ({
+          agentId,
+          data: {
+            upgraded_at: null,
+            upgrade_started_at: null,
+          },
+        }))
+      );
+    }
   }
 
   return {

x-pack/plugins/fleet/server/services/agents/upgrade.ts

@@ -267,6 +267,7 @@ async function _getCancelledActionId(
 ) {
   const res = await esClient.search<FleetServerAgentAction>({
     index: AGENT_ACTIONS_INDEX,
+    ignore_unavailable: true,
     query: {
       bool: {
         must: [
@@ -296,6 +297,7 @@ async function _getCancelledActionId(
 async function _getUpgradeActions(esClient: ElasticsearchClient, now = new Date().toISOString()) {
   const res = await esClient.search<FleetServerAgentAction>({
     index: AGENT_ACTIONS_INDEX,
+    ignore_unavailable: true,
     query: {
       bool: {
         must: [