diff --git a/app.go b/bebida-shaker.go similarity index 65% rename from app.go rename to bebida-shaker.go index 438826e..b077d42 100644 --- a/app.go +++ b/bebida-shaker.go @@ -1,11 +1,14 @@ package main import ( + "flag" + "fmt" "os" "strconv" connectors "github.com/RyaxTech/bebida-shaker/connectors" "github.com/RyaxTech/bebida-shaker/events" + "github.com/RyaxTech/bebida-shaker/utils" "github.com/apex/log" ) @@ -84,11 +87,34 @@ func getStrEnv(envName string, defaultValue string) string { } func main() { - log.Info("Starting Bebida Shaker") - params = Parameters{ - maxPendingPunchJob: getIntEnv("BEBIDA_MAX_PENDING_PUNCH_JOB", 2), - HPCSchedulerType: getStrEnv("BEBIDA_HPC_SCHEDULER_TYPE", "OAR"), + annotateCmd := flag.NewFlagSet("annotate", flag.ExitOnError) + deadline := annotateCmd.String("deadline", "", "App deadline date") + duration := annotateCmd.Int("duration", 900, "App duration in seconds") + cores := annotateCmd.Int("cores", 1, "Number of cores reqired") + memory := annotateCmd.Int("memory", 1024, "Amount of memory reqired in Bytes") + + flag.Parse() + if len(os.Args) < 2 { + fmt.Println("expected 'shaker' or 'annotate' subcommands") + os.Exit(1) + } + switch os.Args[1] { + case "run": + log.Info("Starting Bebida Shaker") + params = Parameters{ + maxPendingPunchJob: getIntEnv("BEBIDA_MAX_PENDING_PUNCH_JOB", 2), + HPCSchedulerType: getStrEnv("BEBIDA_HPC_SCHEDULER_TYPE", "OAR"), + } + log.Infof("Parameters: %+v\n", params) + run() + case "annotate": + annotateCmd.Parse(os.Args[2:]) + err := utils.Annotate(annotateCmd.Arg(0), *deadline, *duration, *cores, *memory) + if err != nil { + panic(err) + } + default: + fmt.Println("expected 'annotate' or 'run' subcommands") + os.Exit(1) } - log.Infof("Parameters: %+v\n", params) - run() } diff --git a/flake.nix b/flake.nix index aa2b524..8c005cb 100644 --- a/flake.nix +++ b/flake.nix @@ -12,7 +12,7 @@ "x86_64-linux" # "aarch64-linux" ]; - inherit (flake-utils.lib) eachSystem filterPackages; + inherit (flake-utils.lib) eachSystem; bebidaShakerPackage = { pkgs }: pkgs.buildGoModule { @@ -27,7 +27,6 @@ #} checkPhase = ""; - # vendorHash = null; # vendorHash = pkgs.lib.fakeHash; vendorHash = "sha256-n+Pe2nVWlwDLPbzaWTSYtMyYLzMpC1H+oilg7YJhftI="; }; @@ -72,7 +71,7 @@ Restart = "always"; RestartSec = "5s"; EnvironmentFile = cfg.environmentFile; - ExecStart = "${cfg.package}/bin/bebida-shaker"; + ExecStart = "${cfg.package}/bin/bebida-shaker run"; }; }; }; diff --git a/go.mod b/go.mod index fc8a15f..af1d578 100644 --- a/go.mod +++ b/go.mod @@ -44,7 +44,7 @@ require ( google.golang.org/protobuf v1.28.1 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/yaml.v2 v2.4.0 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect + gopkg.in/yaml.v3 v3.0.1 k8s.io/api v0.26.0 k8s.io/apimachinery v0.26.0 k8s.io/klog/v2 v2.80.1 // indirect diff --git a/utils/bebida.go b/utils/bebida.go new file mode 100644 index 0000000..bd225a2 --- /dev/null +++ b/utils/bebida.go @@ -0,0 +1,49 @@ +package utils + +import ( + "bytes" + "fmt" + "io" + "io/ioutil" + "strconv" + + "gopkg.in/yaml.v3" +) + +var prefix = "ryax.tech/" + +func Annotate(filePath string, deadline string, duration int, cores int, memory int) error { + + fileContent, err := ioutil.ReadFile(filePath) + if err != nil { + return err + } + decoder := yaml.NewDecoder(bytes.NewBuffer(fileContent)) + for { + data := make(map[interface{}]interface{}) + + err := decoder.Decode(&data) + if err != nil { + if err == io.EOF { + break + } + return err + } + if data["kind"] == "Pod" { + annotations := make(map[interface{}]interface{}) + annotations[prefix+"deadline"] = deadline + annotations[prefix+"duration"] = strconv.Itoa(duration) + annotations[prefix+"resources.cores"] = strconv.Itoa(cores) + annotations[prefix+"resources.memory"] = strconv.Itoa(memory) + annotationsInYaml := data["metadata"].(map[string]interface{}) + annotationsInYaml["annotations"] = annotations + } + strData, err := yaml.Marshal(&data) + if err != nil { + return err + } + fmt.Println(string(strData)) + fmt.Println("---") + } + return nil +}