oc-monitord/main.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"oc-monitord/conf"
	"oc-monitord/models"
	"oc-monitord/workflow_builder"

	oclib "cloud.o-forge.io/core/oc-lib"

	"cloud.o-forge.io/core/oc-lib/logs"
	"cloud.o-forge.io/core/oc-lib/models/utils"
	"cloud.o-forge.io/core/oc-lib/models/workflow_execution"

	"github.com/akamensky/argparse"
	"github.com/google/uuid"
	"github.com/goraz/onion"
	"github.com/rs/zerolog"
)

// Command-line args:
// - url: Loki URL (default: "http://127.0.0.1:3100")
// - execution: Workflow Execution ID (required) to identify the current execution, allows to retrieve Workflow
// - mongo: MongoDB URL (default: "mongodb://127.0.0.1:27017")
// - db: MongoDB database name (default: "DC_myDC")
// - timeout: Execution timeout (default: -1)

var logger zerolog.Logger
var wf_logger zerolog.Logger
var parser argparse.Parser
var monitorLocal bool
var workflowName string

const defaultConfigFile = "/etc/oc/ocmonitord_conf.json"
const localConfigFile = "./conf/local_ocmonitord_conf.json"

func main() {

	os.Setenv("test_service", "true") // Only for service demo, delete before merging on main

	monitorLocal = false
	// Test if monitor is launched outside (with parameters) or in a k8s environment (env variables sets)
	if os.Getenv("KUBERNETES_SERVICE_HOST") == "" {
		// Not in a k8s environment, get conf from parameters
		fmt.Println("Executes outside of k8s")
		parser = *argparse.NewParser("oc-monitord", "Launch the execution of a workflow given as a parameter and sends the produced logs to a loki database")
		loadConfig(false, &parser)
	} else {
		// Executed in a k8s environment
		fmt.Println("Executes inside a k8s")
		monitorLocal = true
		loadConfig(true, nil)
	}
	oclib.InitDaemon("oc-monitord")

	oclib.SetConfig(
		conf.GetConfig().MongoURL,
		conf.GetConfig().Database,
		conf.GetConfig().NatsURL,
		conf.GetConfig().LokiURL,
		conf.GetConfig().Logs,
	)

	logger = logs.CreateLogger("oc-monitord")

	logger.Debug().Msg("Loki URL : " + conf.GetConfig().LokiURL)
	logger.Debug().Msg("Workflow executed : " + conf.GetConfig().ExecutionID)

	wf_id := getWorkflowId(conf.GetConfig().ExecutionID)
	conf.GetConfig().WorkflowID = wf_id

	logger.Debug().Msg("Starting construction of yaml argo for workflow :" + wf_id)

	if _, err := os.Stat("./argo_workflows/"); os.IsNotExist(err) {
		os.Mkdir("./argo_workflows/", 0755)
		logger.Info().Msg("Created argo_workflows/")
	}

	// // create argo
	new_wf := workflow_builder.WorflowDB{}

	err := new_wf.LoadFrom(conf.GetConfig().WorkflowID)
	if err != nil {
		logger.Error().Msg("Could not retrieve workflow " + conf.GetConfig().WorkflowID + " from oc-catalog API")
	}

	argo_file_path, stepMax, err := new_wf.ExportToArgo(conf.GetConfig().Timeout)
	if err != nil {
		logger.Error().Msg("Could not create the Argo file for " + conf.GetConfig().WorkflowID)
		logger.Error().Msg(err.Error())
	}
	logger.Debug().Msg("Created :" + argo_file_path)

	workflowName = getContainerName(argo_file_path)

	wf_logger = logger.With().Str("argo_name", workflowName).Str("workflow_id", conf.GetConfig().WorkflowID).Str("workflow_execution_id", conf.GetConfig().ExecutionID).Logger()
	wf_logger.Debug().Msg("Testing argo name")

	executeWorkflow(argo_file_path, stepMax)

}

