oc-scheduler/controllers/workflow_sheduler.go

package controllers

import (
	"encoding/json"
	"fmt"
	"net/http"
	"oc-scheduler/infrastructure"
	"strings"

	oclib "cloud.o-forge.io/core/oc-lib"
	"cloud.o-forge.io/core/oc-lib/dbs"
	"cloud.o-forge.io/core/oc-lib/tools"
	beego "github.com/beego/beego/v2/server/web"
	"github.com/google/uuid"
	gorillaws "github.com/gorilla/websocket"
)

var orderCollection = oclib.LibDataEnum(oclib.ORDER)
var logger = oclib.GetLogger()

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

// @Title Schedule
// @Description schedule workflow
// @Param	id		path 	string	true		"id execution"
// @Param	body		body 	models.compute	true		"The compute content"
// @Success 200 {workspace} models.workspace
// @router /:id [post]
func (o *WorkflowSchedulerController) Schedule() {
	logger := oclib.GetLogger()

	code := 200
	e := ""
	user, peerID, groups := oclib.ExtractTokenInfo(*o.Ctx.Request)
	wfId := o.Ctx.Input.Param(":id")
	var resp *infrastructure.WorkflowSchedule
	json.Unmarshal(o.Ctx.Input.CopyBody(100000), &resp)

	caller := tools.NewHTTPCaller(map[tools.DataType]map[tools.METHOD]string{ // paths to call other OC services
		tools.PEER: {
			tools.POST: "/status/",
		},
		tools.BOOKING: {
			tools.GET:  "/booking/check/:id/:start_date/:end_date",
			tools.POST: "/booking/",
		},
	})

	logger.Info().Msg("Booking for " + wfId)
	req := oclib.NewRequestAdmin(collection, caller)
	// req := oclib.NewRequest(collection, user, peerID, groups, caller)
	resp.UUID = uuid.New().String()
	fmt.Println(user, peerID, groups)
	sch, _, _, err := resp.Schedules(wfId, &tools.APIRequest{
		Username: user,
		PeerID:   peerID,
		Groups:   groups,
		Caller:   caller,
		Admin:    true,
	})
	if err != nil {
		if sch != nil {
			for _, w := range sch.WorkflowExecution {
				req.DeleteOne(w.GetID())
			}
		}
		o.Data["json"] = map[string]interface{}{
			"data":  nil,
			"code":  409,
			"error": "Error when scheduling your execution(s): " + err.Error(),
		}
		o.ServeJSON()
		return
	}
	o.Data["json"] = map[string]interface{}{
		"data":  sch.WorkflowExecution,
		"code":  code,
		"error": e,
	}
	o.ServeJSON()
}

