tasks/chasm/lib/scheduler/backfiller.go

package scheduler

import (
	"time"

	schedulespb "github.com/hanzoai/tasks/api/schedule/v1"
	"github.com/hanzoai/tasks/chasm"
	"github.com/hanzoai/tasks/chasm/lib/scheduler/gen/schedulerpb/v1"
	schedulescommon "github.com/hanzoai/tasks/common/schedules"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// The Backfiller component is responsible for buffering manually
// requested actions. Each backfill request has its own Backfiller node.
type Backfiller struct {
	chasm.UnimplementedComponent

	*schedulerpb.BackfillerState

	Scheduler chasm.ParentPtr[*Scheduler]
}

type BackfillRequestType int

const (
	RequestTypeTrigger BackfillRequestType = iota
	RequestTypeBackfill
)

// addBackfiller returns an initialized backfiller, adding it to the scheduler's
// Backfillers.
func addBackfiller(
	ctx chasm.MutableContext,
	scheduler *Scheduler,
) *Backfiller {
	id := schedulescommon.GenerateBackfillerID()
	backfiller := newBackfillerWithState(ctx, &schedulerpb.BackfillerState{
		BackfillId:        id,
		LastProcessedTime: timestamppb.New(ctx.Now(scheduler)),
	})

	if scheduler.Backfillers == nil {
		scheduler.Backfillers = make(chasm.Map[string, *Backfiller])
	}
	scheduler.Backfillers[id] = chasm.NewComponentField(ctx, backfiller)

	return backfiller
}

func newBackfillerWithState(ctx chasm.MutableContext, state *schedulerpb.BackfillerState) *Backfiller {
	backfiller := &Backfiller{
		BackfillerState: state,
	}
	backfiller.scheduleTask(ctx, chasm.TaskScheduledTimeImmediate)
	return backfiller
}

// scheduleTask schedules a BackfillerTask at the given time.
func (b *Backfiller) scheduleTask(ctx chasm.MutableContext, scheduledTime time.Time) {
	ctx.AddTask(b, chasm.TaskAttributes{
		ScheduledTime: scheduledTime,
	}, &schedulerpb.BackfillerTask{})
}

func (b *Backfiller) LifecycleState(ctx chasm.Context) chasm.LifecycleState {
	return chasm.LifecycleStateRunning
}

func (b *Backfiller) RequestType() BackfillRequestType {
	if b.GetTriggerRequest() != nil {
		return RequestTypeTrigger
	}

	return RequestTypeBackfill
}

type backfillProgressResult struct {
	// BufferedStarts that should be enqueued to the Invoker.
	BufferedStarts []*schedulespb.BufferedStart

	// High water mark for when state was last updated.
	LastProcessedTime time.Time

	// When true, the backfill has completed and the node can be deleted.
	Complete bool
}