oc-scheduler/infrastructure/scheduler/scheduler.go

package scheduler

import (
	"errors"
	"fmt"
	"oc-scheduler/infrastructure/planner"
	"oc-scheduler/infrastructure/scheduling_resources"
	infUtils "oc-scheduler/infrastructure/utils"
	"strings"
	"time"

	"cloud.o-forge.io/core/oc-lib/models/booking"
	"cloud.o-forge.io/core/oc-lib/models/common/enum"
	"cloud.o-forge.io/core/oc-lib/models/common/pricing"
	"cloud.o-forge.io/core/oc-lib/models/utils"
	"cloud.o-forge.io/core/oc-lib/models/workflow"
	"cloud.o-forge.io/core/oc-lib/models/workflow_execution"
	"cloud.o-forge.io/core/oc-lib/tools"
	"github.com/google/uuid"
	"github.com/robfig/cron"
)

const asapBuffer = 2 * time.Minute

// Schedule holds a resolved start/end pair for a single execution slot.
type Schedule struct {
	Start time.Time
	End   *time.Time
}

// WorkflowSchedule is the flying session object for a scheduling interaction.
// It is never persisted; it lives only for the duration of a WebSocket check session.
type WorkflowSchedule struct {
	UUID              string                                  `json:"id" validate:"required"`
	Workflow          *workflow.Workflow                      `json:"workflow,omitempty"`
	WorkflowExecution []*workflow_execution.WorkflowExecution `json:"workflow_executions,omitempty"`
	Message           string                                  `json:"message,omitempty"`
	Warning           string                                  `json:"warning,omitempty"`
	Start             time.Time                               `json:"start" validate:"required,ltfield=End"`
	End               *time.Time                              `json:"end,omitempty"`
	DurationS         float64                                 `json:"duration_s" default:"-1"`
	Cron              string                                  `json:"cron,omitempty"`

	BookingMode          booking.BookingMode     `json:"booking_mode,omitempty"`
	SelectedInstances    workflow.ConfigItem     `json:"selected_instances"`
	SelectedPartnerships workflow.ConfigItem     `json:"selected_partnerships"`
	SelectedBuyings      workflow.ConfigItem     `json:"selected_buyings"`
	SelectedStrategies   workflow.ConfigItem     `json:"selected_strategies"`
	SelectedBillingStrategy pricing.BillingStrategy `json:"selected_billing_strategy"`

	// Confirm, when true, triggers Schedule() to confirm the drafts held by this session.
	Confirm bool `json:"confirm,omitempty"`
}

// CheckResult is the response payload for an availability check.
type CheckResult struct {
	Available   bool       `json:"available"`
	Start       time.Time  `json:"start"`
	End         *time.Time `json:"end,omitempty"`
	Warnings    []string   `json:"warnings,omitempty"`
	Preemptible bool       `json:"preemptible,omitempty"`
	// SchedulingID is the session UUID the client must supply when confirming.
	SchedulingID string `json:"scheduling_id,omitempty"`
}

// ---------------------------------------------------------------------------
// Check — availability
// ---------------------------------------------------------------------------

// Check verifies whether the requested slot is available across all resource peers.
func (ws *WorkflowSchedule) Check(wfID string, asap bool, preemption bool, request *tools.APIRequest) (*CheckResult, error) {
	fmt.Println("CHECK", asap, "/", preemption)
	obj, code, err := workflow.NewAccessor(request).LoadOne(wfID)
	if code != 200 || err != nil {
		msg := "could not load workflow " + wfID
		if err != nil {
			msg += ": " + err.Error()
		}
		return nil, errors.New(msg)
	}
	wf := obj.(*workflow.Workflow)

	start := ws.Start
	if asap || start.IsZero() {
		start = time.Now().UTC().Add(asapBuffer)
	}

	end := ws.End
	if end == nil {
		if ws.DurationS > 0 {
			e := start.Add(time.Duration(ws.DurationS * float64(time.Second)))
			end = &e
		} else {
			_, longest, _, _, planErr := wf.Planify(
				start, nil,
				ws.SelectedInstances, ws.SelectedPartnerships,
				ws.SelectedBuyings, ws.SelectedStrategies,
				int(ws.BookingMode), nil, request,
			)
			if planErr == nil && longest > 0 {
				e := start.Add(time.Duration(longest) * time.Second)
				end = &e
			}
		}
	}

	checkables := infUtils.CollectBookingResources(wf, ws.SelectedInstances)
	start, end, available, preemptible, warnings := planner.GetPlannerService().FindDate(wfID, checkables, start, end, preemption, asap)

	return &CheckResult{
		Start:       start,
		End:         end,
		Available:   available,
		Preemptible: preemptible,
		Warnings:    warnings,
	}, nil
}

