面试题：Go语言Context在并发任务超时控制中的复杂应用

Question

假设有一个Go程序，需要并发执行多个`heavyTask`函数（`func heavyTask(ctx context.Context) (result string, err error)`），每个任务的执行时间可能不同。要求当总执行时间超过3秒时，所有正在执行的`heavyTask`任务都能被正确取消，并且统计成功完成的任务数量和失败的任务数量。请详细描述实现方案并给出完整的Go代码。

Answer 1

package main

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

func heavyTask(ctx context.Context) (result string, err error) {
    // 模拟一个耗时任务
    select {
    case <-ctx.Done():
        return "", ctx.Err()
    case <-time.After(time.Second * 2):
        return "success", nil
    }
}

func main() {
    var successCount, failCount int
    var wg sync.WaitGroup

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

    tasks := []func() (string, error){
        func() (string, error) {
            result, err := heavyTask(ctx)
            if err != nil {
                failCount++
            } else {
                successCount++
            }
            return result, err
        },
        func() (string, error) {
            result, err := heavyTask(ctx)
            if err != nil {
                failCount++
            } else {
                successCount++
            }
            return result, err
        },
        // 可以添加更多任务
    }

    for _, task := range tasks {
        wg.Add(1)
        go func(t func() (string, error)) {
            defer wg.Done()
            _, _ = t()
        }(task)
    }

    wg.Wait()

    fmt.Printf("成功完成的任务数量: %d\n", successCount)
    fmt.Printf("失败的任务数量: %d\n", failCount)
}

1. 任务函数heavyTask：接受一个context.Context参数，通过select语句，在ctx.Done()通道收到信号（即上下文取消）时，任务取消；否则，模拟一个2秒的耗时操作。
2. 主函数main：
   • 初始化成功和失败任务计数变量successCount、failCount。
   • 使用context.WithTimeout创建一个带有3秒超时的上下文ctx和取消函数cancel，在函数结束时调用cancel确保资源释放。
   • 定义一个任务函数切片，每个任务函数调用heavyTask，并根据返回结果更新成功或失败计数。
   • 使用sync.WaitGroup等待所有任务完成。
   • 最后打印成功和失败的任务数量。