refactor: add lock release in after workflow hook

SalmaElsoly · SalmaElsoly · commit d2f1d3dbca52 · 2025-12-15T08:31:00.000+02:00
diff --git a/backend/internal/api/app/app.go b/backend/internal/api/app/app.go
@@ -96,6 +96,7 @@ func (app *App) registerEWFWorkflows() {
 		app.communication.notificationDispatcher,
 		app.infra.gridClient.GridProxyClient,
 		stripeClient,
+		app.core.locker,
 	)
 }
 
diff --git a/backend/internal/api/handlers/deployment_handler.go b/backend/internal/api/handlers/deployment_handler.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 
+	distributedlocks "kubecloud/internal/core/distributed_locks"
 	"kubecloud/internal/core/models"
 	"kubecloud/internal/core/services"
 	"kubecloud/internal/deployment/kubedeployer"
@@ -240,6 +241,10 @@ func (h *DeploymentHandler) HandleDeployCluster(c *gin.Context) {
 	wfUUID, wfStatus, err := h.svc.AsyncDeployCluster(config, cluster)
 	if err != nil {
 		reqLog.Error().Err(err).Msg("failed to start deployment workflow")
+		if errors.Is(err, distributedlocks.ErrNodeLocked) {
+			Conflict(c, "Node is busy serving another request")
+			return
+		}
 		InternalServerError(c)
 		return
 	}
@@ -412,6 +417,10 @@ func (h *DeploymentHandler) HandleAddNode(c *gin.Context) {
 	wfUUID, wfStatus, err := h.svc.AsyncAddNode(config, cl, cluster.Nodes[0])
 	if err != nil {
 		reqLog.Error().Err(err).Msg("failed to start add node workflow")
+		if errors.Is(err, distributedlocks.ErrNodeLocked) {
+			Conflict(c, "Node is busy serving another request")
+			return
+		}
 		InternalServerError(c)
 		return
 	}
diff --git a/backend/internal/api/handlers/node_handler.go b/backend/internal/api/handlers/node_handler.go
@@ -3,6 +3,7 @@ package handlers
 import (
 	"errors"
 	"fmt"
+	distributedlocks "kubecloud/internal/core/distributed_locks"
 	"kubecloud/internal/core/models"
 	"math/rand/v2"
 	"net/url"
@@ -297,6 +298,10 @@ func (h *NodeHandler) ReserveNodeHandler(c *gin.Context) {
 	wfUUID, err := h.svc.AsyncReserveNode(userID, user.Mnemonic, nodeID)
 	if err != nil {
 		reqLog.Error().Err(err).Msg("failed to start workflow to reserve node")
+		if errors.Is(err, distributedlocks.ErrNodeLocked) {
+			Conflict(c, "Node is busy serving another request")
+			return
+		}
 		InternalServerError(c)
 		return
 	}
diff --git a/backend/internal/core/workflows/deployer_activities.go b/backend/internal/core/workflows/deployer_activities.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	cfg "kubecloud/internal/config"
+	distributedlocks "kubecloud/internal/core/distributed_locks"
 	"kubecloud/internal/core/models"
 	"kubecloud/internal/deployment/kubedeployer"
 	"kubecloud/internal/deployment/statemanager"
@@ -641,7 +642,7 @@ func createAddNodeWorkflowTemplate(notificationDispatcher *notification.Notifica
 	return template
 }
 
-func registerDeploymentActivities(engine *ewf.Engine, metrics *metricsLib.Metrics, clusterRepo models.ClusterRepository, notificationDispatcher *notification.NotificationDispatcher, config cfg.Configuration) {
+func registerDeploymentActivities(engine *ewf.Engine, metrics *metricsLib.Metrics, clusterRepo models.ClusterRepository, notificationDispatcher *notification.NotificationDispatcher, config cfg.Configuration, locker distributedlocks.DistributedLocks) {
 	engine.Register(StepDeployNetwork, DeployNetworkStep(metrics))
 	engine.Register(StepDeployLeaderNode, DeployLeaderNodeStep(metrics))
 	engine.Register(StepBatchDeployAllNodes, BatchDeployAllNodesStep(metrics))
@@ -670,6 +671,7 @@ func registerDeploymentActivities(engine *ewf.Engine, metrics *metricsLib.Metric
 	deployWFTemplate.AfterStepHooks = []ewf.AfterStepHook{
 		notifyStepHook(notificationDispatcher),
 	}
+	deployWFTemplate.AfterWorkflowHooks = append(deployWFTemplate.AfterWorkflowHooks, releaseLocksHook(locker))
 	engine.RegisterTemplate(WorkflowDeployCluster, &deployWFTemplate)
 
 	deleteWFTemplate := createDeployerWorkflowTemplate(notificationDispatcher, engine, metrics)
@@ -695,6 +697,7 @@ func registerDeploymentActivities(engine *ewf.Engine, metrics *metricsLib.Metric
 		{Name: StepVerifyNewNodes, RetryPolicy: longExponentialRetryPolicy},
 		{Name: StepStoreDeployment, RetryPolicy: standardRetryPolicy},
 	}
+	addNodeWFTemplate.AfterWorkflowHooks = append(addNodeWFTemplate.AfterWorkflowHooks, releaseLocksHook(locker))
 	engine.RegisterTemplate(WorkflowAddNode, &addNodeWFTemplate)
 
 	removeNodeWFTemplate := createDeployerWorkflowTemplate(notificationDispatcher, engine, metrics)
diff --git a/backend/internal/core/workflows/hooks.go b/backend/internal/core/workflows/hooks.go
@@ -2,11 +2,13 @@ package workflows
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 
 	"time"
 
+	distributedlocks "kubecloud/internal/core/distributed_locks"
 	"kubecloud/internal/deployment/kubedeployer"
 	"kubecloud/internal/deployment/statemanager"
 	"kubecloud/internal/infrastructure/logger"
@@ -241,3 +243,23 @@ func metricsFailureHook(metrics *metricsLib.Metrics) ewf.AfterWorkflowHook {
 		}
 	}
 }
+
+func releaseLocksHook(locker distributedlocks.DistributedLocks) ewf.AfterWorkflowHook {
+	return func(ctx context.Context, wf *ewf.Workflow, _ error) {
+		log := logger.ForOperation("workflow", "release_locks").With().Str("workflow_name", wf.Name).Logger()
+		lockedKeys, ok := wf.Metadata["locked_keys"]
+		if !ok {
+			return
+		}
+		var lockedKeysJSON map[string]string
+		err := json.Unmarshal([]byte(lockedKeys), &lockedKeysJSON)
+		if err != nil {
+			log.Error().Err(err).Msg("failed to unmarshal locked keys")
+			return
+		}
+		if err := locker.ReleaseLock(ctx, lockedKeysJSON); err != nil {
+			log.Error().Err(err).Msg("failed to release locks")
+			return
+		}
+	}
+}
diff --git a/backend/internal/core/workflows/workflow.go b/backend/internal/core/workflows/workflow.go
@@ -3,6 +3,7 @@ package workflows
 import (
 	"kubecloud/internal/billing"
 	cfg "kubecloud/internal/config"
+	distributedlocks "kubecloud/internal/core/distributed_locks"
 	"kubecloud/internal/core/models"
 	"kubecloud/internal/core/persistence"
 	"kubecloud/internal/infrastructure/kyc"
@@ -31,6 +32,7 @@ func RegisterEWFWorkflows(
 	notificationDispatcher *notification.NotificationDispatcher,
 	proxyClient proxy.Client,
 	stripeClient billing.StripeClient,
+	locker distributedlocks.DistributedLocks,
 ) {
 	userRepo := persistence.NewGormUserRepository(db)
 	clusterRepo := persistence.NewGormClusterRepository(db)
@@ -127,6 +129,7 @@ func RegisterEWFWorkflows(
 		{Name: StepReserveNode, RetryPolicy: &ewf.RetryPolicy{MaxAttempts: 2, BackOff: ewf.ConstantBackoff(2 * time.Second)}},
 		{Name: StepVerifyNodeState, RetryPolicy: &ewf.RetryPolicy{MaxAttempts: 5, BackOff: ewf.ExponentialBackoff(10*time.Second, 2*time.Minute, 2.0)}},
 	}
+	reserveNodeTemplate.AfterWorkflowHooks = append(reserveNodeTemplate.AfterWorkflowHooks, releaseLocksHook(locker))
 	engine.RegisterTemplate(WorkflowReserveNode, &reserveNodeTemplate)
 
 	unreserveNodeTemplate := newKubecloudWorkflowTemplate(notificationDispatcher)
@@ -146,7 +149,7 @@ func RegisterEWFWorkflows(
 	// trackClusterHealthWFTemplate.BeforeWorkflowHooks = []ewf.BeforeWorkflowHook{hookNotificationWorkflowStarted}
 	engine.RegisterTemplate(WorkflowTrackClusterHealth, &trackClusterHealthWFTemplate)
 
-	registerDeploymentActivities(engine, metrics, clusterRepo, notificationDispatcher, config)
+	registerDeploymentActivities(engine, metrics, clusterRepo, notificationDispatcher, config, locker)
 
 	// Email-only workflow for guaranteed email delivery with retries
 	emailNotificationTemplate := ewf.WorkflowTemplate{

Original file line number	Diff line number	Diff line change
`@@ -96,6 +96,7 @@ func (app *App) registerEWFWorkflows() {`
`96`	`96`	`app.communication.notificationDispatcher,`
`97`	`97`	`app.infra.gridClient.GridProxyClient,`
`98`	`98`	`stripeClient,`
	`99`	`+ app.core.locker,`
`99`	`100`	`)`
`100`	`101`	`}`
`101`	`102`