面试题：Go WaitGroup在分布式任务调度中的应用优化

Question

在分布式任务调度系统里，有大量的任务需要并发执行，每个任务执行时间不同且可能失败。使用WaitGroup等待所有任务完成时，可能会出现因部分任务长时间阻塞导致整体效率低下的问题。请分析如何优化这种情况，利用Go语言的特性和WaitGroup提出至少两种优化方案，并说明每种方案的优缺点。

Answer 1

方案一：设置任务超时

1. 实现方式：利用Go语言的context包来设置任务的执行超时。在启动每个任务时，传入带有超时的context.Context。如果任务在规定时间内未完成，context会通知任务停止执行。

package main

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

func task(ctx context.Context, wg *sync.WaitGroup, taskID int) {
    defer wg.Done()
    select {
    case <-ctx.Done():
        fmt.Printf("Task %d was cancelled due to timeout\n", taskID)
    case <-time.After(time.Duration(taskID) * time.Second):
        fmt.Printf("Task %d completed successfully\n", taskID)
    }
}

func main() {
    var wg sync.WaitGroup
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    for i := 1; i <= 5; i++ {
        wg.Add(1)
        go task(ctx, &wg, i)
    }

    wg.Wait()
}

2. 优点：
   • 有效防止长时间阻塞的任务影响整体效率，能控制任务执行的最长时间。
   • 优雅地处理任务取消，避免资源浪费。
3. 缺点：
   • 可能会误杀正常执行但耗时较长的任务，尤其是在设置的超时时间不合理的情况下。
   • 对于需要长时间执行的任务，如果因为超时被取消，可能需要额外的机制来恢复或重新执行。

方案二：任务分流与动态调整

1. 实现方式：将任务按照一定规则（如预估执行时间）分流到不同的协程组中。每个协程组使用独立的WaitGroup等待任务完成。同时，监控每个协程组的任务执行情况，动态调整任务分配。

package main

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

func shortTask(wg *sync.WaitGroup, taskID int) {
    defer wg.Done()
    time.Sleep(time.Duration(taskID) * time.Second)
    fmt.Printf("Short Task %d completed\n", taskID)
}

func longTask(wg *sync.WaitGroup, taskID int) {
    defer wg.Done()
    time.Sleep(time.Duration(taskID) * time.Second)
    fmt.Printf("Long Task %d completed\n", taskID)
}

func main() {
    var shortWG, longWG sync.WaitGroup

    for i := 1; i <= 3; i++ {
        if i <= 2 {
            shortWG.Add(1)
            go shortTask(&shortWG, i)
        } else {
            longWG.Add(1)
            go longTask(&longWG, i)
        }
    }

    go func() {
        shortWG.Wait()
        fmt.Println("All short tasks completed")
    }()

    go func() {
        longWG.Wait()
        fmt.Println("All long tasks completed")
    }()

    time.Sleep(5 * time.Second)
}

2. 优点：
   • 可以提高整体效率，因为短任务不会被长任务阻塞。
   • 具有较好的灵活性，通过动态调整任务分配，可以适应不同的任务负载。
3. 缺点：
   • 实现相对复杂，需要额外的逻辑来进行任务分流和动态调整。
   • 可能需要更多的系统资源来管理多个协程组。