WaitGroup uses Atomic under the hood for counting. WaitGroup can accept work weight. How would you rewrite the semaphore so that methods Lock(weight int), TryLock(weight int), and Release(weight int) accept weight? The goroutine acquires resources immediately for the entire operation, then releases a little after each stage (IO-bound, CPU-bound). How would you implement this?

To implement a semaphore with a weight in the Lock, TryLock, and Release methods, you can use an atomic counter that reflects the current number of occupied resources and a maximum resource limit. A goroutine attempting to acquire resources tries to increase the counter by the weight, provided it does not exceed the limit afterward. When releasing resources, the counter is decreased by the weight.

Example implementation in Go using atomic and sync.Mutex for waiting:

import (
    "sync"
    "sync/atomic"
)

type WeightedSemaphore struct {
    capacity int64
    current  int64
    mu       sync.Mutex
    cond     *sync.Cond
}

func NewWeightedSemaphore(capacity int64) *WeightedSemaphore {
    ws := &WeightedSemaphore{capacity: capacity}
    ws.cond = sync.NewCond(&ws.mu)
    return ws
}

func (ws *WeightedSemaphore) Lock(weight int64) {
    ws.mu.Lock()
    defer ws.mu.Unlock()
    for {
        current := atomic.LoadInt64(&ws.current)
        if current+weight <= ws.capacity {
            atomic.AddInt64(&ws.current, weight)
            return
        }
        ws.cond.Wait()
    }
}

func (ws *WeightedSemaphore) TryLock(weight int64) bool {
    ws.mu.Lock()
    defer ws.mu.Unlock()
    current := atomic.LoadInt64(&ws.current)
    if current+weight <= ws.capacity {
        atomic.AddInt64(&ws.current, weight)
        return true
    }
    return false
}

func (ws *WeightedSemaphore) Release(weight int64) {
    ws.mu.Lock()
    atomic.AddInt64(&ws.current, -weight)
    ws.mu.Unlock()
    ws.cond.Signal()
}

Thus, a goroutine can acquire resources with a certain weight and release part of the resources after each stage (for example, IO-bound, CPU-bound) by calling Release with the corresponding weight. This allows flexible resource management and avoids deadlocks if resources are insufficient.