var wsUpgrader = gorillaws.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// CheckStreamHandler is a plain http.HandlerFunc (registered via beego.Handler
// to avoid Beego's WriteHeader interference with the WebSocket upgrade).
// Path: /oc/:id/check  →  parts = ["", "oc", "<id>", "check"]
// Query params: as_possible=true, preemption=true
func CheckStreamHandler(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(strings.TrimSuffix(r.URL.Path, "/"), "/")
	wfID := parts[len(parts)-2] // second-to-last segment

	q := r.URL.Query()
	asap := q.Get("as_possible") == "true"
	preemption := q.Get("preemption") == "true"

	user, peerID, groups := oclib.ExtractTokenInfo(*r)
	req := &tools.APIRequest{
		Username: user,
		PeerID:   peerID,
		Groups:   groups,
		Caller:   nil,
		Admin:    true,
	}

	// Resolve the peer IDs concerned by this workflow before upgrading so we
	// can abort cleanly with a plain HTTP error if the workflow is not found.
	watchedPeers, err := infrastructure.GetWorkflowPeerIDs(wfID, req)
	fmt.Println("Here my watched peers involved in workflow", watchedPeers)
	if err != nil {
		http.Error(w, `{"code":404,"error":"`+err.Error()+`"}`, http.StatusNotFound)
		return
	}

	// Upgrade to WebSocket.
	conn, err := wsUpgrader.Upgrade(w, r, nil)
	if err != nil {
		// gorilla already wrote the error response
		return
	}

	// Read the schedule parameters sent by the client as the first message.
	var ws infrastructure.WorkflowSchedule
	if err := conn.ReadJSON(&ws); err != nil {
		conn.Close()
		return
	}

	// Subscribe to planner updates for the initially resolved peers and to
	// workflow change notifications (peer list may change on workflow edit).
	plannerCh, plannerUnsub := infrastructure.SubscribePlannerUpdates(watchedPeers)
	wfCh, wfUnsub := infrastructure.SubscribeWorkflowUpdates(wfID)

	// Cleanup on exit: cancel subscriptions, evict planner cache entries,
	// signal PB_CLOSE_PLANNER on NATS for each peer that was being watched.
	defer func() {
		conn.Close()
		plannerUnsub()
		wfUnsub()
		for _, peer := range watchedPeers {
			if b, err := json.Marshal(map[string]interface{}{"peer_id": peer}); err == nil {
				infrastructure.EmitNATS(peer, tools.PropalgationMessage{
					Action:  tools.PB_CLOSE_PLANNER,
					Payload: b,
				})
			}
		}
	}()

	push := func() error {
		result, checkErr := ws.Check(wfID, asap, preemption, req)
		fmt.Println(result, checkErr)
		if checkErr != nil {
			return checkErr
		}
		return conn.WriteJSON(result)
	}

	// Initial check.
	if err := push(); err != nil {
		return
	}

	// Detect client-side close in a separate goroutine.
	closeCh := make(chan struct{})
	go func() {
		defer close(closeCh)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	// Stream loop.
	for {
		select {
		case <-wfCh:
			// The workflow was modified: refresh the peer list and re-subscribe
			// so the stream watches the correct set of planners going forward.
			if newPeers, err := infrastructure.GetWorkflowPeerIDs(wfID, req); err == nil {
				plannerUnsub()
				watchedPeers = newPeers
				plannerCh, plannerUnsub = infrastructure.SubscribePlannerUpdates(newPeers)
			}
			if err := push(); err != nil {
				return
			}
		case <-plannerCh:
			if err := push(); err != nil {
				return
			}
		case <-closeCh:
			return
		}
	}
}

// @Title UnSchedule
// @Description schedule workflow
// @Param	id		path 	string	true		"id execution"
// @Param	body		body 	models.compute	true		"The compute content"
// @Success 200 {workspace} models.workspace
// @router /:id [delete]
func (o *WorkflowSchedulerController) UnSchedule() {
	// user, peerID, groups := oclib.ExtractTokenInfo(*o.Ctx.Request)
	id := o.Ctx.Input.Param(":id")
	// TODO UNSCHEDULER
	filter := &dbs.Filters{
		And: map[string][]dbs.Filter{
			"workflow_id": {{Operator: dbs.EQUAL.String(), Value: id}},
		},
	}
	o.Data["json"] = oclib.NewRequestAdmin(collection, nil).Search(filter, "", true)

	// o.Data["json"] = oclib.NewRequest(collection, user, peerID, groups, nil).Search(filter, "", true)
	o.ServeJSON()
}

// @Title SearchScheduledDraftOrder
// @Description schedule workflow
// @Param	id		path 	string	true		"id execution"
// @Success 200 {workspace} models.workspace
// @router /:id/order [get]
func (o *WorkflowSchedulerController) SearchScheduledDraftOrder() {
	user, peerID, groups := oclib.ExtractTokenInfo(*o.Ctx.Request)
	id := o.Ctx.Input.Param(":id")
	fmt.Println(user, peerID, groups)
	filter := &dbs.Filters{
		And: map[string][]dbs.Filter{
			"workflow_id": {{Operator: dbs.EQUAL.String(), Value: id}},
			"order_by":    {{Operator: dbs.EQUAL.String(), Value: peerID}},
		},
	}
	o.Data["json"] = oclib.NewRequestAdmin(orderCollection, nil).Search(filter, "", true)

	//o.Data["json"] = oclib.NewRequest(orderCollection, user, peerID, groups, nil).Search(filter, "", true)
	o.ServeJSON()
}