1package scheduler
2
3import (
4 "context"
5 "fmt"
6 "sort"
7 "strings"
8 "sync"
9 "time"
10
11 "github.com/google/uuid"
12)
13
14// Task is one recurring prompt owned by a session branch. Tasks are kept only
15// in process memory; callers are responsible for stopping the engine with the
16// host process.
17type Task struct {
18 ID string
19 SessionID string
20 Cron string
21 Timezone string
22 Message string
23 CreatedAt time.Time
24 NextRun time.Time
25 LastRun time.Time
26 RunCount int
27 LastError string
28
29 schedule Cron
30 queued bool
31 running bool
32}
33
34// NewTaskInput describes a task to add to an Engine. Location is captured at
35// creation rather than consulting time.Local during later evaluations.
36type NewTaskInput struct {
37 SessionID string
38 Cron string
39 Message string
40 Location *time.Location
41}
42
43// Engine owns due-task claiming through one timer loop. Executors run outside
44// the lock, allowing independent session tasks to progress concurrently; the
45// host remains responsible for serializing work for one session.
46type Engine struct {
47 mu sync.Mutex
48 now func() time.Time
49 tasks map[string]Task
50 wake chan struct{}
51}
52
53// NewEngine creates an empty scheduler. now is injectable for tests; nil uses
54// time.Now.
55func NewEngine(now func() time.Time) *Engine {
56 if now == nil {
57 now = time.Now
58 }
59 return &Engine{now: now, tasks: make(map[string]Task), wake: make(chan struct{}, 1)}
60}
61
62// Add validates and records a new recurring task.
63func (e *Engine) Add(input NewTaskInput) (Task, error) {
64 if e == nil {
65 return Task{}, fmt.Errorf("scheduler is not initialized")
66 }
67 input.SessionID = strings.TrimSpace(input.SessionID)
68 input.Cron = strings.TrimSpace(input.Cron)
69 input.Message = strings.TrimSpace(input.Message)
70 if input.SessionID == "" {
71 return Task{}, fmt.Errorf("session ID is required")
72 }
73 if input.Message == "" {
74 return Task{}, fmt.Errorf("scheduled message is required")
75 }
76 if input.Location == nil {
77 return Task{}, fmt.Errorf("schedule timezone is required")
78 }
79 schedule, err := ParseCron(input.Cron, input.Location)
80 if err != nil {
81 return Task{}, err
82 }
83 now := e.now()
84 next := schedule.Next(now)
85 if next.IsZero() {
86 return Task{}, fmt.Errorf("cron has no occurrence in the next eight years")
87 }
88 task := Task{
89 ID: uuid.NewString(),
90 SessionID: input.SessionID,
91 Cron: input.Cron,
92 Timezone: input.Location.String(),
93 Message: input.Message,
94 CreatedAt: now,
95 NextRun: next,
96 schedule: schedule,
97 }
98 e.mu.Lock()
99 e.tasks[task.ID] = task
100 e.mu.Unlock()
101 e.signalWake()
102 return task, nil
103}
104
105// Cancel removes a scheduled task. It returns false when the ID is unknown.
106func (e *Engine) Cancel(id string) bool {
107 if e == nil {
108 return false
109 }
110 e.mu.Lock()
111 _, ok := e.tasks[id]
112 if ok {
113 delete(e.tasks, id)
114 }
115 e.mu.Unlock()
116 if ok {
117 e.signalWake()
118 }
119 return ok
120}
121
122// Get returns a task snapshot by ID.
123func (e *Engine) Get(id string) (Task, bool) {
124 if e == nil {
125 return Task{}, false
126 }
127 e.mu.Lock()
128 defer e.mu.Unlock()
129 task, ok := e.tasks[id]
130 return task, ok
131}
132
133// List returns snapshots ordered by next execution and then ID.
134func (e *Engine) List() []Task {
135 if e == nil {
136 return nil
137 }
138 e.mu.Lock()
139 tasks := make([]Task, 0, len(e.tasks))
140 for _, task := range e.tasks {
141 tasks = append(tasks, task)
142 }
143 e.mu.Unlock()
144 sort.Slice(tasks, func(i, j int) bool {
145 if tasks[i].NextRun.Equal(tasks[j].NextRun) {
146 return tasks[i].ID < tasks[j].ID
147 }
148 return tasks[i].NextRun.Before(tasks[j].NextRun)
149 })
150 return tasks
151}
152
153// Run waits for due tasks until ctx is cancelled. Each claimed task receives
154// ctx and executes asynchronously; cancellation prevents new callbacks and is
155// propagated to callbacks already in flight.
156func (e *Engine) Run(ctx context.Context, executor func(context.Context, Task) error) error {
157 if e == nil {
158 return fmt.Errorf("scheduler is not initialized")
159 }
160 if executor == nil {
161 return fmt.Errorf("scheduler executor is required")
162 }
163 for {
164 if ctx.Err() != nil {
165 return nil
166 }
167 now := e.now()
168 e.dispatchDue(ctx, now, executor)
169 next, ok := e.nextRun()
170 if !ok {
171 select {
172 case <-ctx.Done():
173 return nil
174 case <-e.wake:
175 }
176 continue
177 }
178 wait := time.Until(next)
179 if wait < 0 {
180 wait = 0
181 }
182 timer := time.NewTimer(wait)
183 select {
184 case <-ctx.Done():
185 if !timer.Stop() {
186 select {
187 case <-timer.C:
188 default:
189 }
190 }
191 return nil
192 case <-e.wake:
193 if !timer.Stop() {
194 select {
195 case <-timer.C:
196 default:
197 }
198 }
199 case <-timer.C:
200 }
201 }
202}
203
204// dispatchDue reserves every due task before launching an asynchronous
205// dispatch. A reservation prevents concurrent Run callers from selecting the
206// same task. dispatchTask claims it immediately before invoking the executor,
207// so Cancel can still win while a task is waiting to begin.
208func (e *Engine) dispatchDue(ctx context.Context, now time.Time, executor func(context.Context, Task) error) {
209 if e == nil || executor == nil || ctx.Err() != nil {
210 return
211 }
212 for _, task := range e.takeDue(now) {
213 go e.dispatchTask(ctx, task, executor)
214 }
215}
216
217func (e *Engine) takeDue(now time.Time) []Task {
218 e.mu.Lock()
219 defer e.mu.Unlock()
220 due := make([]Task, 0)
221 for id, task := range e.tasks {
222 if task.queued || task.running || task.NextRun.After(now) {
223 continue
224 }
225 task.queued = true
226 e.tasks[id] = task
227 due = append(due, task)
228 }
229 sort.Slice(due, func(i, j int) bool {
230 if due[i].NextRun.Equal(due[j].NextRun) {
231 return due[i].ID < due[j].ID
232 }
233 return due[i].NextRun.Before(due[j].NextRun)
234 })
235 return due
236}
237
238func (e *Engine) dispatchTask(ctx context.Context, queued Task, executor func(context.Context, Task) error) {
239 if ctx.Err() != nil {
240 e.releaseQueued(queued.ID, queued.NextRun)
241 return
242 }
243 e.mu.Lock()
244 task, ok := e.tasks[queued.ID]
245 if !ok || !task.queued || !task.NextRun.Equal(queued.NextRun) {
246 e.mu.Unlock()
247 return // cancelled or superseded before it was claimed
248 }
249 task.queued = false
250 task.running = true
251 e.tasks[task.ID] = task
252 e.mu.Unlock()
253 if ctx.Err() != nil {
254 e.finish(task.ID, task.NextRun, ctx.Err())
255 return
256 }
257 err := executor(ctx, task)
258 e.finish(task.ID, task.NextRun, err)
259}
260
261func (e *Engine) releaseQueued(id string, scheduledFor time.Time) {
262 e.mu.Lock()
263 task, ok := e.tasks[id]
264 if !ok || !task.queued || !task.NextRun.Equal(scheduledFor) {
265 e.mu.Unlock()
266 return
267 }
268 task.queued = false
269 e.tasks[id] = task
270 e.mu.Unlock()
271 e.signalWake()
272}
273
274func (e *Engine) finish(id string, scheduledFor time.Time, err error) {
275 e.mu.Lock()
276 task, ok := e.tasks[id]
277 if !ok || !task.NextRun.Equal(scheduledFor) {
278 e.mu.Unlock()
279 return // cancelled or replaced while the executor ran
280 }
281 task.running = false
282 task.LastRun = scheduledFor
283 task.RunCount++
284 if err != nil {
285 task.LastError = err.Error()
286 } else {
287 task.LastError = ""
288 }
289 // Calculate from completion time, rather than repeatedly from the prior
290 // deadline, so a sleeping machine or a slow task never creates a burst of
291 // stale catch-up executions.
292 task.NextRun = task.schedule.Next(e.now())
293 if task.NextRun.IsZero() {
294 delete(e.tasks, id)
295 } else {
296 e.tasks[id] = task
297 }
298 e.mu.Unlock()
299 e.signalWake()
300}
301
302func (e *Engine) nextRun() (time.Time, bool) {
303 e.mu.Lock()
304 defer e.mu.Unlock()
305 var next time.Time
306 for _, task := range e.tasks {
307 if task.queued || task.running {
308 continue
309 }
310 if next.IsZero() || task.NextRun.Before(next) {
311 next = task.NextRun
312 }
313 }
314 return next, !next.IsZero()
315}
316
317func (e *Engine) signalWake() {
318 select {
319 case e.wake <- struct{}{}:
320 default:
321 }
322}