Luigit
repositories / bugabinga.net

bugabinga.net

personal infrastructure for bugabinga!

owned by admin

services/luci/internal/scheduler/scheduler.go

Raw
package scheduler

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync/atomic"
	"time"

	"bugabinga.net/luci/internal/ciconfig"
	"bugabinga.net/luci/internal/gitrepo"
	"bugabinga.net/luci/internal/history"
	"bugabinga.net/luci/internal/inbox"
	"bugabinga.net/luci/internal/refscan"
	"bugabinga.net/luci/internal/trigger"
)

const (
	minimumInterval    = 15 * time.Minute
	maxQueuedSchedules = 64
)

var (
	lastRunID  atomic.Int64
	writeInbox = inbox.WriteID
)

type Config struct {
	DataDir   string
	RepoRoots []string
	InboxDir  string
	Now       func() time.Time
}

type state struct {
	Repo          string                 `json:"repo"`
	Job           string                 `json:"job"`
	LastEvaluated string                 `json:"last_evaluated,omitempty"`
	LastEnqueued  string                 `json:"last_enqueued,omitempty"`
	PendingID     string                 `json:"pending_id,omitempty"`
	Pending       *trigger.ScheduleEvent `json:"pending,omitempty"`
}

func Poll(cfg Config) ([]string, error) {
	minute := time.Now().UTC().Truncate(time.Minute)
	if cfg.Now != nil {
		minute = cfg.Now().UTC().Truncate(time.Minute)
	}
	var errs []error
	if err := recoverPending(cfg); err != nil {
		errs = append(errs, err)
	}
	refs, scanErr := refscan.Scan(cfg.RepoRoots)
	if scanErr != nil {
		errs = append(errs, scanErr)
	}
	repos := uniqueRepos(refs)
	pending, _, err := inbox.List(cfg.InboxDir)
	if err != nil {
		return nil, errors.Join(append(errs, err)...)
	}
	outstanding, scheduledCount := scheduledJobs(pending)
	reserved, reservedCount, err := pendingReservations(cfg.DataDir)
	if err != nil {
		return nil, errors.Join(append(errs, err)...)
	}
	for key := range reserved {
		outstanding[key] = true
	}
	scheduledCount += reservedCount
	var ids []string
	for _, repo := range repos {
		repoPath, err := gitrepo.RepoPath(cfg.RepoRoots, repo)
		if err != nil {
			errs = append(errs, err)
			continue
		}
		ref, rev, err := gitrepo.DefaultBranch(repoPath)
		if err != nil {
			errs = append(errs, fmt.Errorf("schedule %s: %w", repo, err))
			continue
		}
		workspace := filepath.Join(cfg.DataDir, "workspaces", ".scheduler", repo)
		if err := gitrepo.Checkout(repoPath, rev, workspace); err != nil {
			errs = append(errs, fmt.Errorf("schedule %s: %w", repo, err))
			continue
		}
		ci, err := ciconfig.Load(workspace)
		_ = os.RemoveAll(workspace)
		if err != nil {
			if recordErr := recordConfigFailure(cfg.DataDir, repo, ref, rev, minute, err); recordErr != nil {
				errs = append(errs, recordErr)
			}
			errs = append(errs, fmt.Errorf("schedule %s@%s: %w", repo, rev, err))
			continue
		}
		for _, job := range ci.Jobs {
			if job.Trigger.Schedule == nil || !ciconfig.ScheduleMatches(job.Trigger.Schedule.Expression, minute) {
				continue
			}
			current, path, err := loadState(cfg.DataDir, repo, job.Name)
			if err != nil {
				errs = append(errs, err)
				continue
			}
			slot := minute.Format("2006-01-02T15:04Z")
			if current.Pending != nil || current.LastEvaluated >= slot {
				continue
			}
			current.Repo, current.Job, current.LastEvaluated = repo, job.Name, slot
			if tooSoon(current.LastEnqueued, minute) || outstanding[repo+"\x00"+job.Name] || scheduledCount >= maxQueuedSchedules {
				if err := saveState(path, current); err != nil {
					errs = append(errs, err)
				}
				continue
			}
			event := trigger.ScheduleEvent{Kind: "schedule", Repo: repo, Job: job.Name, Rev: rev, Ref: ref, Minute: slot}
			id := nextRunID()
			current.LastEnqueued, current.PendingID, current.Pending = slot, id, &event
			if err := saveState(path, current); err != nil {
				errs = append(errs, err)
				continue
			}
			outstanding[repo+"\x00"+job.Name] = true
			scheduledCount++
			if err := publishPending(cfg, path, &current); err != nil {
				errs = append(errs, err)
				continue
			}
			ids = append(ids, id)
		}
	}
	return ids, errors.Join(errs...)
}

func nextRunID() string {
	for {
		previous := lastRunID.Load()
		next := time.Now().UnixNano()
		if next <= previous {
			next = previous + 1
		}
		if lastRunID.CompareAndSwap(previous, next) {
			return fmt.Sprintf("%d", next)
		}
	}
}

func recoverPending(cfg Config) error {
	paths, err := statePaths(cfg.DataDir)
	if err != nil {
		return err
	}
	var errs []error
	for _, path := range paths {
		current, err := readState(path)
		if err != nil {
			errs = append(errs, err)
			continue
		}
		if current.Pending != nil {
			if err := publishPending(cfg, path, &current); err != nil {
				errs = append(errs, err)
			}
		}
	}
	return errors.Join(errs...)
}

