Go语言中带超时机制的信号量实现指南

package main

import (
    "fmt"
    "time"
)

// NewSemaphore 创建一个指定容量的信号量
func NewSemaphore(capacity int) chan struct{} {
    return make(chan struct{}, capacity)
}

// Acquire 尝试获取信号量
func Acquire(sem chan struct{}) {
    sem <- struct{}{} // 阻塞直到有空闲位置
}

// Release 释放信号量
func Release(sem chan struct{}) {
    <-sem // 释放一个位置
}

func main() {
    sem := NewSemaphore(3) // 容量为3的信号量

    for i := 0; i < 5; i++ {
        go func(id int) {
            Acquire(sem)
            fmt.Printf("Worker %d acquired semaphore.\n", id)
            time.Sleep(time.Second) // 模拟工作
            fmt.Printf("Worker %d released semaphore.\n", id)
            Release(sem)
        }(i)
    }

    time.Sleep(5 * time.Second) // 等待所有协程完成
    fmt.Println("Main finished.")
}

package main

import (
    "fmt"
    "time"
)

// AcquireWithTimeout 尝试在指定超时时间内获取信号量
// 如果成功获取，返回 true；如果超时，返回 false。
func AcquireWithTimeout(sem chan struct{}, timeout time.Duration) bool {
    select {
    case sem <- struct{}{}:
        return true // 成功获取信号量
    case <-time.After(timeout):
        return false // 获取信号量超时
    }
}

// Release 释放信号量
func Release(sem chan struct{}) {
    <-sem
}

func main() {
    sem := make(chan struct{}, 2) // 容量为2的信号量

    for i := 0; i < 5; i++ {
        go func(id int) {
            fmt.Printf("Worker %d trying to acquire semaphore...\n", id)
            if AcquireWithTimeout(sem, 500*time.Millisecond) {
                fmt.Printf("Worker %d acquired semaphore.\n", id)
                time.Sleep(1 * time.Second) // 模拟工作
                fmt.Printf("Worker %d released semaphore.\n", id)
                Release(sem)
            } else {
                fmt.Printf("Worker %d failed to acquire semaphore within timeout.\n", id)
            }
        }(i)
    }

    time.Sleep(3 * time.Second)
    fmt.Println("Main finished.")
}

package main

import (
    "context"
    "fmt"
    "math/rand"
    "sync"
    "time"
)

// NewSemaphore 创建一个指定容量的信号量
func NewSemaphore(capacity int) chan struct{} {
    return make(chan struct{}, capacity)
}

// Worker 模拟一个需要执行的任务
func Worker(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup, taskID int) {
    defer wg.Done() // 任务完成时通知 WaitGroup

    // 确保无论任务如何结束，信号量都会被释放
    defer func() {
        fmt.Printf("Worker %d releasing semaphore.\n", taskID)
        <-sem
    }()

    // 模拟任务执行时间，随机在1到5秒之间
    workDuration := time.Duration(rand.Intn(5)+1) * time.Second
    fmt.Printf("Worker %d started, expected work duration: %s\n", taskID, workDuration)

    select {
    case <-time.After(workDuration):
        // 任务在超时前完成
        fmt.Printf("Worker %d completed successfully.\n", taskID)
    case <-ctx.Done():
        // 任务被外部上下文取消（例如，因为超时）
        fmt.Printf("Worker %d cancelled: %v\n", taskID, ctx.Err())
    }
}

func main() {
    rand.Seed(time.Now().UnixNano()) // 初始化随机数种子

    sem := NewSemaphore(3) // 信号量容量为3
    var wg sync.WaitGroup

    totalTasks := 10
    taskExecutionTimeout := 3 * time.Second // 每个任务的最大执行时间

    for i := 0; i < totalTasks; i++ {
        // 尝试获取信号量，如果获取失败，则不启动任务
        select {
        case sem <- struct{}{}:
            // 成功获取信号量
            wg.Add(1)
            // 为每个任务创建一个带超时的上下文
            ctx, cancel := context.WithTimeout(context.Background(), taskExecutionTimeout)
            go Worker(ctx, sem, &wg, i)
            // 注意：cancel() 函数在 Worker 协程结束后应该被调用，
            // 但因为 ctx 是为每个 Worker 独立创建的，且 Worker 结束后会释放资源，
            // 在此示例中，我们不显式调用 cancel，让 GC 处理。
            // 在更复杂的场景中，通常会在协程结束后显式调用 cancel()。
            _ = cancel // 避免 linter 警告 unused variable
        case <-time.After(1 * time.Second): // 信号量获取超时
            fmt.Printf("Task %d failed to acquire semaphore within 1 second timeout. Skipping.\n", i)
        }
    }

    wg.Wait() // 等待所有任务完成
    fmt.Println("All tasks finished.")
}