Decentralized

controllers/workflow_sheduler.go

@@ -3,6 +3,7 @@ package controllers
 import (
 	"encoding/json"
 	"fmt"
+	"net/http"
 	"oc-scheduler/infrastructure"
 	"slices"
 
@@ -15,6 +16,7 @@ import (
 	"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)
@@ -108,6 +110,132 @@ func (o *WorkflowSchedulerController) Schedule() {
 	o.ServeJSON()
 }
 
+var wsUpgrader = gorillaws.Upgrader{
+	CheckOrigin: func(r *http.Request) bool { return true },
+}
+
+// @Title CheckStream
+// @Description WebSocket stream of slot availability for a workflow.
+// After the handshake the client sends one JSON frame containing the
+// WorkflowSchedule parameters (start, end, booking_mode, duration_s, …).
+// The server responds with a CheckResult frame immediately and again each time
+// a planner for one of the workflow's storage/compute peers is updated.
+// When the stream is interrupted the cache entries for those peers are evicted
+// and a PB_CLOSE_PLANNER event is emitted on NATS.
+// Query params:
+//   - as_possible=true  ignore start date, search from now
+//   - preemption=true   validate anyway, raise warnings
+//
+// @Param	id			path	string	true	"workflow id"
+// @Param	as_possible	query	bool	false	"find nearest free slot from now"
+// @Param	preemption	query	bool	false	"validate anyway, raise warnings"
+// @Success 101
+// @router /:id/check [get]
+func (o *WorkflowSchedulerController) CheckStream() {
+	wfID := o.Ctx.Input.Param(":id")
+	asap, _ := o.GetBool("as_possible", false)
+	preemption, _ := o.GetBool("preemption", false)
+
+	user, peerID, groups := oclib.ExtractTokenInfo(*o.Ctx.Request)
+	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)
+	if err != nil {
+		o.Data["json"] = map[string]interface{}{"code": 404, "error": err.Error()}
+		o.ServeJSON()
+		return
+	}
+
+	// Upgrade to WebSocket.
+	conn, err := wsUpgrader.Upgrade(o.Ctx.ResponseWriter, o.Ctx.Request, nil)
+	if err != nil {
+		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)
+		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"