func publishPending(cfg Config, path string, current *state) error {
	exists, err := runExists(cfg, current.PendingID)
	if err != nil {
		return err
	}
	if !exists {
		if err := writeInbox(cfg.InboxDir, current.PendingID, *current.Pending); err != nil && !errors.Is(err, os.ErrExist) {
			return err
		}
	} else if err := inbox.Sync(cfg.InboxDir); err != nil {
		return err
	}
	current.PendingID, current.Pending = "", nil
	return saveState(path, *current)
}

func runExists(cfg Config, id string) (bool, error) {
	for _, suffix := range []string{".json", ".json.processing", ".done", ".failed"} {
		if _, err := os.Stat(filepath.Join(cfg.InboxDir, id+suffix)); err == nil {
			return true, nil
		} else if !errors.Is(err, os.ErrNotExist) {
			return false, err
		}
	}
	return (history.Store{DataDir: cfg.DataDir}).Contains(id, "run")
}

func recordConfigFailure(dataDir, repo, ref, rev string, minute time.Time, configErr error) error {
	sum := sha256.Sum256([]byte(repo))
	path := filepath.Join(dataDir, "schedules", "errors", hex.EncodeToString(sum[:])+".json")
	data, err := os.ReadFile(path)
	if err == nil && strings.TrimSpace(string(data)) == rev {
		return nil
	}
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}
	id := nextRunID()
	detail := configErr.Error()
	if len(detail) > 4096 {
		detail = detail[:4096] + "..."
	}
	_, err = (history.Store{DataDir: dataDir}).AppendOnce(history.Event{RunID: id, ChildID: "run", Repo: repo, Job: "schedule", Rev: rev, Ref: ref, Trigger: "schedule", Status: "failed", Detail: detail, Time: minute})
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	return os.WriteFile(path, []byte(rev+"\n"), 0o644)
}

func uniqueRepos(refs []refscan.Ref) []string {
	seen := map[string]bool{}
	for _, ref := range refs {
		seen[ref.Repo] = true
	}
	repos := make([]string, 0, len(seen))
	for repo := range seen {
		repos = append(repos, repo)
	}
	sort.Strings(repos)
	return repos
}

func pendingReservations(dataDir string) (map[string]bool, int, error) {
	paths, err := statePaths(dataDir)
	if err != nil {
		return nil, 0, err
	}
	out := map[string]bool{}
	count := 0
	for _, path := range paths {
		current, err := readState(path)
		if err != nil {
			return nil, 0, err
		}
		if current.Pending != nil {
			out[current.Repo+"\x00"+current.Job] = true
			count++
		}
	}
	return out, count, nil
}

func statePaths(dataDir string) ([]string, error) {
	dir := filepath.Join(dataDir, "schedules")
	entries, err := os.ReadDir(dir)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	paths := make([]string, 0, len(entries))
	for _, entry := range entries {
		if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".json") {
			paths = append(paths, filepath.Join(dir, entry.Name()))
		}
	}
	sort.Strings(paths)
	return paths, nil
}

func scheduledJobs(pending []inbox.Pending) (map[string]bool, int) {
	out := map[string]bool{}
	count := 0
	for _, item := range pending {
		if event, ok := item.Event.(trigger.ScheduleEvent); ok {
			out[event.Repo+"\x00"+event.Job] = true
			count++
		}
	}
	return out, count
}

func tooSoon(last string, minute time.Time) bool {
	if last == "" {
		return false
	}
	previous, err := time.Parse("2006-01-02T15:04Z", last)
	return err == nil && minute.Sub(previous) < minimumInterval
}

func loadState(dataDir, repo, job string) (state, string, error) {
	sum := sha256.Sum256([]byte(repo + "\x00" + job))
	path := filepath.Join(dataDir, "schedules", hex.EncodeToString(sum[:])+".json")
	current, err := readState(path)
	if errors.Is(err, os.ErrNotExist) {
		return state{Repo: repo, Job: job}, path, nil
	}
	if err != nil {
		return state{}, path, err
	}
	if current.Repo != repo || current.Job != job {
		return state{}, path, fmt.Errorf("schedule state identity mismatch for %s/%s", repo, job)
	}
	return current, path, nil
}

func readState(path string) (state, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return state{}, err
	}
	var current state
	if err := json.Unmarshal(data, &current); err != nil {
		return state{}, fmt.Errorf("read schedule state %s: %w", filepath.Base(path), err)
	}
	return current, nil
}

func saveState(path string, current state) error {
	if strings.TrimSpace(current.Repo) == "" || strings.TrimSpace(current.Job) == "" {
		return fmt.Errorf("schedule state requires repo and job")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	data, err := json.Marshal(current)
	if err != nil {
		return err
	}
	data = append(data, '\n')
	file, err := os.CreateTemp(filepath.Dir(path), ".schedule-*.tmp")
	if err != nil {
		return err
	}
	tmp := file.Name()
	defer os.Remove(tmp)
	if err := file.Chmod(0o644); err != nil {
		file.Close()
		return err
	}
	if _, err := file.Write(data); err != nil {
		file.Close()
		return err
	}
	if err := file.Sync(); err != nil {
		file.Close()
		return err
	}
	if err := file.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmp, path); err != nil {
		return err
	}
	dir, err := os.Open(filepath.Dir(path))
	if err != nil {
		return err
	}
	defer dir.Close()
	return dir.Sync()
}