From 73071b6d714b77b271caa901a8aa2893664dce5e Mon Sep 17 00:00:00 2001
From: Mykola Kobets <mykola_kobets@epam.com>
Date: Mon, 2 Sep 2024 11:44:39 +0300
Subject: [PATCH] [umcontroller] Handle node status "unprovisioned"

Signed-off-by: Mykola Kobets <mykola_kobets@epam.com>
---
 umcontroller/umcontroller.go | 92 +++++++++++++++++++++++-------------
 1 file changed, 60 insertions(+), 32 deletions(-)

diff --git a/umcontroller/umcontroller.go b/umcontroller/umcontroller.go
index cb242402..d13f96f1 100644
--- a/umcontroller/umcontroller.go
+++ b/umcontroller/umcontroller.go
@@ -29,6 +29,8 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/exp/slices"
+
 	"github.com/aosedge/aos_common/aoserrors"
 	"github.com/aosedge/aos_common/image"
 	"github.com/aosedge/aos_common/spaceallocator"
@@ -526,22 +528,8 @@ func (umCtrl *Controller) handleNewConnection(umID string, handler *umHandler, s
 		return
 	}
 
-	for _, value := range umCtrl.connections {
-		if value.handler == nil {
-			return
-		}
-	}
-
-	if umCtrl.fsm.Current() == stateInit {
-		log.Debug("All connection to Ums established")
-
-		umCtrl.connectionMonitor.stopConnectionTimer()
-
-		if err := umCtrl.getUpdateComponentsFromStorage(); err != nil {
-			log.Error("Can't read update components from storage: ", err)
-		}
-
-		umCtrl.generateFSMEvent(evAllClientsConnected)
+	if umCtrl.allNodesConnected() {
+		umCtrl.processAllNodesConnected()
 	}
 }
 
@@ -1124,29 +1112,47 @@ func (umCtrl *Controller) createConnections(nodeInfoProvider NodeInfoProvider) e
 }
 
 func (umCtrl *Controller) handleNodeInfoChange(nodeInfo cloudprotocol.NodeInfo) {
-	if nodeInfo.Status == "unprovisioned" {
+	switch nodeInfo.Status {
+	case cloudprotocol.NodeStatusPaused:
 		return
-	}
+	case cloudprotocol.NodeStatusUnprovisioned:
+		ind := slices.IndexFunc(umCtrl.connections, func(conn umConnection) bool {
+			return nodeInfo.NodeID == conn.umID
+		})
+
+		if ind >= 0 {
+			if umCtrl.connections[ind].handler != nil {
+				umCtrl.connections[ind].handler.Close(ConnectionClose)
+			}
 
-	nodeHasUM, err := umCtrl.nodeHasUMComponent(nodeInfo)
-	if err != nil {
-		log.WithField("nodeID", nodeInfo.NodeID).Errorf("Failed to check UM component: %v", err)
-	}
+			umCtrl.connections = append(umCtrl.connections[:ind], umCtrl.connections[ind+1:]...)
+		}
 
-	if !nodeHasUM {
-		return
-	}
+		if umCtrl.allNodesConnected() {
+			umCtrl.processAllNodesConnected()
+		}
+
+	case cloudprotocol.NodeStatusProvisioned:
+		nodeHasUM, err := umCtrl.nodeHasUMComponent(nodeInfo)
+		if err != nil {
+			log.WithField("nodeID", nodeInfo.NodeID).Errorf("Failed to check UM component: %v", err)
+		}
 
-	for _, connection := range umCtrl.connections {
-		if connection.umID == nodeInfo.NodeID {
+		if !nodeHasUM {
 			return
 		}
-	}
 
-	umCtrl.connections = append(umCtrl.connections, umConnection{
-		umID:          nodeInfo.NodeID,
-		isLocalClient: nodeInfo.NodeID == umCtrl.currentNodeID, updatePriority: lowestUpdatePriority, handler: nil,
-	})
+		for _, connection := range umCtrl.connections {
+			if connection.umID == nodeInfo.NodeID {
+				return
+			}
+		}
+
+		umCtrl.connections = append(umCtrl.connections, umConnection{
+			umID:          nodeInfo.NodeID,
+			isLocalClient: nodeInfo.NodeID == umCtrl.currentNodeID, updatePriority: lowestUpdatePriority, handler: nil,
+		})
+	}
 }
 
 func (status systemComponentStatus) String() string {
@@ -1195,3 +1201,25 @@ func (umCtrl *Controller) nodeHasUMComponent(
 
 	return false, nil
 }
+
+func (umCtrl *Controller) allNodesConnected() bool {
+	notConnectedNodeInd := slices.IndexFunc(umCtrl.connections, func(conn umConnection) bool {
+		return conn.handler == nil
+	})
+
+	return notConnectedNodeInd < 0
+}
+
+func (umCtrl *Controller) processAllNodesConnected() {
+	if umCtrl.fsm.Current() == stateInit {
+		log.Debug("All connection to Ums established")
+
+		umCtrl.connectionMonitor.stopConnectionTimer()
+
+		if err := umCtrl.getUpdateComponentsFromStorage(); err != nil {
+			log.Error("Can't read update components from storage: ", err)
+		}
+
+		umCtrl.generateFSMEvent(evAllClientsConnected)
+	}
+}