Add controller package that schedules collection and listens for events

Author: joelrebel

internal/controller/checkpoint.go

@@ -0,0 +1,139 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	v1cclient "github.com/metal-toolbox/conditionorc/pkg/api/v1/client"
+	v1ctypes "github.com/metal-toolbox/conditionorc/pkg/api/v1/types"
+	cptypes "github.com/metal-toolbox/conditionorc/pkg/types"
+)
+
+var (
+	ErrCheckpointSet    = errors.New("error setting task checkpoint")
+	ErrConditionOrcResp = errors.New("error in conditionorc response")
+)
+
+// TaskCheckpointer checkpoints a task by updating its status
+//
+// Going ahead long running tasks are to set agreed 'checkpoints' where it can resume from in cases of failure.
+type TaskCheckpointer interface {
+	Set(ctx context.Context, task *Task, state cptypes.ConditionState, status string) error
+	Get() error
+}
+
+// OrcCheckpointer implements the TaskCheckpointer interface for the conditionorc-api and the events subsystem.
+type OrcCheckpointer struct {
+	// tasksLocker is the list of tasks being worked on by the orchestrator under a mutex.
+	tasksLocker *TasksLocker
+
+	// orcclient is the condition orchestrator client
+	orcclient *v1cclient.Client // condition orchestrator client
+}
+
+// NewTaskCheckpointer returns a task checkpointer to persist task progress.
+func NewTaskCheckpointer(serverAddress string, tasks *TasksLocker) (TaskCheckpointer, error) {
+	orcclient, err := v1cclient.NewClient(serverAddress)
+	if err != nil {
+		return nil, err
+	}
+
+	return &OrcCheckpointer{tasksLocker: tasks, orcclient: orcclient}, nil
+}
+
+// Get returns current checkpoint information on a task
+//
+// TO be implemented.
+func (c *OrcCheckpointer) Get() error {
+	return nil
+}
+
+// SetTaskProgress updates task progress in the events subsystem, the condition orchestrator and the local TasksLocker data under mutex.
+//
+// The NATS streaming subsystem has to be acked to makes sure it does not redeliver the same message.
+func (c *OrcCheckpointer) Set(ctx context.Context, task *Task, state cptypes.ConditionState, status string) error {
+	previousState := task.State
+	task.State = state
+	task.UpdatedAt = time.Now()
+
+	if status != "" {
+		task.Status = status
+	}
+
+	// update condition orchestrator
+	if err := c.updateCondition(
+		ctx,
+		task.Urn.ResourceID,
+		cptypes.ConditionKind(task.Data.EventType),
+		state,
+		statusInfoJSON(status),
+	); err != nil {
+		return err
+	}
+
+	switch task.State {
+	case cptypes.Pending:
+		// mark task as in progress in the events subsystem
+		// resetting the event subsystem timer for this task.
+		if err := task.Msg.InProgress(); err != nil {
+			return err
+		}
+
+		c.tasksLocker.Add(*task)
+
+	case cptypes.Active:
+		// mark task as in progress in the events subsystem
+		// resetting the event subsystem timer for this task.
+		if err := task.Msg.InProgress(); err != nil {
+			return err
+		}
+
+		if previousState != cptypes.Active {
+			c.tasksLocker.Add(*task)
+		} else {
+			c.tasksLocker.Update(*task)
+		}
+
+	case cptypes.Succeeded, cptypes.Failed:
+		if err := task.Msg.Ack(); err != nil {
+			return err
+		}
+
+		c.tasksLocker.Purge(task.ID)
+	}
+
+	return nil
+}
+
+func statusInfoJSON(s string) json.RawMessage {
+	return []byte(fmt.Sprintf("{%q: %q}", "output", s))
+}
+
+func (c *OrcCheckpointer) updateCondition(ctx context.Context, serverID uuid.UUID, kind cptypes.ConditionKind, state cptypes.ConditionState, status json.RawMessage) error {
+	response, err := c.orcclient.ServerConditionGet(ctx, serverID, kind)
+	if err != nil {
+		return err
+	}
+
+	if response == nil || response.Record == nil || response.Record.Condition == nil {
+		return ErrConditionOrcResp
+	}
+
+	update := v1ctypes.ConditionUpdate{
+		State:           state,
+		Status:          status,
+		ResourceVersion: response.Record.Condition.ResourceVersion,
+	}
+
+	// TODO: add retries for resource version mismatch errors
+	_, err = c.orcclient.ServerConditionUpdate(ctx, serverID, kind, update)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

internal/controller/controller.go

@@ -0,0 +1,201 @@
+package controller
+
+import (
+	"context"
+	"math/rand"
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
+
+	"github.com/metal-toolbox/alloy/internal/app"
+	"github.com/metal-toolbox/alloy/internal/model"
+	"github.com/metal-toolbox/alloy/internal/store"
+
+	cptypes "github.com/metal-toolbox/conditionorc/pkg/types"
+	"github.com/sirupsen/logrus"
+	"go.hollow.sh/toolbox/events"
+)
+
+var (
+	concurrency          = 10
+	collectInterval      = 1 * time.Hour
+	collectIntervalSplay = 10 * time.Minute
+	fetchEventsInterval  = 10 * time.Second
+	AckActiveTimeout     = 3 * time.Minute
+	ackActiveInterval    = 1 * time.Minute
+	TaskTimeout          = 180 * time.Minute
+)
+
+type Controller struct {
+	repository        store.Repository
+	streamBroker      events.StreamBroker
+	checkpointHelper  TaskCheckpointer
+	cfg               *app.Configuration
+	tasksLocker       *TasksLocker
+	syncWG            *sync.WaitGroup
+	logger            *logrus.Logger
+	iterCollectActive bool
+}
+
+func New(
+	ctx context.Context,
+	streamBroker events.StreamBroker,
+	cfg *app.Configuration,
+	syncWG *sync.WaitGroup,
+	logger *logrus.Logger,
+) (*Controller, error) {
+	tasksLocker := NewTasksLocker()
+
+	checkpointHelper, err := NewTaskCheckpointer("http://conditionorc-api:9001", tasksLocker)
+	if err != nil {
+		return nil, err
+	}
+
+	if cfg.Concurrency == 0 {
+		cfg.Concurrency = concurrency
+	}
+
+	if cfg.CollectInterval == 0 {
+		cfg.CollectInterval = collectInterval
+	}
+
+	if cfg.CollectIntervalSplay == 0 {
+		cfg.CollectIntervalSplay = collectIntervalSplay
+	}
+
+	repository, err := store.NewRepository(ctx, cfg.StoreKind, model.AppKindOutOfBand, cfg, logger)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Controller{
+		cfg:              cfg,
+		tasksLocker:      tasksLocker,
+		syncWG:           syncWG,
+		logger:           logger,
+		repository:       repository,
+		streamBroker:     streamBroker,
+		checkpointHelper: checkpointHelper,
+	}, nil
+}
+
+func (c *Controller) connectStream(ctx context.Context) (events.MsgCh, error) {
+	if err := c.streamBroker.Open(); err != nil {
+		return nil, err
+	}
+
+	return c.streamBroker.Subscribe(ctx)
+}
+
+func (c *Controller) Run(ctx context.Context) {
+	// TODO: implement stream reconnect loop
+	eventCh, err := c.connectStream(ctx)
+	if err != nil {
+		c.logger.WithError(err).Error("event stream connection error")
+
+		c.loopWithoutEventstream(ctx)
+
+		return
+	}
+
+	c.logger.Info("connected to event stream.")
+
+	c.loopWithEventstream(ctx, eventCh)
+}
+
+func (c *Controller) loopWithEventstream(ctx context.Context, eventCh events.MsgCh) {
+	tickerFetchEvent := time.NewTicker(fetchEventsInterval).C
+	tickerAckActive := time.NewTicker(ackActiveInterval).C
+	tickerCollectAll := time.NewTicker(c.splayInterval()).C
+
+	// to kick alloy to collect for all assets.
+	sigHupCh := make(chan os.Signal, 1)
+	signal.Notify(sigHupCh, syscall.SIGHUP)
+
+	for {
+		select {
+		case <-ctx.Done():
+			c.streamBroker.Close()
+
+		case <-tickerFetchEvent:
+			if c.maximumActive() {
+				continue
+			}
+
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.fetchEvent(ctx, eventCh) }()
+
+		case <-tickerAckActive:
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.ackActive(ctx) }()
+
+		case <-tickerCollectAll:
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.iterCollectOutofband(ctx) }()
+
+		case <-sigHupCh:
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.iterCollectOutofband(ctx) }()
+
+		case event := <-eventCh:
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.processEvent(ctx, event) }()
+		}
+	}
+}
+
+func (c *Controller) loopWithoutEventstream(ctx context.Context) {
+	tickerCollectAll := time.NewTicker(c.splayInterval()).C
+
+	// to kick alloy to collect for all assets.
+	sigHupCh := make(chan os.Signal, 1)
+	signal.Notify(sigHupCh, syscall.SIGHUP)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+
+		case <-sigHupCh:
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.iterCollectOutofband(ctx) }()
+
+		case <-tickerCollectAll:
+			c.syncWG.Add(1)
+
+			go func() { defer c.syncWG.Done(); c.iterCollectOutofband(ctx) }()
+		}
+	}
+}
+
+func (c *Controller) splayInterval() time.Duration {
+	// randomize to given splay value and add to interval
+	rand.Seed(time.Now().UnixNano())
+
+	// nolint:gosec // Ideally this should be using crypto/rand,
+	//                 although the generated random value here is just used to add jitter/splay to
+	//                 the interval value and is not used outside of this context.
+	return c.cfg.CollectInterval + time.Duration(rand.Int63n(int64(c.cfg.CollectIntervalSplay)))
+}
+
+func (c *Controller) maximumActive() bool {
+	active := c.tasksLocker.List()
+
+	var found int
+
+	for idx := range active {
+		if !cptypes.ConditionStateFinalized(active[idx].State) {
+			found++
+		}
+	}
+
+	return found >= concurrency
+}