// ---------------------------------------------------------------------------
// GetBuyAndBook — generate scheduling resources
// ---------------------------------------------------------------------------

// GetBuyAndBook runs Planify to generate the purchases and bookings for this session.
func (ws *WorkflowSchedule) GetBuyAndBook(wfID string, request *tools.APIRequest) (
	bool,
	*workflow.Workflow,
	[]*workflow_execution.WorkflowExecution,
	[]scheduling_resources.SchedulerObject,
	[]scheduling_resources.SchedulerObject,
	error,
) {
	res, code, err := workflow.NewAccessor(request).LoadOne(wfID)
	if code != 200 {
		return false, nil, nil, nil, nil,
			errors.New("could not load the workflow: " + err.Error())
	}
	wf := res.(*workflow.Workflow)
	isPreemptible, longest, priceds, wf, err := wf.Planify(
		ws.Start, ws.End,
		ws.SelectedInstances, ws.SelectedPartnerships,
		ws.SelectedBuyings, ws.SelectedStrategies,
		int(ws.BookingMode), nil, request,
	)
	if err != nil {
		return false, wf, nil, nil, nil, err
	}
	ws.DurationS = longest
	ws.Message = "We estimate that the workflow will start at " + ws.Start.String() +
		" and last " + fmt.Sprintf("%v", ws.DurationS) + " seconds."
	if ws.End != nil && ws.Start.Add(time.Duration(longest)*time.Second).After(*ws.End) {
		ws.Warning = "The workflow may be too long to be executed in the given time frame, we will try to book it anyway\n"
	}

	execs, err := ws.GenerateExecutions(wf, isPreemptible)
	if err != nil {
		return false, wf, nil, nil, nil, err
	}

	var purchased, bookings []scheduling_resources.SchedulerObject
	for _, exec := range execs {
		for _, obj := range exec.Buy(ws.SelectedBillingStrategy, ws.UUID, wfID, priceds) {
			purchased = append(purchased, scheduling_resources.ToSchedulerObject(tools.PURCHASE_RESOURCE, obj))
		}
		for _, obj := range exec.Book(ws.UUID, wfID, priceds) {
			bookings = append(bookings, scheduling_resources.ToSchedulerObject(tools.BOOKING, obj))
		}
	}
	return true, wf, execs, purchased, bookings, nil
}

// ---------------------------------------------------------------------------
// GenerateExecutions / GetDates
// ---------------------------------------------------------------------------

// GenerateExecutions expands the cron schedule into WorkflowExecution instances.
func (ws *WorkflowSchedule) GenerateExecutions(wf *workflow.Workflow, isPreemptible bool) ([]*workflow_execution.WorkflowExecution, error) {
	dates, err := ws.GetDates()
	if err != nil {
		return nil, err
	}
	var executions []*workflow_execution.WorkflowExecution
	for _, date := range dates {
		obj := &workflow_execution.WorkflowExecution{
			AbstractObject: utils.AbstractObject{
				UUID: uuid.New().String(),
				Name: wf.Name + "_execution_" + date.Start.String(),
			},
			Priority:     1,
			ExecutionsID: ws.UUID,
			ExecDate:     date.Start,
			EndDate:      date.End,
			State:        enum.DRAFT,
			WorkflowID:   wf.GetID(),
		}
		if ws.BookingMode != booking.PLANNED {
			obj.Priority = 0
		}
		if ws.BookingMode == booking.PREEMPTED && isPreemptible {
			obj.Priority = 7
		}
		executions = append(executions, obj)
	}
	return executions, nil
}

// GetDates parses the cron expression and returns execution date slots.
func (ws *WorkflowSchedule) GetDates() ([]Schedule, error) {
	var schedule []Schedule
	if len(ws.Cron) > 0 {
		if ws.End == nil {
			return schedule, errors.New("a cron task should have an end date")
		}
		if ws.DurationS <= 0 {
			ws.DurationS = ws.End.Sub(ws.Start).Seconds()
		}
		cronStr := strings.Split(ws.Cron, " ")
		if len(cronStr) < 6 {
			return schedule, errors.New("Bad cron message: (" + ws.Cron + "). Should be at least ss mm hh dd MM dw")
		}
		subCron := strings.Join(cronStr[:6], " ")
		specParser := cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow)
		sched, err := specParser.Parse(subCron)
		if err != nil {
			return schedule, errors.New("Bad cron message: " + err.Error())
		}
		for s := sched.Next(ws.Start); !s.IsZero() && s.Before(*ws.End); s = sched.Next(s) {
			e := s.Add(time.Duration(ws.DurationS) * time.Second)
			schedule = append(schedule, Schedule{Start: s, End: &e})
		}
	} else {
		schedule = append(schedule, Schedule{Start: ws.Start, End: ws.End})
	}
	return schedule, nil
}