diff --git a/config/config.go b/config/config.go index db7ad63..d9ef92d 100644 --- a/config/config.go +++ b/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); diff --git a/docker/scale.go b/docker/scale.go index f4e2ed2..208f9b7 100644 --- a/docker/scale.go +++ b/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 - time.Sleep(15 * time.Second); + waitForScale(serviceId, ctx, 0); } 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 - time.Sleep(15 * time.Second); + ctx, cancel := context.WithTimeout(context.Background(), theConfig.Constants.ServiceScaleTimeout); + 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)); + } } diff --git a/main.go b/main.go index 68cc3c9..e22ed9f 100644 --- a/main.go +++ b/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())); }