当前位置：首页 > article >正文

40分钟学 Go 语言高并发：超时控制与取消机制

article 2024/11/29 6:45:24

超时控制与取消机制

一、知识要点概述

知识模块	重要程度	掌握要求
超时策略	⭐⭐⭐⭐⭐	深入理解context包的超时机制，掌握不同场景下的超时控制方法
取消传播	⭐⭐⭐⭐⭐	熟练运用context的取消传播机制，实现优雅的任务取消
资源清理	⭐⭐⭐⭐	掌握goroutine退出时的资源清理方法，避免资源泄露
级联取消	⭐⭐⭐⭐	理解并实现多层级的取消控制，确保子任务能正确响应取消信号

二、详细内容

1. 超时策略

超时控制是高并发系统中非常重要的一环，它能够防止系统资源被长时间占用，提高系统的可用性。

package main

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

// 模拟一个耗时操作
func longRunningOperation(ctx context.Context) error {
    // 创建一个channel用于接收操作结果
    done := make(chan bool)
    
    go func() {
        // 模拟耗时操作
        time.Sleep(2 * time.Second)
        done <- true
    }()
    
    // 使用select实现超时控制
    select {
    case <-done:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

// 使用context.WithTimeout实现超时控制
func withTimeoutExample() {
    // 创建一个1秒超时的context
    ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
    defer cancel()
    
    if err := longRunningOperation(ctx); err != nil {
        fmt.Printf("操作超时: %v\n", err)
        return
    }
    fmt.Println("操作成功完成")
}

// 使用time.After实现超时控制
func withTimeAfterExample() {
    done := make(chan bool)
    
    go func() {
        time.Sleep(2 * time.Second)
        done <- true
    }()
    
    select {
    case <-done:
        fmt.Println("操作成功完成")
    case <-time.After(1 * time.Second):
        fmt.Println("操作超时")
    }
}

// 带有重试机制的超时控制
func withRetryTimeout(maxRetries int, timeout time.Duration) error {
    for i := 0; i < maxRetries; i++ {
        ctx, cancel := context.WithTimeout(context.Background(), timeout)
        err := longRunningOperation(ctx)
        cancel()
        
        if err == nil {
            return nil
        }
        
        fmt.Printf("第%d次重试失败: %v\n", i+1, err)
        time.Sleep(time.Second) // 重试间隔
    }
    return fmt.Errorf("达到最大重试次数")
}

func main() {
    fmt.Println("=== 使用context.WithTimeout示例 ===")
    withTimeoutExample()
    
    fmt.Println("\n=== 使用time.After示例 ===")
    withTimeAfterExample()
    
    fmt.Println("\n=== 带重试机制的超时控制示例 ===")
    err := withRetryTimeout(3, time.Second)
    if err != nil {
        fmt.Printf("最终结果: %v\n", err)
    }
}

package main

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

// 模拟数据处理任务
type Task struct {
    ID int
    ProcessTime time.Duration
}

// 工作协程
func worker(ctx context.Context, id int, tasks <-chan Task, wg *sync.WaitGroup) {
    defer wg.Done()
    
    for {
        select {
        case task, ok := <-tasks:
            if !ok {
                fmt.Printf("Worker %d: 通道已关闭，退出\n", id)
                return
            }
            
            // 检查是否已取消
            select {
            case <-ctx.Done():
                fmt.Printf("Worker %d: 收到取消信号，停止处理任务\n", id)
                return
            default:
                // 继续处理任务
                fmt.Printf("Worker %d: 开始处理任务 %d\n", id, task.ID)
                time.Sleep(task.ProcessTime)
                fmt.Printf("Worker %d: 完成任务 %d\n", id, task.ID)
            }
            
        case <-ctx.Done():
            fmt.Printf("Worker %d: 收到取消信号，退出\n", id)
            return
        }
    }
}

// 任务分发器
func dispatcher(ctx context.Context, numWorkers int) {
    var wg sync.WaitGroup
    tasks := make(chan Task, numWorkers)
    
    // 启动工作协程
    for i := 0; i < numWorkers; i++ {
        wg.Add(1)
        go worker(ctx, i, tasks, &wg)
    }
    
    // 模拟任务生成
    taskID := 0
    for {
        select {
        case <-ctx.Done():
            close(tasks)
            wg.Wait()
            fmt.Println("调度器: 所有工作协程已退出")
            return
        default:
            // 生成新任务
            taskID++
            task := Task{
                ID: taskID,
                ProcessTime: time.Duration(taskID*100) * time.Millisecond,
            }
            tasks <- task
            time.Sleep(200 * time.Millisecond) // 控制任务生成速率
        }
    }
}

func main() {
    // 创建可取消的context
    ctx, cancel := context.WithCancel(context.Background())
    
    // 启动调度器
    go dispatcher(ctx, 3)
    
    // 运行一段时间后取消
    time.Sleep(2 * time.Second)
    fmt.Println("主程序: 发送取消信号")
    cancel()
    
    // 等待一段时间确保清理完成
    time.Sleep(time.Second)
    fmt.Println("主程序: 退出")
}

package main

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

// 模拟数据库连接
type DBConnection struct {
    id     int
    closed bool
    mu     sync.Mutex
}

func (db *DBConnection) Query(ctx context.Context) error {
    select {
    case <-ctx.Done():
        return ctx.Err()
    case <-time.After(500 * time.Millisecond):
        return nil
    }
}

func (db *DBConnection) Close() error {
    db.mu.Lock()
    defer db.mu.Unlock()
    
    if !db.closed {
        fmt.Printf("关闭数据库连接 %d\n", db.id)
        db.closed = true
    }
    return nil
}

// 资源管理器
type ResourceManager struct {
    connections []*DBConnection
    mu          sync.Mutex
}

func NewResourceManager() *ResourceManager {
    return &ResourceManager{
        connections: make([]*DBConnection, 0),
    }
}

func (rm *ResourceManager) AcquireConnection() *DBConnection {
    rm.mu.Lock()
    defer rm.mu.Unlock()
    
    conn := &DBConnection{
        id: len(rm.connections),
    }
    rm.connections = append(rm.connections, conn)
    fmt.Printf("创建新的数据库连接 %d\n", conn.id)
    return conn
}

func (rm *ResourceManager) ReleaseAll() {
    rm.mu.Lock()
    defer rm.mu.Unlock()
    
    for _, conn := range rm.connections {
        conn.Close()
    }
    rm.connections = nil
}

// 工作任务
func doWork(ctx context.Context, rm *ResourceManager) error {
    // 获取数据库连接
    conn := rm.AcquireConnection()
    
    // 确保在函数退出时释放资源
    defer func() {
        if err := conn.Close(); err != nil {
            fmt.Printf("关闭连接失败: %v\n", err)
        }
    }()
    
    // 执行查询
    return conn.Query(ctx)
}

func main() {
    rm := NewResourceManager()
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    
    var wg sync.WaitGroup
    
    // 启动多个工作协程
    for i := 0; i < 3; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            
            for j := 0; j < 2; j++ {
                select {
                case <-ctx.Done():
                    fmt.Printf("工作协程 %d: 收到取消信号\n", id)
                    return
                default:
                    if err := doWork(ctx, rm); err != nil {
                        fmt.Printf("工作协程 %d: 执行失败 - %v\n", id, err)
                        return
                    }
                    fmt.Printf("工作协程 %d: 完成一次查询\n", id)
                }
            }
        }(i)
    }
    
    // 等待所有工作完成
    wg.Wait()
    
    // 清理所有资源
    fmt.Println("清理所有资源")
    rm.ReleaseAll()
}

