Golang实现基础任务调度工具实例

context.Context

stopChan

package main

import (
    "context"
    "fmt"
    "sync"
    "time"
)

// Task 定义了一个要执行的任务的结构
type Task struct {
    ID        string
    Interval  time.Duration
    Handler   func(ctx context.Context) // 任务执行函数，接收一个context用于取消
    cancel    context.CancelFunc        // 用于取消当前任务的context
}

// Scheduler 是一个基础的任务调度器
type Scheduler struct {
    tasks map[string]*Task
    mu    sync.Mutex // 保护tasks map的并发访问，避免竞态条件
    wg    sync.WaitGroup // 等待所有任务goroutine完成，以便优雅停机
}

// NewScheduler 创建并返回一个新的Scheduler实例
func NewScheduler() *Scheduler {
    return &Scheduler{
        tasks: make(map[string]*Task),
    }
}

// AddTask 添加一个新任务到调度器。它会立即启动一个goroutine来执行这个任务。
func (s *Scheduler) AddTask(id string, interval time.Duration, handler func(ctx context.Context)) error {
    s.mu.Lock()
    defer s.mu.Unlock()

    if _, exists := s.tasks[id]; exists {
        return fmt.Errorf("任务ID %s 已存在，请使用唯一的ID", id)
    }

    // 为每个任务创建一个独立的context，以便单独控制其生命周期
    ctx, cancel := context.WithCancel(context.Background())

    task := &Task{
        ID:        id,
        Interval:  interval,
        Handler:   handler,
        cancel:    cancel,
    }
    s.tasks[id] = task

    s.wg.Add(1) // 增加WaitGroup计数，表示有一个任务正在运行
    go s.runTask(ctx, task)

    fmt.Printf(">> 调度器: 任务 '%s' 已添加并启动，间隔 %v\n", id, interval)
    return nil
}

// runTask 负责在一个独立的goroutine中执行任务的逻辑
func (s *Scheduler) runTask(ctx context.Context, task *Task) {
    defer s.wg.Done() // 任务goroutine退出时减少WaitGroup计数

    ticker := time.NewTicker(task.Interval)
    defer ticker.Stop() // 确保ticker在函数退出时被停止，避免资源泄露

    // 首次执行任务，这通常是一个好的实践
    fmt.Printf(">> 调度器: [%s] 任务首次执行...\n", task.ID)
    task.Handler(ctx)

    for {
        select {
        case <-ticker.C:
            // 避免在context被取消后还执行任务
            if ctx.Err() != nil {
                fmt.Printf(">> 调度器: [%s] Context已取消，停止执行。\n", task.ID)
                return
            }
            fmt.Printf(">> 调度器: [%s] 任务定时执行...\n", task.ID)
            task.Handler(ctx)
        case <-ctx.Done():
            // 接收到取消信号，优雅地退出任务goroutine
            fmt.Printf(">> 调度器: [%s] 任务接收到取消信号，正在停止...\n", task.ID)
            return
        }
    }
}

// RemoveTask 从调度器中移除一个任务
func (s *Scheduler) RemoveTask(id string) error {
    s.mu.Lock()
    defer s.mu.Unlock()

    task, exists := s.tasks[id]
    if !exists {
        return fmt.Errorf("任务ID %s 不存在，无法移除", id)
    }

    // 发送取消信号给任务的context
    task.cancel()
    delete(s.tasks, id) // 从map中移除任务

    fmt.Printf(">> 调度器: 任务 '%s' 已标记为移除，等待其停止...\n", id)
    return nil
}

// Wait 等待所有正在运行的任务goroutine完成。通常在程序退出前调用，以确保所有任务都已优雅停机。
func (s *Scheduler) Wait() {
    s.wg.Wait()
    fmt.Println(">> 调度器: 所有任务已停止。")
}

func main() {
    scheduler := NewScheduler()

    // 添加一个每隔2秒执行一次的任务
    scheduler.AddTask("clean-logs", 2*time.Second, func(ctx context.Context) {
        fmt.Println(">>> 任务 [clean-logs]: 正在清理日志...")
        // 模拟一个可能耗时或出错的操作
        time.Sleep(500 * time.Millisecond)
        if ctx.Err() != nil { // 在任务执行过程中检查取消信号
            fmt.Println(">>> 任务 [clean-logs]: 在执行中被取消了！")
            return
        }
        fmt.Println(">>> 任务 [clean-logs]: 清理完成。")
    })

    // 添加一个每隔3秒执行一次的数据同步任务
    scheduler.AddTask("sync-data", 3*time.Second, func(ctx context.Context) {
        fmt.Println("--- 任务 [sync-data]: 正在同步数据...")
        time.Sleep(1 * time.Second)
        if ctx.Err() != nil {
            fmt.Println("--- 任务 [sync-data]: 在执行中被取消了！")
            return
        }
        fmt.Println("--- 任务 [sync-data]: 数据同步完成。")
    })

    // 让调度器运行一段时间，观察任务执行情况
    time.Sleep(7 * time.Second)

    // 模拟需要移除一个任务的场景
    fmt.Println("\n--- 7秒后，移除 'clean-logs' 任务 ---")
    scheduler.RemoveTask("clean-logs")

    // 再运行一段时间，看看剩下的任务是否还在正常执行
    time.Sleep(5 * time.Second)

    fmt.Println("\n--- 5秒后，等待所有任务彻底停止 ---")
    scheduler.Wait() // 等待所有任务goroutine完成

    fmt.Println("主程序退出。")
}

context.Context

sync.Mutex

sync.WaitGroup

tasks

sync.Mutex

tasks