// Return the Workflow ID associated to a workflow execution object
func getWorkflowId(exec_id string) string {
	res := oclib.LoadOne(oclib.LibDataEnum(oclib.WORKFLOW_EXECUTION), exec_id)
	if res.Code != 200 {
		logger.Error().Msg("Could not retrieve workflow ID from execution ID " + exec_id)
		return ""
	}
	wf_exec := res.ToWorkflowExecution()

	return wf_exec.WorkflowID
}

// So far we only log the output from

func executeWorkflow(argo_file_path string, stepMax int) {
	// var stdout, stderr, stdout_logs, stderr_logs 	io.ReadCloser
	var stdout, stderr io.ReadCloser
	// var stderr 	io.ReadCloser
	var err error
	cmd := exec.Command("argo", "submit", "--log", "./argo_workflows/"+argo_file_path, "--serviceaccount=argo", "-n", "argo")
	fmt.Println(cmd)
	if stdout, err = cmd.StdoutPipe(); err != nil {
		wf_logger.Error().Msg("Could not retrieve stdoutpipe " + err.Error())
		return
	}
	if err := cmd.Start(); err != nil {
		panic(err)
	}
	var wg sync.WaitGroup
	go logWorkflow(argo_file_path, stepMax, stdout, &wg)

	if err := cmd.Wait(); err != nil {
		wf_logger.Error().Msg("Could not execute argo submit")
		wf_logger.Error().Msg(err.Error() + bufio.NewScanner(stderr).Text())
		updateStatus("fatal")
	}

	wg.Wait()
}

// We could improve this function by creating an object with the same attribute as the output
// and only send a new log if the current object has different values than the previous
func logWorkflow(argo_file_path string, stepMax int, pipe io.ReadCloser, wg *sync.WaitGroup) {
	var current_watch, previous_watch *models.ArgoWatch
	split := strings.Split(argo_file_path, "_")
	argoLogs := models.NewArgoLogs(split[0], "argo", stepMax)

	watch_output := make([]string, 0)
	scanner := bufio.NewScanner(pipe)

	for scanner.Scan() {
		log := scanner.Text()
		watch_output = append(watch_output, log)
		if strings.Contains(log, "Progress:") {
			current_watch = argoLogs.StopStepRecording(watch_output)
			watch_output = []string{}
		} else if strings.Contains(log, "sub-process exited") {
			current_watch = argoLogs.StopStepRecording(watch_output)
		}
		if current_watch != nil && !current_watch.Equals(previous_watch) && current_watch.Name != "" {
			wg.Add(1)
			checkStatus(current_watch, previous_watch)
			jsonified, err := json.Marshal(current_watch)
			if err != nil {
				logger.Error().Msg("Could not create watch log")
			}
			if current_watch.Status == "Failed" {
				wf_logger.Error().Msg(string(jsonified))
			} else {
				wf_logger.Info().Msg(string(jsonified))
			}
			previous_watch = current_watch
			current_watch = &models.ArgoWatch{}
			watch_output = []string{}
			wg.Done()
		}
	}
}

func loadConfig(is_k8s bool, parser *argparse.Parser) {

	var o *onion.Onion

	o = initOnion(o)
	// These variables can only be retrieved in the onion
	// Variables that don't depend on the environment (from conf file), can be loaded after
	// We can't use underscore in the env variable names because it's the delimitor with OCMONITOR too
	setConf(is_k8s, o, parser)

	if !IsValidUUID(conf.GetConfig().ExecutionID) {
		logger.Fatal().Msg("Provided ID is not an UUID")
	}
}

