rara/packages/go/workflow/worker.go

package workflow

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"git.toki-labs.com/toki/rara/packages/go/config"
	"git.toki-labs.com/toki/rara/packages/go/database"
	"git.toki-labs.com/toki/rara/packages/go/database/dbgen"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"go.uber.org/fx"
	"go.uber.org/zap"
)

type Worker struct {
	store         *database.Store
	registry      *Registry
	logger        *zap.Logger
	workerID      string
	executorKind  string
	pollInterval  time.Duration
	leaseDuration time.Duration
	retryDelay    time.Duration
}

type stepInput struct {
	Operation string          `json:"operation"`
	Payload   json.RawMessage `json:"payload"`
}

func NewWorker(
	lifecycle fx.Lifecycle,
	cfg config.Config,
	store *database.Store,
	registry *Registry,
	logger *zap.Logger,
) (*Worker, error) {
	pollInterval, err := cfg.WorkerPollInterval()
	if err != nil {
		return nil, err
	}
	leaseDuration, err := cfg.WorkerLeaseDuration()
	if err != nil {
		return nil, err
	}
	retryDelay, err := cfg.WorkerRetryDelay()
	if err != nil {
		return nil, err
	}
	worker := &Worker{
		store:         store,
		registry:      registry,
		logger:        logger,
		workerID:      cfg.Worker.ID,
		executorKind:  cfg.Worker.ExecutorKind,
		pollInterval:  pollInterval,
		leaseDuration: leaseDuration,
		retryDelay:    retryDelay,
	}
	var cancel context.CancelFunc
	var done chan struct{}
	lifecycle.Append(fx.Hook{
		OnStart: func(context.Context) error {
			runContext, stop := context.WithCancel(context.Background())
			cancel = stop
			done = make(chan struct{})
			go func() {
				defer close(done)
				worker.run(runContext)
			}()
			logger.Info(
				"workflow worker started",
				zap.String("worker_id", worker.workerID),
				zap.String("executor_kind", worker.executorKind),
			)
			return nil
		},
		OnStop: func(ctx context.Context) error {
			if cancel != nil {
				cancel()
			}
			if done == nil {
				return nil
			}
			select {
			case <-done:
				return nil
			case <-ctx.Done():
				return ctx.Err()
			}
		},
	})
	return worker, nil
}

func (worker *Worker) run(ctx context.Context) {
	ticker := time.NewTicker(worker.pollInterval)
	defer ticker.Stop()
	for {
		if err := worker.poll(ctx); err != nil && !errors.Is(err, context.Canceled) {
			worker.logger.Error("workflow poll failed", zap.Error(err))
		}
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
	}
}

func (worker *Worker) poll(ctx context.Context) error {
	step, err := worker.store.Queries.LeaseNextStepRun(ctx, dbgen.LeaseNextStepRunParams{
		WorkerID:              pgtype.Text{String: worker.workerID, Valid: true},
		LeaseSeconds:          int32(worker.leaseDuration / time.Second),
		RequestedExecutorKind: worker.executorKind,
	})
	if errors.Is(err, pgx.ErrNoRows) {
		return nil
	}
	if err != nil {
		return err
	}
	return worker.execute(ctx, step)
}

func (worker *Worker) execute(ctx context.Context, step dbgen.StepRun) error {
	workerID := pgtype.Text{String: worker.workerID, Valid: true}
	if _, err := worker.store.Queries.MarkStepRunRunning(ctx, dbgen.MarkStepRunRunningParams{
		ID:       step.ID,
		LeasedBy: workerID,
	}); err != nil {
		return err
	}

	var input stepInput
	if err := json.Unmarshal(step.Input, &input); err != nil {
		return worker.fail(ctx, step, workerID, err, false)
	}
	output, err := worker.registry.Execute(ctx, input.Operation, input.Payload)
	if err != nil {
		return worker.fail(ctx, step, workerID, err, !errors.Is(err, ErrExecutorNotFound))
	}
	_, err = worker.store.Queries.CompleteStepRun(ctx, dbgen.CompleteStepRunParams{
		ID:       step.ID,
		LeasedBy: workerID,
		Output:   output,
	})
	return err
}

func (worker *Worker) fail(
	ctx context.Context,
	step dbgen.StepRun,
	workerID pgtype.Text,
	executionError error,
	retryable bool,
) error {
	payload, _ := json.Marshal(map[string]string{"message": executionError.Error()})
	_, err := worker.store.Queries.FailStepRun(ctx, dbgen.FailStepRunParams{
		Retryable:         retryable,
		ErrorPayload:      payload,
		RetryDelaySeconds: int32(worker.retryDelay / time.Second),
		ID:                step.ID,
		WorkerID:          workerID,
	})
	if err != nil {
		return err
	}
	worker.logger.Warn(
		"workflow step failed",
		zap.String("step_key", step.StepKey),
		zap.Bool("retryable", retryable),
		zap.Error(executionError),
	)
	return nil
}