package main

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

// Pipeline 代表一个处理管道
type Pipeline struct {
    name     string
    process  func(context.Context) error
    children []*Pipeline
    wg       sync.WaitGroup
}

// NewPipeline 创建新的处理管道
func NewPipeline(name string, process func(context.Context) error) *Pipeline {
    return &Pipeline{
        name:     name,
        process:  process,
        children: make([]*Pipeline, 0),
    }
}

// AddChild 添加子管道
func (p *Pipeline) AddChild(child *Pipeline) {
    p.children = append(p.children, child)
}

// Execute 执行当前管道及其所有子管道
func (p *Pipeline) Execute(ctx context.Context) error {
    // 创建一个子context，用于控制子管道
    childCtx, cancel := context.WithCancel(ctx)
    defer cancel() // 确保所有子管道都会被取消

    // 启动所有子管道
    errChan := make(chan error, len(p.children))
    for _, child := range p.children {
        p.wg.Add(1)
        go func(pipeline *Pipeline) {
            defer p.wg.Done()
            if err := pipeline.Execute(childCtx); err != nil {
                errChan <- err
                cancel() // 如果有错误发生，取消所有其他子管道
            }
        }(child)
    }

    // 执行当前管道的处理函数
    processDone := make(chan error, 1)
    go func() {
        processDone <- p.process(childCtx)
    }()

    // 等待处理完成或者context取消
    select {
    case err := <-processDone:
        if err != nil {
            cancel() // 如果处理出错，取消所有子管道
            fmt.Printf("[%s] 处理错误: %v\n", p.name, err)
            return err
        }
    case err := <-errChan:
        fmt.Printf("[%s] 子管道错误: %v\n", p.name, err)
        return err
    case <-ctx.Done():
        fmt.Printf("[%s] 被取消\n", p.name)
        return ctx.Err()
    }

    // 等待所有子管道完成
    waitCh := make(chan struct{})
    go func() {
        p.wg.Wait()
        close(waitCh)
    }()

    select {
    case <-waitCh:
        fmt.Printf("[%s] 及其所有子管道已完成\n", p.name)
        return nil
    case <-ctx.Done():
        fmt.Printf("[%s] 等待子管道完成时被取消\n", p.name)
        return ctx.Err()
    }
}

