mirror of
https://github.com/tiennm99/miti99bot.git
synced 2026-09-08 18:17:54 +00:00
feat(cron): add in-process cron scheduler
Add robfig/cron-based task scheduler with per-UTC-date idempotency guard. Enables scheduled jobs (e.g., daily price checks) without external scheduler dependency. Module interface updated to support cron registration.
This commit is contained in:
@@ -0,0 +1,81 @@
|
||||
// Package cron runs module crons in-process for the self-hosted (long-lived
|
||||
// container) deployment. On AWS, EventBridge Scheduler hit /cron/{name}; off
|
||||
// AWS there is no external trigger, so this scheduler reads each registered
|
||||
// cron's Schedule field and fires its handler on time, in UTC.
|
||||
package cron
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime/debug"
|
||||
"time"
|
||||
|
||||
"github.com/robfig/cron/v3"
|
||||
|
||||
"github.com/tiennm99/miti99bot/internal/log"
|
||||
"github.com/tiennm99/miti99bot/internal/modules"
|
||||
)
|
||||
|
||||
// cronTimeout caps a single in-process cron fire. Defined locally (rather than
|
||||
// imported from internal/server) to avoid a server→cron layering inversion;
|
||||
// it intentionally matches internal/server's defaultCronTimeout so the
|
||||
// in-process path and the HTTP /cron path share the same 60s budget.
|
||||
const cronTimeout = 60 * time.Second
|
||||
|
||||
// Run starts an in-process scheduler that fires every registered cron with a
|
||||
// non-empty Schedule. Each fire is wrapped with a timeout, panic recovery, and
|
||||
// structured logging — modules.DispatchScheduled provides none of these (it is
|
||||
// only a registry lookup + handler call), so the wrapping lives here, mirroring
|
||||
// the server-package cronHandler.
|
||||
//
|
||||
// A malformed schedule string is fatal (returned as an error) so a config typo
|
||||
// fails fast at startup rather than silently never firing. The returned stop
|
||||
// function halts the scheduler and waits for any in-flight job to finish; it is
|
||||
// safe to call exactly once (e.g. via defer).
|
||||
func Run(ctx context.Context, reg *modules.Registry) (stop func(), err error) {
|
||||
c := cron.New(cron.WithLocation(time.UTC))
|
||||
|
||||
registered := 0
|
||||
for _, cr := range reg.Crons() {
|
||||
if cr.Schedule == "" {
|
||||
log.Warn("cron has empty schedule; it will never fire", "name", cr.Name)
|
||||
continue
|
||||
}
|
||||
name := cr.Name // capture per iteration
|
||||
if _, err := c.AddFunc(cr.Schedule, func() { fire(ctx, name, reg) }); err != nil {
|
||||
return func() {}, fmt.Errorf("cron %q has invalid schedule %q: %w", cr.Name, cr.Schedule, err)
|
||||
}
|
||||
registered++
|
||||
}
|
||||
|
||||
c.Start()
|
||||
log.Info("cron scheduler started", "registered", registered)
|
||||
|
||||
stop = func() {
|
||||
// c.Stop() returns a context that is done once running jobs complete.
|
||||
stopCtx := c.Stop()
|
||||
<-stopCtx.Done()
|
||||
}
|
||||
return stop, nil
|
||||
}
|
||||
|
||||
// fire dispatches one scheduled cron with its own timeout and panic recovery so
|
||||
// a slow or panicking handler never wedges or kills the scheduler.
|
||||
func fire(ctx context.Context, name string, reg *modules.Registry) {
|
||||
runCtx, cancel := context.WithTimeout(ctx, cronTimeout)
|
||||
defer cancel()
|
||||
|
||||
log.Info("cron triggered", "source", "scheduler", "name", name)
|
||||
defer func() {
|
||||
if rec := recover(); rec != nil {
|
||||
log.Error("cron handler panic",
|
||||
"source", "scheduler",
|
||||
"name", name,
|
||||
"panic", rec,
|
||||
"stack", string(debug.Stack()))
|
||||
}
|
||||
}()
|
||||
if err := modules.DispatchScheduled(runCtx, name, reg); err != nil {
|
||||
log.Error("cron failed", "source", "scheduler", "name", name, "err", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package cron
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/tiennm99/miti99bot/internal/modules"
|
||||
"github.com/tiennm99/miti99bot/internal/storage"
|
||||
)
|
||||
|
||||
// buildReg builds a registry with a single module exposing one cron whose
|
||||
// handler runs onFire. schedule is the cron expression under test.
|
||||
func buildReg(t *testing.T, schedule string, onFire func()) *modules.Registry {
|
||||
t.Helper()
|
||||
factories := map[string]modules.Factory{
|
||||
"ticker": func(_ modules.Deps) modules.Module {
|
||||
return modules.Module{
|
||||
Name: "ticker",
|
||||
Crons: []modules.Cron{{
|
||||
Name: "tick",
|
||||
Schedule: schedule,
|
||||
Handler: func(_ context.Context, _ modules.Deps) error {
|
||||
onFire()
|
||||
return nil
|
||||
},
|
||||
}},
|
||||
}
|
||||
},
|
||||
}
|
||||
reg, err := modules.Build([]string{"ticker"}, factories, storage.NewMemoryProvider(), modules.BuildOptions{})
|
||||
if err != nil {
|
||||
t.Fatalf("modules.Build: %v", err)
|
||||
}
|
||||
return reg
|
||||
}
|
||||
|
||||
func TestRun_FiresHandlerOnSchedule(t *testing.T) {
|
||||
fired := make(chan struct{}, 1)
|
||||
reg := buildReg(t, "@every 1s", func() {
|
||||
select {
|
||||
case fired <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
})
|
||||
|
||||
stop, err := Run(context.Background(), reg)
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
defer stop()
|
||||
|
||||
select {
|
||||
case <-fired:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("cron handler did not fire within 5s")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_BadScheduleIsFatal(t *testing.T) {
|
||||
reg := buildReg(t, "not-a-cron-expression", func() {})
|
||||
if _, err := Run(context.Background(), reg); err == nil {
|
||||
t.Fatal("Run with invalid schedule: got nil error, want a parse error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRun_StopHaltsScheduler(t *testing.T) {
|
||||
reg := buildReg(t, "@every 1s", func() {})
|
||||
stop, err := Run(context.Background(), reg)
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
stop()
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("stop() did not return within 5s")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRun_RecoversHandlerPanic asserts a panicking cron handler does not crash
|
||||
// the scheduler; the next tick still fires.
|
||||
func TestRun_RecoversHandlerPanic(t *testing.T) {
|
||||
calls := make(chan struct{}, 4)
|
||||
reg := buildReg(t, "@every 1s", func() {
|
||||
calls <- struct{}{}
|
||||
panic("boom")
|
||||
})
|
||||
stop, err := Run(context.Background(), reg)
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
defer stop()
|
||||
|
||||
// Two separate fires prove the scheduler survived the first panic.
|
||||
for i := 0; i < 2; i++ {
|
||||
select {
|
||||
case <-calls:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatalf("expected fire #%d after panic recovery", i+1)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
|
||||
"github.com/tiennm99/miti99bot/internal/log"
|
||||
"github.com/tiennm99/miti99bot/internal/modules"
|
||||
"github.com/tiennm99/miti99bot/internal/storage"
|
||||
)
|
||||
|
||||
// terminalKind classifies a permanent send failure by blast radius.
|
||||
@@ -76,15 +77,24 @@ func classifyTerminal(err error) terminalKind {
|
||||
return terminalNone
|
||||
}
|
||||
|
||||
// dailyPushCronName is the cron route segment + EventBridge schedule key.
|
||||
// dailyPushCronName is the cron route segment + in-process scheduler key.
|
||||
// Must match the regex in internal/server/router.go (^[a-z0-9_]{1,32}$).
|
||||
const dailyPushCronName = "lolschedule_daily_push"
|
||||
|
||||
// dailyPushSchedule is documentation only — the real fire time lives in
|
||||
// EventBridge Scheduler. Cron expression is UTC;
|
||||
// 01:00 UTC == 08:00 ICT.
|
||||
// dailyPushSchedule drives the in-process scheduler (internal/cron) on
|
||||
// self-host; it was also the documented EventBridge time on AWS. Cron
|
||||
// expression is UTC; 01:00 UTC == 08:00 ICT.
|
||||
const dailyPushSchedule = "0 1 * * *"
|
||||
|
||||
// lastPushDateKey records the UTC date (YYYY-MM-DD) of the most recent
|
||||
// completed daily push. The handler claims this key before fanning out and
|
||||
// no-ops if it is already today's date, making the push idempotent per UTC
|
||||
// date. This defends against every double-fire window — cutover overlap
|
||||
// (EventBridge still live while the container's scheduler runs), rolling
|
||||
// deploys that briefly run two containers, and operator misconfiguration —
|
||||
// none of which a single trigger source can prevent.
|
||||
const lastPushDateKey = "daily_push:last_date"
|
||||
|
||||
// telegramRateLimitThreshold is the subscriber count above which we throttle
|
||||
// sends to stay under Telegram's global 30 msg/sec cap. Below it we send hot.
|
||||
const telegramRateLimitThreshold = 30
|
||||
@@ -119,6 +129,50 @@ func (s *state) dailyPushHandler(ctx context.Context, deps modules.Deps) error {
|
||||
return runDailyPush(ctx, s, deps.Bot)
|
||||
}
|
||||
|
||||
// claimDailyPush atomically records that today's UTC push is happening and
|
||||
// reports whether THIS caller won the claim. It is the idempotency primitive
|
||||
// for the daily push: a winner proceeds to fan out; a loser (another trigger
|
||||
// already claimed today) returns false and sends nothing.
|
||||
//
|
||||
// The claim is a compare-and-swap on lastPushDateKey so two simultaneous
|
||||
// triggers cannot both win. Every real KV backend implements
|
||||
// CompareAndSwapStore; if a bare store without CAS is used (only in narrow
|
||||
// tests), it falls back to a plain Put — losing the atomic guarantee but
|
||||
// preserving the "no-op if already today" behaviour.
|
||||
func claimDailyPush(ctx context.Context, kv storage.KVStore, today string) (bool, error) {
|
||||
current, err := kv.Get(ctx, lastPushDateKey)
|
||||
switch {
|
||||
case err == nil:
|
||||
if string(current) == today {
|
||||
return false, nil // already pushed today
|
||||
}
|
||||
case errors.Is(err, storage.ErrNotFound):
|
||||
current = nil // never pushed
|
||||
default:
|
||||
return false, err
|
||||
}
|
||||
|
||||
cas, ok := kv.(storage.CompareAndSwapStore)
|
||||
if !ok {
|
||||
if err := kv.Put(ctx, lastPushDateKey, []byte(today)); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
var expected []byte
|
||||
if len(current) > 0 {
|
||||
expected = current
|
||||
}
|
||||
if err := cas.CompareAndSwap(ctx, lastPushDateKey, expected, []byte(today)); err != nil {
|
||||
if errors.Is(err, storage.ErrConflict) {
|
||||
return false, nil // another trigger claimed today first
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// runDailyPush is the testable core: fetch subscribers, fetch today's matches,
|
||||
// fan out to every subscriber. Per-chat send failures are logged but do not
|
||||
// abort the batch — one bad chat does not deny the rest.
|
||||
@@ -144,6 +198,20 @@ func runDailyPush(ctx context.Context, s *state, sender messageSender) error {
|
||||
filtered := FilterMajor(events)
|
||||
text := RenderToday(filtered, from)
|
||||
|
||||
// Idempotency gate: claim today's push before sending. A lost claim means
|
||||
// another trigger already pushed (or is pushing) for this UTC date, so we
|
||||
// send nothing. Placed after the fetch so a transient fetch failure does
|
||||
// not consume the day's claim.
|
||||
today := s.now().UTC().Format("2006-01-02")
|
||||
won, err := claimDailyPush(ctx, s.kv, today)
|
||||
if err != nil {
|
||||
return fmt.Errorf("lolschedule daily push: claim date: %w", err)
|
||||
}
|
||||
if !won {
|
||||
log.Info("lolschedule daily push: already pushed today, skipping", "date", today)
|
||||
return nil
|
||||
}
|
||||
|
||||
throttle := len(subs) > telegramRateLimitThreshold
|
||||
var sent, failed int
|
||||
deadChats := map[int64]struct{}{}
|
||||
|
||||
@@ -167,6 +167,33 @@ func TestRunDailyPush_ForwardsMessageThreadID(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunDailyPush_IdempotentPerDate locks in the double-fire guard: invoking
|
||||
// the handler twice on the same UTC date sends each subscriber exactly one
|
||||
// digest. Defends against cutover overlap, rolling-deploy overlap, and operator
|
||||
// misconfiguration (all double-fire windows the daily push must survive).
|
||||
func TestRunDailyPush_IdempotentPerDate(t *testing.T) {
|
||||
s := newTestState(t)
|
||||
seedFreshCache(t, s.kv, nil)
|
||||
|
||||
chatIDs := []int64{100, 200, 300}
|
||||
for _, id := range chatIDs {
|
||||
if _, err := addSubscriber(context.Background(), s.kv, id, 0); err != nil {
|
||||
t.Fatalf("addSubscriber %d: %v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
sender := &fakeSender{}
|
||||
for i := 0; i < 2; i++ {
|
||||
if err := runDailyPush(context.Background(), s, sender); err != nil {
|
||||
t.Fatalf("runDailyPush call %d: %v", i+1, err)
|
||||
}
|
||||
}
|
||||
if len(sender.calls) != len(chatIDs) {
|
||||
t.Errorf("two same-date pushes sent %d messages, want %d (one per subscriber)",
|
||||
len(sender.calls), len(chatIDs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunDailyPush_PartialFailureContinues(t *testing.T) {
|
||||
s := newTestState(t)
|
||||
seedFreshCache(t, s.kv, nil)
|
||||
|
||||
@@ -28,9 +28,11 @@ const (
|
||||
// purely for logging/metrics, not flow control.
|
||||
type CommandHandler func(ctx context.Context, b *bot.Bot, update *models.Update) error
|
||||
|
||||
// CronHandler runs when EventBridge Scheduler hits /cron/{name}. Crons receive the
|
||||
// per-module-prefixed Deps via the registry; handlers should not capture the
|
||||
// base Deps from the factory closure or KV writes will collide across modules.
|
||||
// CronHandler runs when a cron fires — driven by the in-process scheduler
|
||||
// (internal/cron) on self-host, or by a POST to /cron/{name} for manual
|
||||
// triggers. Crons receive the per-module-prefixed Deps via the registry;
|
||||
// handlers should not capture the base Deps from the factory closure or KV
|
||||
// writes will collide across modules.
|
||||
type CronHandler func(ctx context.Context, deps Deps) error
|
||||
|
||||
// Command is a single Telegram bot command exposed by a module.
|
||||
@@ -43,7 +45,7 @@ type Command struct {
|
||||
|
||||
// Cron is a single scheduled job exposed by a module.
|
||||
type Cron struct {
|
||||
Schedule string // documentation only; real schedule lives in EventBridge Scheduler
|
||||
Schedule string // 5-field cron expr (UTC); the in-process scheduler fires the handler on it
|
||||
Name string // unique within module
|
||||
Handler CronHandler // required
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user