diff --git a/config/config.go b/config/config.go index 69c9dfc..db7ad63 100644 --- a/config/config.go +++ b/config/config.go @@ -1,9 +1,12 @@ -package config; +package config import ( - "os" + "encoding/base64" "encoding/json" "errors" + "os" + + "github.com/docker/docker/api/types/registry" ); type Config struct { @@ -15,10 +18,12 @@ type Config struct { LocalBasePath string BlazenaImageUrl string RegistryAuth RegistryAuth + EncodedRegistryAuth string Constants struct{ OverlayNetworkName string HelperServiceName string StorageContainerName string + PrepullImageServiceName string } } @@ -39,7 +44,7 @@ func GetConfig()(Config, error){ cfg.Constants.OverlayNetworkName = "blazenaPohar"; cfg.Constants.HelperServiceName = "blazenaHelper"; cfg.Constants.StorageContainerName = "blazenaStorage"; - + cfg.Constants.PrepullImageServiceName = "blazenaPrepull"; err = json.Unmarshal(rawConfig, &cfg); @@ -47,5 +52,19 @@ func GetConfig()(Config, error){ return cfg, errors.New("Failed to unmarshal config: " + err.Error()); } + authConfig := registry.AuthConfig{ + Username: cfg.RegistryAuth.Username, + Password: cfg.RegistryAuth.Password, + } + + authJSON, err := json.Marshal(authConfig) + + if err != nil { + panic("Failed to marshal auth config!"+ err.Error()); + } + + cfg.EncodedRegistryAuth = base64.StdEncoding.EncodeToString(authJSON); + + return cfg, err; } diff --git a/docker/docker.go b/docker/docker.go index fa9b452..76cadbd 100644 --- a/docker/docker.go +++ b/docker/docker.go @@ -27,6 +27,8 @@ var scale sync.Map; var token string = "12345"; var theConfig cfg.Config; +var theShutdownFuncPointer *context.CancelFunc; + type aService struct{ ServiceId string `json:"serviceId"`; VolumeNames []string `json:"volumeNames"`; @@ -39,6 +41,7 @@ func Run(Config cfg.Config){ theConfig = Config; ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM); + theShutdownFuncPointer = &stop; var err error; ApiClient, err = client.NewClientWithOpts(client.FromEnv); @@ -62,6 +65,7 @@ func Run(Config cfg.Config){ Addr: ":1234", } + http.HandleFunc("/prepull", prepullImage); http.HandleFunc("/services", listServices); http.HandleFunc("/scale/up", scaleUp); http.HandleFunc("/scale/down", scaleDown); diff --git a/docker/prepullImage.go b/docker/prepullImage.go new file mode 100644 index 0000000..f3c98a6 --- /dev/null +++ b/docker/prepullImage.go @@ -0,0 +1,78 @@ +package docker + +import ( + "context" + "fmt" + "log/slog" + "net/http" + "os" + "time" + + "github.com/docker/docker/api/types/swarm" +) + +func prepullImage(w http.ResponseWriter, r *http.Request){ + if r.Method != http.MethodPost{ + w.WriteHeader(http.StatusMethodNotAllowed); + fmt.Fprint(w, "Method Not Allowed"); + return; + } + + if !bearerAuth(w, r){return;} + + sc, err := ApiClient.ServiceCreate(context.Background(), swarm.ServiceSpec{ + Annotations: swarm.Annotations{ + Name: theConfig.Constants.PrepullImageServiceName, + Labels: map[string]string{"blazena.prepull": "true"}, + }, + Mode: swarm.ServiceMode{ + GlobalJob: &swarm.GlobalJob{}, + }, + TaskTemplate: swarm.TaskSpec{ + ContainerSpec: &swarm.ContainerSpec{ + Image: theConfig.BlazenaImageUrl, + Command: []string{"sleep", "1s"}, + }, + }, + }, swarm.ServiceCreateOptions{ + QueryRegistry: true, + EncodedRegistryAuth: theConfig.EncodedRegistryAuth, + }); + + if err != nil { + slog.Error("Failed to create prepull service", slog.Any("propagatedError", err)); + os.Exit(3); + } + + slog.Info("Started Prepull of blazena image."); + + go func(){ + CHECKLOOP: for { + time.Sleep(3*time.Second); + tasks, err := ApiClient.TaskList(context.Background(), swarm.TaskListOptions{}); + + if err != nil { + slog.Error("Failed to list tasks.", slog.Any("propagatedError", err)); + os.Exit(3); + } + + for _, task := range tasks { + if task.ServiceID != sc.ID {continue}; + + if task.Status.State != "complete" { + time.Sleep(1*time.Second); + continue CHECKLOOP + }; + } + + err = ApiClient.ServiceRemove(context.Background(), sc.ID); + + if err != nil { + slog.Warn("Failed to remove prepull service.", slog.Any("propagatedError", err)); + } + + slog.Info("Prepull Finished"); + break CHECKLOOP; + } + }(); +} diff --git a/host/host.go b/host/host.go index 1a480eb..5763bf7 100644 --- a/host/host.go +++ b/host/host.go @@ -41,6 +41,7 @@ func Run(Config cfg.Config) { sshKeyPair := shared.GenerateSSHKeypair(); sshHostPkPem := exchangeKeys(Config, string(sshKeyPair.Public)); + go prepullImage(Config); createStorageContainer(Config, DockerClient, sshKeyPair.Private, sshHostPkPem); services := getServices(Config); diff --git a/host/prepullImage.go b/host/prepullImage.go new file mode 100644 index 0000000..eeb0fc7 --- /dev/null +++ b/host/prepullImage.go @@ -0,0 +1,30 @@ +package host + +import ( + "bytes" + "log/slog" + "net/http" + "os" + + cfg "github.com/rony5394/blazena/config" +); + +func prepullImage(Config cfg.Config){ + rq, err := http.NewRequest("POST", Config.DockerManagerBaseUrl + "/prepull", bytes.NewBufferString("{}")); + + if err != nil{ + slog.Error("Failed to create request", slog.Any("propagatedError", err), slog.String("note", "not send just create the object")); + os.Exit(1); + } + + rq.Header.Set("Authorization", "Bearer "+ token); + rq.Close = true; + rs, err := http.DefaultClient.Do(rq); + + if err != nil{ + slog.Error("Failed to send http request", slog.Any("propagatedError", err)); + os.Exit(1); + } + + defer rs.Body.Close(); +} diff --git a/host/scale.go b/host/scale.go index 8c5cd00..cfd587a 100644 --- a/host/scale.go +++ b/host/scale.go @@ -37,7 +37,9 @@ func scale(Config cfg.Config, serviceId string, up bool)bool{ rq.Header.Set("Authorization", "Bearer "+ token); rq.Close = true; rs, err := http.DefaultClient.Do(rq); - defer rs.Body.Close(); + if err == nil { + rs.Body.Close(); + } if err != nil{ panic("Failed to send http request"+ err.Error()); diff --git a/main.go b/main.go index effb001..68cc3c9 100644 --- a/main.go +++ b/main.go @@ -47,6 +47,9 @@ func main() { case "host": host.Run(config); break; + case "pull": + os.Exit(0); + break; default: panic("Invalid runtime mode!"); }