Docker OC-MONITORD

tools/kubernetes.go

@@ -7,9 +7,13 @@ import (
 	"fmt"
 	"io"
 	"oc-monitord/conf"
+	"oc-monitord/models"
+	"oc-monitord/utils"
 	"os"
 	"sync"
+	"time"
 
+	"cloud.o-forge.io/core/oc-lib/models/common/enum"
 	wfv1 "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
 	"github.com/argoproj/argo-workflows/v3/pkg/client/clientset/versioned"
 	"github.com/google/uuid"
@@ -54,57 +58,75 @@ func NewKubernetesTool() (Tool, error) {
 	}, nil
 }
 
-func (k *KubernetesTools) LogWorkflow(namespace string, workflowName string, argoFilePath string, stepMax int,
-	logFunc func(argoFilePath string, stepMax int, pipe io.ReadCloser, wg *sync.WaitGroup)) error {
+func (k *KubernetesTools) LogWorkflow(execID string, namespace string, workflowName string, argoFilePath string, stepMax int, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs,
+	seen []string, logFunc func(argoFilePath string, stepMax int, pipe io.ReadCloser, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string, wg *sync.WaitGroup)) error {
+	exec := utils.GetExecution(execID)
+	if exec == nil {
+		return errors.New("Could not retrieve workflow ID from execution ID " + execID)
+	}
+	if exec.State == enum.FAILURE || exec.State == enum.SUCCESS {
+		return nil
+	}
+	k.logWorkflow(namespace, workflowName, argoFilePath, stepMax, current_watch, previous_watch, argoLogs, seen, logFunc)
+	return k.LogWorkflow(execID, namespace, workflowName, argoFilePath, stepMax, current_watch, previous_watch, argoLogs, seen, logFunc)
+}
+
+func (k *KubernetesTools) logWorkflow(namespace string, workflowName string, argoFilePath string, stepMax int, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs,
+	seen []string,
+	logFunc func(argoFilePath string, stepMax int, pipe io.ReadCloser, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string, wg *sync.WaitGroup)) error {
 	// List pods related to the Argo workflow
 	labelSelector := fmt.Sprintf("workflows.argoproj.io/workflow=%s", workflowName)
-	pods, err := k.Set.CoreV1().Pods(namespace).List(context.TODO(), metav1.ListOptions{
-		LabelSelector: labelSelector,
-	})
-	if err != nil {
-		panic(fmt.Sprintf("failed to list pods: %v", err))
-	}
-
-	if len(pods.Items) == 0 {
-		return errors.New("no pods found for the workflow")
-	}
-	var wg *sync.WaitGroup
-	// Stream logs from all matching pods
-	wg.Add(len(pods.Items))
-	for _, pod := range pods.Items {
-		for _, container := range pod.Spec.Containers {
-			fmt.Printf("streaming logs for Pod: %s, Container: %s\n", pod.Name, container.Name)
-			go k.streamLogs(namespace, pod.Name, container.Name, argoFilePath, stepMax, wg, logFunc)
+	for retries := 0; retries < 10; retries++ { // Retry for up to ~20 seconds
+		// List workflow pods
+		wfPods, err := k.Set.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{
+			LabelSelector: labelSelector,
+		})
+		if err != nil {
+			return err
 		}
+		// If we found pods, stream logs
+		if len(wfPods.Items) > 0 {
+			var wg sync.WaitGroup
+			// Stream logs from all matching pods
+			for _, pod := range wfPods.Items {
+				for _, container := range pod.Spec.Containers {
+					wg.Add(1)
+					go k.streamLogs(namespace, pod.Name, container.Name, argoFilePath, stepMax, &wg, current_watch, previous_watch, argoLogs, seen, logFunc)
+				}
+			}
+			wg.Wait()
+			return nil
+		}
+		time.Sleep(2 * time.Second) // Wait before retrying
 	}
-	wg.Wait()
-	return nil
+	return errors.New("no pods found for the workflow")
 }
 
 // Function to stream logs
 func (k *KubernetesTools) streamLogs(namespace string, podName string, containerName string,
-	argoFilePath string, stepMax int, wg *sync.WaitGroup,
-	logFunc func(argo_file_path string, stepMax int, pipe io.ReadCloser, wg *sync.WaitGroup)) {
+	argoFilePath string, stepMax int, wg *sync.WaitGroup, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string,
+	logFunc func(argo_file_path string, stepMax int, pipe io.ReadCloser, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string, wg *sync.WaitGroup)) {
 	req := k.Set.CoreV1().Pods(namespace).GetLogs(podName, &corev1.PodLogOptions{
-		Container: containerName,
-		Follow:    true, // Equivalent to -f flag in kubectl logs
+		Container: containerName, // Main container
+		Follow:    true,          // Equivalent to -f flag in kubectl logs
 	})
-
+	defer wg.Done()
 	// Open stream
-	stream, err := req.Stream(context.TODO())
+	stream, err := req.Stream(context.Background())
 	if err != nil {
-		fmt.Printf("Error opening log stream for pod %s: %v\n", podName, err)
 		return
 	}
 	defer stream.Close()
-	logFunc(argoFilePath, stepMax, stream, wg)
+	var internalWg sync.WaitGroup
+	logFunc(argoFilePath, stepMax, stream, current_watch, previous_watch, argoLogs, seen, &internalWg)
+	internalWg.Wait()
 }
 
-func (k *KubernetesTools) CreateArgoWorkflow(path string) error {
+func (k *KubernetesTools) CreateArgoWorkflow(path string, ns string) (string, error) {
 	// Read workflow YAML file
 	workflowYAML, err := os.ReadFile(path)
 	if err != nil {
-		return err
+		return "", err
 	}
 	// Decode the YAML into a Workflow struct
 	scheme := runtime.NewScheme()
@@ -114,21 +136,21 @@ func (k *KubernetesTools) CreateArgoWorkflow(path string) error {
 
 	obj, _, err := decode(workflowYAML, nil, nil)
 	if err != nil {
-		return errors.New("failed to decode YAML: " + err.Error())
+		return "", errors.New("failed to decode YAML: " + err.Error())
 	}
 
 	workflow, ok := obj.(*wfv1.Workflow)
 	if !ok {
-		return errors.New("decoded object is not a Workflow")
+		return "", errors.New("decoded object is not a Workflow")
 	}
 
 	// Create the workflow in the "argo" namespace
-	createdWf, err := k.VersionedSet.ArgoprojV1alpha1().Workflows("argo").Create(context.TODO(), workflow, metav1.CreateOptions{})
+	createdWf, err := k.VersionedSet.ArgoprojV1alpha1().Workflows(ns).Create(context.Background(), workflow, metav1.CreateOptions{})
 	if err != nil {
-		return errors.New("failed to create workflow: " + err.Error())
+		return "", errors.New("failed to create workflow: " + err.Error())
 	}
 	fmt.Printf("workflow %s created in namespace %s\n", createdWf.Name, "argo")
-	return nil
+	return createdWf.Name, nil
 }
 
 func (k *KubernetesTools) CreateAccessSecret(ns string, login string, password string) (string, error) {
@@ -151,7 +173,7 @@ func (k *KubernetesTools) CreateAccessSecret(ns string, login string, password s
 		Data: secretData,
 	}
 	// Create the Secret in Kubernetes
-	_, err := k.Set.CoreV1().Secrets(namespace).Create(context.TODO(), secret, metav1.CreateOptions{})
+	_, err := k.Set.CoreV1().Secrets(namespace).Create(context.Background(), secret, metav1.CreateOptions{})
 	if err != nil {
 		return "", errors.New("Error creating secret: " + err.Error())
 	}