oc-scheduler/infrastructure/session/session.go

package session

import (
	"encoding/json"
	"fmt"
	"oc-scheduler/infrastructure/execution"
	"oc-scheduler/infrastructure/scheduling_resources"
	"sync"
	"time"

	"cloud.o-forge.io/core/oc-lib/dbs"
	"cloud.o-forge.io/core/oc-lib/models/bill"
	"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/order"
	"cloud.o-forge.io/core/oc-lib/models/resources/purchase_resource"
	"cloud.o-forge.io/core/oc-lib/models/utils"
	"cloud.o-forge.io/core/oc-lib/models/workflow_execution"
	"cloud.o-forge.io/core/oc-lib/tools"

	oclib "cloud.o-forge.io/core/oc-lib"
)

type SessionExecutionsService struct {
	Mu                  sync.RWMutex
	ExecutionsSessionID string
}

func NewSessionExecutionsService(sessionID string) *SessionExecutionsService {
	return &SessionExecutionsService{ExecutionsSessionID: sessionID}
}

// ---------------------------------------------------------------------------
// DB helpers
// ---------------------------------------------------------------------------

func (s *SessionExecutionsService) sessionIDFilter(field, id string) *dbs.Filters {
	return &dbs.Filters{
		And: map[string][]dbs.Filter{
			field: {{Operator: dbs.EQUAL.String(), Value: id}},
		},
	}
}

func (s *SessionExecutionsService) loadSession(dt tools.DataType) []scheduling_resources.SchedulerObject {
	results := oclib.NewRequestAdmin(oclib.LibDataEnum(dt), nil).Search(
		s.sessionIDFilter("executions_id", s.ExecutionsSessionID), "", true)
	out := make([]scheduling_resources.SchedulerObject, 0, len(results.Data))
	for _, obj := range results.Data {
		out = append(out, scheduling_resources.ToSchedulerObject(dt, obj))
	}
	return out
}

func (s *SessionExecutionsService) LoadSessionExecs() []*workflow_execution.WorkflowExecution {
	adminReq := &tools.APIRequest{Admin: true}
	results, _, _ := workflow_execution.NewAccessor(adminReq).Search(
		s.sessionIDFilter("executions_id", s.ExecutionsSessionID), "", true)
	out := make([]*workflow_execution.WorkflowExecution, 0)
	for _, obj := range results {
		if exec, ok := obj.(*workflow_execution.WorkflowExecution); ok {
			out = append(out, exec)
		}
	}
	return out
}

func (s *SessionExecutionsService) loadSessionOrder() *order.Order {
	adminReq := &tools.APIRequest{Admin: true}
	results, _, _ := order.NewAccessor(adminReq).Search(
		s.sessionIDFilter("executions_id", s.ExecutionsSessionID), "", true)
	for _, obj := range results {
		if o, ok := obj.(*order.Order); ok {
			return o
		}
	}
	return nil
}

// ---------------------------------------------------------------------------
// Session upsert
// ---------------------------------------------------------------------------

func (s *SessionExecutionsService) UpsertSessionDrafts(
	purchases, bookings []scheduling_resources.SchedulerObject,
	execs []*workflow_execution.WorkflowExecution,
	request *tools.APIRequest,
) {
	adminReq := &tools.APIRequest{Admin: true}

	for dt, datas := range map[tools.DataType][]scheduling_resources.SchedulerObject{
		tools.BOOKING:           bookings,
		tools.PURCHASE_RESOURCE: purchases,
	} {
		existing := map[string]scheduling_resources.SchedulerObject{}
		seen := map[string]bool{}
		for _, bk := range s.loadSession(dt) {
			existing[bk.GetKey()] = bk
		}
		s.upsertDrafts(dt, datas, existing, seen, request)
		for key, prev := range existing {
			if !seen[key] {
				scheduling_resources.GetService().Delete(dt, prev, request)
			}
		}
	}

	for _, old := range s.LoadSessionExecs() {
		execution.UnregisterExecLock(old.GetID())
		workflow_execution.NewAccessor(adminReq).DeleteOne(old.GetID())
	}
	for _, exec := range execs {
		exec.ExecutionsID = s.ExecutionsSessionID
		exec.IsDraft = true
		ex, _, err := utils.GenericStoreOne(exec, workflow_execution.NewAccessor(adminReq))
		if err == nil {
			execution.RegisterExecLock(ex.GetID())
			go execution.WatchDeadline(ex.GetID(), s.ExecutionsSessionID, exec.ExecDate, request)
		}
	}

	if existing := s.loadSessionOrder(); existing == nil {
		GenerateOrder(purchases, bookings, s.ExecutionsSessionID, request)
	} else {
		for _, purch := range purchases {
			existing.Purchases = append(existing.Purchases,
				scheduling_resources.FromSchedulerObject(tools.PURCHASE_RESOURCE, purch).(*purchase_resource.PurchaseResource))
		}
		for _, b := range bookings {
			existing.Bookings = append(existing.Bookings,
				scheduling_resources.FromSchedulerObject(tools.BOOKING, b).(*booking.Booking))
		}
		utils.GenericRawUpdateOne(existing, existing.GetID(), order.NewAccessor(adminReq))
	}
}

