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)

	logger.Info().Msg("Booking for " + wfId)
	req := oclib.NewRequestAdmin(collection, nil)
	// 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:   nil,
		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 },
}

// @Title CheckStream
// @Description WebSocket stream for slot availability checking.
// @Param	id			path	string	true	"workflow id"
// @Param	as_possible	query	bool	false	"search from now"
// @Param	preemption	query	bool	false	"force-valid, surface warnings"
// @router /:id/check [get]
func (o *WorkflowSchedulerController) CheckStream() {
	CheckStreamHandler(o.Ctx.ResponseWriter, o.Ctx.Request)
}

// CheckStreamHandler is the WebSocket handler for slot availability checking.
// It is invoked via the CheckStream controller method.
// Query params: as_possible=true, preemption=true
func CheckStreamHandler(w http.ResponseWriter, r *http.Request) {
	wfID := strings.TrimSuffix(
		strings.TrimPrefix(r.URL.Path, "/oc/"),
		"/check",
	)

	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)

	// Unique ID for this check session — used to track refresh ownership.
	sessionID := uuid.New().String()

	// Request a fresh planner snapshot for every concerned peer.
	// The first session to claim a peer becomes its refresh owner; others skip
	// the duplicate PB_PLANNER emission. ownedPeers grows if the workflow's
	// peer list changes (wfCh).
	ownedPeers := infrastructure.RequestPlannerRefresh(watchedPeers, sessionID)

	// Cleanup on exit (clean or forced): release refresh ownership for the
	// peers this session claimed, which resets Refreshing state and emits
	// PB_CLOSE_PLANNER so oc-discovery stops the planner stream.
	defer func() {
		conn.Close()
		plannerUnsub()
		wfUnsub()
		infrastructure.ReleaseRefreshOwnership(ownedPeers, sessionID)
	}()

	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
	}

	// Read loop: detect client-side close and parse schedule parameter
	// updates (date changes, booking mode changes, …) sent by the client.
	updateCh := make(chan infrastructure.WorkflowSchedule, 1)
	closeCh := make(chan struct{})
	go func() {
		defer close(closeCh)
		for {
			var updated infrastructure.WorkflowSchedule
			if err := conn.ReadJSON(&updated); err != nil {
				// Connection closed or unrecoverable read error.
				return
			}
			// Drop the oldest pending update if the consumer hasn't caught up.
			select {
			case updateCh <- updated:
			default:
				<-updateCh
				updateCh <- updated
			}
		}
	}()

	// Stream loop.
	for {
		select {
		case updated := <-updateCh:
			// The client changed the requested date/params: adopt the new
			// schedule and re-run the check immediately.
			ws = updated
			if err := push(); err != nil {
				return
			}
		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)
				// Claim refresh ownership for any newly added peers.
				newOwned := infrastructure.RequestPlannerRefresh(newPeers, sessionID)
				ownedPeers = append(ownedPeers, newOwned...)
			}
			if err := push(); err != nil {
				return
			}
		case <-plannerCh:
			// A planner snapshot arrived (or was evicted): re-evaluate.
			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()
}