From eabd5ade54965496d6519fd1563c64b6e9b884ba Mon Sep 17 00:00:00 2001
From: Frikky <frikky@shuffler.io>
Date: Tue, 23 Jan 2024 01:32:42 +0100
Subject: [PATCH] Upgrade worker details

---
 functions/onprem/worker/go.mod    |  2 +-
 functions/onprem/worker/worker.go | 16 +++++++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/functions/onprem/worker/go.mod b/functions/onprem/worker/go.mod
index 20a427a1..d3eed236 100644
--- a/functions/onprem/worker/go.mod
+++ b/functions/onprem/worker/go.mod
@@ -8,7 +8,7 @@ require (
 	github.com/docker/docker v23.0.3+incompatible
 	github.com/gorilla/mux v1.8.0
 	github.com/satori/go.uuid v1.2.0
-	github.com/shuffle/shuffle-shared v0.5.61
+	github.com/shuffle/shuffle-shared v0.5.63
 	k8s.io/api v0.28.3
 	k8s.io/apimachinery v0.28.3
 	k8s.io/client-go v0.28.3
diff --git a/functions/onprem/worker/worker.go b/functions/onprem/worker/worker.go
index 8f212721..0afcc325 100644
--- a/functions/onprem/worker/worker.go
+++ b/functions/onprem/worker/worker.go
@@ -2485,7 +2485,7 @@ func findActiveSwarmNodes(dockercli *dockerclient.Client) (int64, error) {
 	ctx := context.Background()
 	nodes, err := dockercli.NodeList(ctx, types.NodeListOptions{})
 	if err != nil {
-		return 0, err
+		return 1, err
 	}
 
 	nodeCount := int64(0)
@@ -2496,6 +2496,20 @@ func findActiveSwarmNodes(dockercli *dockerclient.Client) (int64, error) {
 		}
 	}
 
+	// Check for SHUFFLE_MAX_NODES
+	maxNodesString := os.Getenv("SHUFFLE_MAX_SWARM_NODES")
+	// Make it into a number and check if it's lower than nodeCount
+	if len(maxNodesString) > 0 {
+		maxNodes, err := strconv.ParseInt(maxNodesString, 10, 64)
+		if err != nil {
+			return nodeCount, err
+		}
+
+		if nodeCount > maxNodes {
+			nodeCount = maxNodes
+		}
+	}
+
 	return nodeCount, nil
 
 	/*