func (s *SessionExecutionsService) upsertDrafts(
	dt tools.DataType,
	datas []scheduling_resources.SchedulerObject,
	existing map[string]scheduling_resources.SchedulerObject,
	seen map[string]bool,
	request *tools.APIRequest,
) {
	self := scheduling_resources.GetService().Self()
	fmt.Println("upsertDrafts", len(datas), len(existing))
	for _, bk := range datas {
		if self != nil {
			bk.SetSchedulerPeerID(self.PeerID)
		}
		bk.SetExecutionsID(s.ExecutionsSessionID)
		seen[bk.GetKey()] = true
		if prev, ok := existing[bk.GetKey()]; ok {
			bk.SetID(prev.GetID())
			bk.SetIsDraft(false)
			needsConsiders := scheduling_resources.GetService().PropagateWrite(
				scheduling_resources.FromSchedulerDBObject(dt, bk), bk.GetDestPeer(), dt, request)
			if needsConsiders {
				if payload, err := json.Marshal(execution.ConsidersPayload{ID: bk.GetID()}); err == nil {
					go execution.UpdateExecutionState(payload, dt)
				}
			}
		} else {
			errCh := make(chan error, 1)
			scheduling_resources.GetService().PropagateCreate(
				scheduling_resources.FromSchedulerDBObject(dt, bk), bk.GetDestPeer(), dt, request, errCh)
			<-errCh
		}
	}
}

// ---------------------------------------------------------------------------
// Session lifecycle
// ---------------------------------------------------------------------------

func (s *SessionExecutionsService) CleanupSession(request *tools.APIRequest) {
	adminReq := &tools.APIRequest{Admin: true}
	for _, exec := range s.LoadSessionExecs() {
		execution.Unschedule(exec.GetID(), request)
		workflow_execution.NewAccessor(adminReq).DeleteOne(exec.GetID())
	}
	if o := s.loadSessionOrder(); o != nil {
		order.NewAccessor(adminReq).DeleteOne(o.GetID())
	}
}

func GenerateOrder(
	purchases, bookings []scheduling_resources.SchedulerObject,
	executionsID string,
	request *tools.APIRequest,
) (string, error) {
	newOrder := &order.Order{
		AbstractObject: utils.AbstractObject{
			Name:    "order_" + request.PeerID + "_" + time.Now().UTC().Format("2006-01-02T15:04:05"),
			IsDraft: true,
		},
		ExecutionsID: executionsID,
		Purchases:    []*purchase_resource.PurchaseResource{},
		Bookings:     []*booking.Booking{},
		Status:       enum.PENDING,
	}
	for _, purch := range purchases {
		newOrder.Purchases = append(newOrder.Purchases,
			scheduling_resources.FromSchedulerObject(tools.PURCHASE_RESOURCE, purch).(*purchase_resource.PurchaseResource))
	}
	for _, b := range bookings {
		newOrder.Bookings = append(newOrder.Bookings,
			scheduling_resources.FromSchedulerObject(tools.BOOKING, b).(*booking.Booking))
	}
	res, _, err := order.NewAccessor(request).StoreOne(newOrder)
	if err != nil {
		return "", err
	}
	if _, err := bill.DraftFirstBill(res.(*order.Order), request); err != nil {
		return res.GetID(), err
	}
	return res.GetID(), nil
}

func (s *SessionExecutionsService) ConfirmSession(request *tools.APIRequest) error {
	for _, dt := range []tools.DataType{tools.BOOKING, tools.PURCHASE_RESOURCE} {
		for _, bk := range s.loadSession(dt) {
			bk.SetIsDraft(false)
			needsConsiders := scheduling_resources.GetService().PropagateWrite(
				scheduling_resources.FromSchedulerDBObject(dt, bk), bk.GetDestPeer(), dt, request)
			if needsConsiders {
				if payload, err := json.Marshal(execution.ConsidersPayload{ID: bk.GetID()}); err == nil {
					go execution.UpdateExecutionState(payload, dt)
				}
			}
		}
	}
	return nil
}