Go - semaphore

 

 

 

 

 

 

 

 

 

package main

import (
    "context"
    "fmt"
    "os"
    "strconv"
    "time"

    "golang.org/x/sync/semaphore"
)

func worker(n int) int {
    square := n * n
    time.Sleep(time.Second)
    return square
}

func main() {
    if len(os.Args) != 2 {
        fmt.Println("Need #jobs!")
        os.Exit(1)
    }

    var workers = 4
    var sem = semaphore.NewWeighted(int64(workers))

    nJobs, err := strconv.Atoi(os.Args[1])
    if err != nil {
        fmt.Println(err)
        os.Exit(1)
    }

    // Where to store the results
    var results = make([]int, nJobs)
    // Needed by Acquire()
    ctx := context.TODO()
    for i := range results {
        err = sem.Acquire(ctx, 1)
        if err != nil {
            fmt.Println("Cannot acquire semaphore:", err)
            break
        }

        go func(i int) {
            defer sem.Release(1)
            temp := worker(i)
            results[i] = temp
        }(i)
    }

    err = sem.Acquire(ctx, int64(workers))
    if err != nil {
        fmt.Println(err)
    }

    for i, v := range results {
        fmt.Println(i, "->", v)
    }
}

 

zzh@ZZHPC:/zdata/Github/ztest$ go run main.go 5
0 -> 0
1 -> 1
2 -> 4
3 -> 9
4 -> 16