Added more proper waiting mechanism.

config/config.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"os"
+	"time"
 
 	"github.com/docker/docker/api/types/registry"
 );
@@ -24,6 +25,7 @@ type Config struct {
 		HelperServiceName string
 		StorageContainerName string
 		PrepullImageServiceName string
+		ServiceScaleTimeout time.Duration
 	}
 }
 
@@ -45,6 +47,7 @@ func GetConfig()(Config, error){
 	cfg.Constants.HelperServiceName = "blazenaHelper";
 	cfg.Constants.StorageContainerName = "blazenaStorage";
 	cfg.Constants.PrepullImageServiceName = "blazenaPrepull";
+	cfg.Constants.ServiceScaleTimeout = time.Second * 15; 
 
 	err = json.Unmarshal(rawConfig, &cfg);
 

docker/scale.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"os"
 	"strconv"
 	"time"
 
@@ -61,9 +62,10 @@ func scaleDown(w http.ResponseWriter, r *http.Request){
 	if(err != nil){
 		panic("Failed to update service."+ err.Error());
 	}
+	ctx, cancel := context.WithTimeout(context.Background(), theConfig.Constants.ServiceScaleTimeout);
+	defer cancel();	
 
-	//TODO: Add proper wait system
+	waitForScale(serviceId, ctx, 0);
-	time.Sleep(15 * time.Second);
 }
 
 func scaleUp(w http.ResponseWriter, r *http.Request){
@@ -118,6 +120,44 @@ func scaleUp(w http.ResponseWriter, r *http.Request){
 
 	ApiClient.ServiceUpdate(context.Background(), serviceId, inspectresoult.Version, updatedSpec, swarm.ServiceUpdateOptions{});
 
-	//TODO: Add proper wait system
+	ctx, cancel := context.WithTimeout(context.Background(), theConfig.Constants.ServiceScaleTimeout);
-	time.Sleep(15 * time.Second);
+	defer cancel();	
+
+	waitForScale(serviceId, ctx, int(originalScaleChecked));
+}
+
+func waitForScale(serviceId string, ctx context.Context, desiredCount int){
+	startTime := time.Now();
+	for ctx.Err() == nil {
+		tasks, err := ApiClient.TaskList(context.Background(), swarm.TaskListOptions{});
+
+		if err != nil {
+			slog.Error("Failed to list tasks.", slog.Any("propagatedError", err));
+			os.Exit(1);
+		}
+
+		var running int;
+		for _, task := range tasks {
+			if task.ServiceID != serviceId {
+				continue;
+			}
+
+			if task.Status.State == swarm.TaskStateRunning{
+				running ++;
+			}
+		}
+
+		if running == desiredCount {
+			slog.Debug("Rescaled Service",
+				slog.String("serviceId", serviceId),
+				slog.Any("took", time.Since(startTime)),
+				slog.Any("targetScale", desiredCount),
+			);
+			return;
+		}
+		time.Sleep(1*time.Second);
+	}
+	if ctx.Err() == context.DeadlineExceeded{
+		slog.Error("Failed to rescale service in given time.", slog.Any("serviceId", serviceId));
+	}
 }

main.go

@@ -1,12 +1,14 @@
 package main
 
 import (
+	"log/slog"
 	"os"
+	"time"
+
+	cfg "github.com/rony5394/blazena/config"
 	"github.com/rony5394/blazena/docker"
 	"github.com/rony5394/blazena/host"
-	cfg "github.com/rony5394/blazena/config"
+)
-	"log/slog"
-);
 
 /*
 	If the exit code is X then it means Y:
@@ -39,6 +41,8 @@ func main() {
 
 	slog.Debug("Config", slog.Any("Value", config));
 
+	startTime := time.Now();
+
 	mode := os.Args[1];
 	switch mode {
 		case "docker":
@@ -53,4 +57,6 @@ func main() {
 		default:
 			panic("Invalid runtime mode!");
 	}
+
+	slog.Debug("Whole run took", slog.String("time", time.Since(startTime).String()));
 }