In-Process Task Scheduler Using a Min-Heap

package main

import (
	"container/heap"
	"context"
	"fmt"
	"log"
	"sync"
	"time"
)

type Task struct {
	ID string
	At time.Time
	Do func(context.Context) error
}

type taskHeap []*Task

func (h taskHeap) Len() int           { return len(h) }
func (h taskHeap) Less(i, j int) bool { return h[i].At.Before(h[j].At) }
func (h taskHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *taskHeap) Push(x any)        { *h = append(*h, x.(*Task)) }
func (h *taskHeap) Pop() any {
	old := *h
	n := len(old)
	t := old[n-1]
	old[n-1] = nil
	*h = old[:n-1]
	return t
}
func (h taskHeap) Peek() *Task { if len(h) == 0 { return nil }; return h[0] }

type Scheduler struct {
	mu   sync.Mutex
	h    taskHeap
	wake chan struct{}
}

func NewScheduler() *Scheduler {
	s := &Scheduler{wake: make(chan struct{}, 1)}
	heap.Init(&s.h)
	return s
}

func (s *Scheduler) Schedule(t *Task) {
	s.mu.Lock()
	heap.Push(&s.h, t)
	s.mu.Unlock()
	select {
	case s.wake <- struct{}{}:
	default:
	}
}

func (s *Scheduler) Run(ctx context.Context) error {
	for {
		s.mu.Lock()
		next := s.h.Peek()
		s.mu.Unlock()

		if next == nil {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-s.wake:
				continue
			}
		}

		wait := next.At.Sub(time.Now())
		if wait > 0 {
			timer := time.NewTimer(wait)
			select {
			case <-ctx.Done():
				timer.Stop()
				return ctx.Err()
			case <-s.wake:
				timer.Stop()
				continue
			case <-timer.C:
			}
		}

		s.mu.Lock()
		next = s.h.Peek()
		if next != nil && !next.At.After(time.Now()) {
			heap.Pop(&s.h)
		} else {
			s.mu.Unlock()
			continue
		}
		s.mu.Unlock()

		if err := next.Do(ctx); err != nil {
			return fmt.Errorf("task %s failed: %w", next.ID, err)
		}
	}
}

func main() {
	s := NewScheduler()
	s.Schedule(&Task{
		ID: "task-1",
		At: time.Now().Add(300 * time.Millisecond),
		Do: func(ctx context.Context) error {
			log.Printf("ran task-1")
			return nil
		},
	})
	s.Schedule(&Task{
		ID: "task-2",
		At: time.Now().Add(120 * time.Millisecond),
		Do: func(ctx context.Context) error {
			log.Printf("ran task-2")
			return nil
		},
	})

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	if err := s.Run(ctx); err != nil && err != context.DeadlineExceeded {
		log.Fatalf("scheduler stopped: %v", err)
	}
	log.Printf("scheduler done")
}