oc-scheduler/controllers/loki.go

package controllers

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"cloud.o-forge.io/core/oc-lib/config"
	beego "github.com/beego/beego/v2/server/web"
	gorillaws "github.com/gorilla/websocket"
)

// Operations about workflow
type LokiController struct {
	beego.Controller
}

type LokiInfo struct {
	Start string `json:"start"`
	End   string `json:"end"`
}

// @Title GetLogs
// @Description get logs
// @Param	body		body 	models.compute	true		"The compute content"
// @Success 200 {workspace} models.workspace
// @router /:id [post]
func (o *LokiController) GetLogs() {
	id := o.Ctx.Input.Param(":id")
	var resp map[string]interface{}
	json.Unmarshal(o.Ctx.Input.CopyBody(100000), &resp)

	path := "/loki/api/v1/query_range"
	if len(resp) > 0 {
		start := fmt.Sprintf("%v", resp["start"])
		if len(start) > 10 {
			start = start[0:10]
		}
		end := fmt.Sprintf("%v", resp["end"])
		if len(end) > 10 {
			end = end[0:10]
		}
		query := []string{
			"workflow_execution_id=\"" + id + "\"",
		}
		for k, v := range resp {
			if k == "start" || k == "end" {
				continue
			}
			query = append(query, fmt.Sprintf("%v=\"%v\"", k, v))
		}
		if len(query) == 0 || len(start) < 10 || len(end) < 10 {
			o.Ctx.ResponseWriter.WriteHeader(403)
			o.Data["json"] = map[string]string{"error": "Query error, missing data : start, end or query"}
			o.ServeJSON()
			return
		}
		path += "?query={" + strings.Join(query, ", ") + "}&start=" + start + "&end=" + end

		resp, err := http.Get(config.GetConfig().LokiUrl + path) // CALL
		if err != nil {
			o.Ctx.ResponseWriter.WriteHeader(422)
			o.Data["json"] = map[string]string{"error": err.Error()}
			o.ServeJSON()
			return
		}
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		var result map[string]interface{}
		// Unmarshal: string → []byte → object
		err = json.Unmarshal(body, &result)
		if err != nil {
			o.Ctx.ResponseWriter.WriteHeader(403)
			o.Data["json"] = map[string]string{"error": err.Error()}
			o.ServeJSON()
			return
		}
		o.Data["json"] = result
		o.ServeJSON()
		return
	}

	o.Ctx.ResponseWriter.WriteHeader(403)
	o.Data["json"] = map[string]string{"error": "Query error"}
	o.ServeJSON()
}

// LogsStreamHandler streams Loki logs over WebSocket.
//
// The client sends one JSON message with the same format as GetLogs:
//
//	{"start": "<unix-seconds>", "label1": "val1", ...}
//
// The server connects to Loki's /loki/api/v1/tail WebSocket endpoint and
// forwards every message it receives until the client disconnects.
func LogsStreamHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("LogsStreamHandler")
	execID := strings.TrimSuffix(
		strings.TrimPrefix(r.URL.Path, "/oc/logs/"),
		"",
	)
	conn, err := wsUpgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("LogsStreamHandler", err)
		return
	}
	defer conn.Close()
	/*
		var query map[string]interface{}
		if err := conn.ReadJSON(&query); err != nil {
			fmt.Println("LogsStreamHandler ReadJSON", err)
			return
		}
	*/

	start := time.Now().UTC().UnixNano()
	labels := []string{
		"workflow_execution_id=\"" + execID + "\"",
	}
	fmt.Println("LOKI START", start, labels)
	// Build Loki tail WS URL (http→ws, https→wss).
	lokiBase := config.GetConfig().LokiUrl
	lokiBase = strings.Replace(lokiBase, "https://", "wss://", 1)
	lokiBase = strings.Replace(lokiBase, "http://", "ws://", 1)
	lokiURL := lokiBase + "/loki/api/v1/tail?" + url.Values{
		"query": {"{" + strings.Join(labels, ", ") + "}"},
		"start": {fmt.Sprintf("%v", start)},
	}.Encode()

	headers := http.Header{}
	headers.Set("X-Scope-OrgID", "1")

	lokiConn, resp, err := gorillaws.DefaultDialer.Dial(lokiURL, headers)
	fmt.Println("LOKI LISTEN", lokiBase, err)
	if err != nil {
		if resp != nil {
			body, _ := io.ReadAll(resp.Body)
			fmt.Printf("Handshake failed: status=%d body=%s", resp.StatusCode, string(body))
		}
		_ = conn.WriteJSON(map[string]string{"error": "loki: " + err.Error()})
		return
	}
	defer lokiConn.Close()

	errCh := make(chan error, 2)

	// Forward Loki → client.
	go func() {
		for {
			_, msg, err := lokiConn.ReadMessage()
			if err != nil {
				errCh <- err
				return
			}
			var result map[string]interface{}
			if json.Unmarshal(msg, &result) == nil {
				fmt.Println(result)
				if err := conn.WriteJSON(result); err != nil {
					errCh <- err
					return
				}
			}
		}
	}()

	// Detect client disconnect (read pump).
	go func() {
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				errCh <- err
				return
			}
		}
	}()

	<-errCh
}