oc-scheduler/controllers/sheduler.go

package controllers

import (
	"fmt"
	"net/http"
	"oc-scheduler/infrastructure"
	"reflect"
	"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
}

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

// CheckStreamHandler is the WebSocket handler for slot availability checking.
// 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,
	}

	watchedPeers, err := infrastructure.GetWorkflowPeerIDs(wfID, req)
	fmt.Println("Watched peers for workflow", wfID, ":", watchedPeers)
	if err != nil {
		http.Error(w, `{"code":404,"error":"`+err.Error()+`"}`, http.StatusNotFound)
		return
	}

	conn, err := wsUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}

	var ws infrastructure.WorkflowSchedule
	if err := conn.ReadJSON(&ws); err != nil {
		conn.Close()
		return
	}

	plannerCh, plannerUnsub := infrastructure.SubscribePlannerUpdates(watchedPeers)
	wfCh, wfUnsub := infrastructure.SubscribeWorkflowUpdates(wfID)

	executionsID := uuid.New().String()
	ownedPeers := infrastructure.RequestPlannerRefresh(watchedPeers, executionsID)

	self, err := oclib.GetMySelf()
	if err != nil || self == nil {
		logger.Err(err).Msg("could not resolve self peer")
		conn.Close()
		return
	}
	selfPeerID := self.PeerID

	scheduled := false
	confirmed := false

	defer func() {
		conn.Close()
		plannerUnsub()
		wfUnsub()
		infrastructure.ReleaseRefreshOwnership(ownedPeers, executionsID)
		if !confirmed {
			infrastructure.CleanupSession(executionsID, req)
		}
	}()

	pushCheck := func(reschedule bool) error {
		result, checkErr := ws.Check(wfID, asap, preemption, req)
		if checkErr != nil {
			return checkErr
		}
		if result.Available && reschedule {
			ws.Start = result.Start
			if result.End != nil {
				ws.End = result.End
			}
			_, _, execs, purchases, bookings, err := ws.GetBuyAndBook(wfID, req)
			if err != nil {
				return err
			}
			infrastructure.UpsertSessionDrafts(executionsID, execs, purchases, bookings, req)
			scheduled = true
		}
		result.SchedulingID = executionsID
		return conn.WriteJSON(result)
	}

	if err := pushCheck(true); err != nil {
		return
	}

	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 {
				return
			}
			select {
			case updateCh <- updated:
			default:
				<-updateCh
				updateCh <- updated
			}
		}
	}()

	for {
		select {
		case updated := <-updateCh:
			if updated.Confirm {
				ws.UUID = executionsID
				_, _, _, schedErr := infrastructure.Schedule(&ws, wfID, req)
				if schedErr != nil {
					_ = conn.WriteJSON(map[string]interface{}{
						"error": schedErr.Error(),
					})
					return
				}
				confirmed = true
				return
			}
			changed := updated.Cron != ws.Cron ||
				!updated.Start.Equal(ws.Start) ||
				updated.DurationS != ws.DurationS ||
				(updated.End == nil) != (ws.End == nil) ||
				(updated.End != nil && ws.End != nil && !updated.End.Equal(*ws.End)) ||
				updated.BookingMode != ws.BookingMode ||
				!reflect.DeepEqual(updated.SelectedBillingStrategy, ws.SelectedBillingStrategy) ||
				!reflect.DeepEqual(updated.SelectedInstances, ws.SelectedInstances) ||
				!reflect.DeepEqual(updated.SelectedPartnerships, ws.SelectedPartnerships) ||
				!reflect.DeepEqual(updated.SelectedBuyings, ws.SelectedBuyings) ||
				!reflect.DeepEqual(updated.SelectedStrategies, ws.SelectedStrategies)
			infrastructure.CleanupSession(executionsID, req)
			ws = updated
			if err := pushCheck(changed || !scheduled); err != nil {
				return
			}

		case remotePeerID := <-plannerCh:
			if remotePeerID == selfPeerID {
				if scheduled {
					continue
				}
				result, checkErr := ws.Check(wfID, asap, preemption, req)
				if checkErr == nil {
					result.SchedulingID = executionsID
					_ = conn.WriteJSON(result)
				}
				continue
			}
			if err := pushCheck(scheduled); err != nil {
				return
			}

		case <-wfCh:
			if newPeers, err := infrastructure.GetWorkflowPeerIDs(wfID, req); err == nil {
				plannerUnsub()
				watchedPeers = newPeers
				plannerCh, plannerUnsub = infrastructure.SubscribePlannerUpdates(newPeers)
				newOwned := infrastructure.RequestPlannerRefresh(newPeers, executionsID)
				ownedPeers = append(ownedPeers, newOwned...)
			}
			if err := pushCheck(false); err != nil {
				return
			}

		case <-closeCh:
			return
		}
	}
}

// @Title UnSchedule
// @Description unschedule a workflow execution: deletes its bookings on all peers then deletes the execution.
// @Param	id		path 	string	true		"execution id"
// @Success 200 {object} map[string]interface{}
// @router /:id [delete]
func (o *WorkflowSchedulerController) UnSchedule() {
	user, peerID, groups := oclib.ExtractTokenInfo(*o.Ctx.Request)
	executionID := o.Ctx.Input.Param(":id")
	req := &tools.APIRequest{
		Username: user,
		PeerID:   peerID,
		Groups:   groups,
		Admin:    true,
	}
	if err := infrastructure.UnscheduleExecution(executionID, req); err != nil {
		o.Data["json"] = map[string]interface{}{"code": 404, "error": err.Error()}
	} else {
		o.Data["json"] = map[string]interface{}{"code": 200, "error": ""}
	}
	o.ServeJSON()
}

// @Title SearchScheduledDraftOrder
// @Description search draft order for a workflow
// @Param	id		path 	string	true		"id execution"
// @Success 200 {workspace} models.workspace
// @router /:id/order [get]
func (o *WorkflowSchedulerController) SearchScheduledDraftOrder() {
	_, peerID, _ := oclib.ExtractTokenInfo(*o.Ctx.Request)
	id := o.Ctx.Input.Param(":id")
	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.ServeJSON()
}