func setConf(is_k8s bool, o *onion.Onion, parser *argparse.Parser) {
	if is_k8s {
		conf.GetConfig().LokiURL = o.GetStringDefault("lokiurl", "http://127.0.0.1:3100")
		i, err := strconv.Atoi(o.GetString("timeout"))
		if err == nil {
			conf.GetConfig().Timeout = i
		} else {
			logger.Error().Msg("Could not parse timeout, using default value")
		}
		conf.GetConfig().ExecutionID = o.GetString("workflow")
		mongo := o.GetStringDefault("mongourl", "mongodb://127.0.0.1:27017")
		db := o.GetStringDefault("database", "DC_myDC")

		conf.GetConfig().MongoURL = mongo
		conf.GetConfig().Database = db
	} else {
		url := parser.String("u", "url", &argparse.Options{Required: true, Default: "http://127.0.0.1:3100", Help: "Url to the Loki database logs will be sent to"})
		execution := parser.String("e", "execution", &argparse.Options{Required: true, Help: "Execution ID of the workflow to request from oc-catalog API"})
		mongo := parser.String("m", "mongo", &argparse.Options{Required: true, Default: "mongodb://127.0.0.1:27017", Help: "URL to reach the MongoDB"})
		db := parser.String("d", "database", &argparse.Options{Required: true, Default: "DC_myDC", Help: "Name of the database to query in MongoDB"})
		timeout := parser.Int("t", "timeout", &argparse.Options{Required: false, Default: -1, Help: "Timeout for the execution of the workflow"})
		err := parser.Parse(os.Args)
		if err != nil {
			fmt.Println(parser.Usage(err))
			os.Exit(1)
		}

		conf.GetConfig().Logs = "debug"

		conf.GetConfig().LokiURL = *url
		conf.GetConfig().MongoURL = *mongo
		conf.GetConfig().Database = *db
		conf.GetConfig().Timeout = *timeout
		conf.GetConfig().ExecutionID = *execution
	}

}

func initOnion(o *onion.Onion) *onion.Onion {
	logger = logs.CreateLogger("oc-monitord")
	configFile := ""

	l3 := onion.NewEnvLayerPrefix("_", "OCMONITORD")
	l2, err := onion.NewFileLayer(defaultConfigFile, nil)
	if err == nil {
		logger.Info().Msg("Config file found : " + defaultConfigFile)
		configFile = defaultConfigFile
	}
	l1, err := onion.NewFileLayer(localConfigFile, nil)
	if err == nil {
		logger.Info().Msg("Local config file found " + localConfigFile + ", overriding default file")
		configFile = localConfigFile
	}
	if configFile == "" {
		logger.Info().Msg("No config file found, using env")
		o = onion.New(l3)
	} else if l1 == nil && l2 == nil {
		o = onion.New(l1, l2, l3)
	} else if l1 == nil {
		o = onion.New(l2, l3)
	} else if l2 == nil {
		o = onion.New(l1, l3)
	}
	return o
}

func IsValidUUID(u string) bool {
	_, err := uuid.Parse(u)
	return err == nil
}

func getContainerName(argo_file string) string {
	regex := "([a-zA-Z]+-[a-zA-Z]+)"
	re := regexp.MustCompile(regex)

	container_name := re.FindString(argo_file)
	return container_name
}

// Uses the ArgoWatch object to update status of the workflow execution object
func checkStatus(current *models.ArgoWatch, previous *models.ArgoWatch) {
	if previous != nil && current.Status != previous.Status {
		updateStatus(current.Status)
	}
}

func updateStatus(status string) {
	exec_id := conf.GetConfig().ExecutionID

	wf_exec := &workflow_execution.WorkflowExecution{AbstractObject: utils.AbstractObject{UUID: conf.GetConfig().ExecutionID}}
	wf_exec.ArgoStatusToState(status)

	serialized := wf_exec.Serialize()

	res := oclib.UpdateOne(oclib.LibDataEnum(oclib.WORKFLOW_EXECUTION), serialized, exec_id)

	if res.Code != 200 {
		logger.Error().Msg("Could not update status for workflow execution " + exec_id)
	}
}