From 9b539609060c3d9bc75c45afdc84bac1d1b844bf Mon Sep 17 00:00:00 2001
From: LP B <xAt0mZ@users.noreply.github.com>
Date: Thu, 2 Feb 2023 16:04:58 +0100
Subject: [PATCH] fix(api/edgestack): update deployments count when env
 relations are updated (#8433)

---
 .../endpointrelation/endpointrelation.go      | 74 ++++++++++++++++++-
 api/datastore/services.go                     |  1 +
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/api/dataservices/endpointrelation/endpointrelation.go b/api/dataservices/endpointrelation/endpointrelation.go
index 2726ee011..eb7089e2b 100644
--- a/api/dataservices/endpointrelation/endpointrelation.go
+++ b/api/dataservices/endpointrelation/endpointrelation.go
@@ -16,13 +16,18 @@ const (
 
 // Service represents a service for managing environment(endpoint) relation data.
 type Service struct {
-	connection portainer.Connection
+	connection    portainer.Connection
+	updateStackFn func(ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error
 }
 
 func (service *Service) BucketName() string {
 	return BucketName
 }
 
+func (service *Service) RegisterUpdateStackFunction(updateFunc func(ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error) {
+	service.updateStackFn = updateFunc
+}
+
 // NewService creates a new instance of a service.
 func NewService(connection portainer.Connection) (*Service, error) {
 	err := connection.SetServiceName(BucketName)
@@ -80,20 +85,36 @@ func (service *Service) Create(endpointRelation *portainer.EndpointRelation) err
 
 // UpdateEndpointRelation updates an Environment(Endpoint) relation object
 func (service *Service) UpdateEndpointRelation(endpointID portainer.EndpointID, endpointRelation *portainer.EndpointRelation) error {
+	previousRelationState, _ := service.EndpointRelation(endpointID)
+
 	identifier := service.connection.ConvertToKey(int(endpointID))
 	err := service.connection.UpdateObject(BucketName, identifier, endpointRelation)
 	cache.Del(endpointID)
+	if err != nil {
+		return err
+	}
 
-	return err
+	updatedRelationState, _ := service.EndpointRelation(endpointID)
+
+	service.updateEdgeStacksAfterRelationChange(previousRelationState, updatedRelationState)
+
+	return nil
 }
 
 // DeleteEndpointRelation deletes an Environment(Endpoint) relation object
 func (service *Service) DeleteEndpointRelation(endpointID portainer.EndpointID) error {
+	deletedRelation, _ := service.EndpointRelation(endpointID)
+
 	identifier := service.connection.ConvertToKey(int(endpointID))
 	err := service.connection.DeleteObject(BucketName, identifier)
 	cache.Del(endpointID)
+	if err != nil {
+		return err
+	}
 
-	return err
+	service.updateEdgeStacksAfterRelationChange(deletedRelation, nil)
+
+	return nil
 }
 
 func (service *Service) InvalidateEdgeCacheForEdgeStack(edgeStackID portainer.EdgeStackID) {
@@ -111,3 +132,50 @@ func (service *Service) InvalidateEdgeCacheForEdgeStack(edgeStackID portainer.Ed
 		}
 	}
 }
+
+func (service *Service) updateEdgeStacksAfterRelationChange(previousRelationState *portainer.EndpointRelation, updatedRelationState *portainer.EndpointRelation) {
+	relations, _ := service.EndpointRelations()
+
+	stacksToUpdate := map[portainer.EdgeStackID]bool{}
+
+	if previousRelationState != nil {
+		for stackId, enabled := range previousRelationState.EdgeStacks {
+			// flag stack for update if stack is not in the updated relation state
+			// = stack has been removed for this relation
+			// or this relation has been deleted
+			if enabled && (updatedRelationState == nil || !updatedRelationState.EdgeStacks[stackId]) {
+				stacksToUpdate[stackId] = true
+			}
+		}
+	}
+
+	if updatedRelationState != nil {
+		for stackId, enabled := range updatedRelationState.EdgeStacks {
+			// flag stack for update if stack is not in the previous relation state
+			// = stack has been added for this relation
+			if enabled && (previousRelationState == nil || !previousRelationState.EdgeStacks[stackId]) {
+				stacksToUpdate[stackId] = true
+			}
+		}
+	}
+
+	// for each stack referenced by the updated relation
+	// list how many time this stack is referenced in all relations
+	// in order to update the stack deployments count
+	for refStackId, refStackEnabled := range stacksToUpdate {
+		if refStackEnabled {
+			numDeployments := 0
+			for _, r := range relations {
+				for sId, enabled := range r.EdgeStacks {
+					if enabled && sId == refStackId {
+						numDeployments += 1
+					}
+				}
+			}
+
+			service.updateStackFn(refStackId, func(edgeStack *portainer.EdgeStack) {
+				edgeStack.NumDeployments = numDeployments
+			})
+		}
+	}
+}
diff --git a/api/datastore/services.go b/api/datastore/services.go
index e687fd0e6..d186493fe 100644
--- a/api/datastore/services.go
+++ b/api/datastore/services.go
@@ -104,6 +104,7 @@ func (store *Store) initServices() error {
 		return err
 	}
 	store.EdgeStackService = edgeStackService
+	endpointRelationService.RegisterUpdateStackFunction(edgeStackService.UpdateEdgeStackFunc)
 
 	edgeGroupService, err := edgegroup.NewService(store.connection)
 	if err != nil {