// 模拟处理函数，带有随机延时
func createProcessor(name string, duration time.Duration) func(context.Context) error {
    return func(ctx context.Context) error {
        fmt.Printf("[%s] 开始处理\n", name)
        select {
        case <-time.After(duration):
            fmt.Printf("[%s] 处理完成\n", name)
            return nil
        case <-ctx.Done():
            fmt.Printf("[%s] 处理被取消\n", name)
            return ctx.Err()
        }
    }
}

func main() {
    // 创建根管道和子管道
    root := NewPipeline("ROOT", createProcessor("ROOT", 1*time.Second))
    
    // 第一层子管道
    stage1A := NewPipeline("1A", createProcessor("1A", 2*time.Second))
    stage1B := NewPipeline("1B", createProcessor("1B", 1500*time.Millisecond))
    root.AddChild(stage1A)
    root.AddChild(stage1B)
    
    // 第二层子管道
    stage2A := NewPipeline("2A", createProcessor("2A", 1*time.Second))
    stage2B := NewPipeline("2B", createProcessor("2B", 800*time.Millisecond))
    stage1A.AddChild(stage2A)
    stage1B.AddChild(stage2B)
    
    // 第三层子管道
    stage3A := NewPipeline("3A", createProcessor("3A", 500*time.Millisecond))
    stage2A.AddChild(stage3A)

    // 创建根context，设置3秒超时
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    // 启动处理管道
    fmt.Println("开始执行管道...")
    startTime := time.Now()
    
    if err := root.Execute(ctx); err != nil {
        fmt.Printf("管道执行失败: %v\n", err)
    }
    
    fmt.Printf("总执行时间: %v\n", time.Since(startTime))
}