Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Better go-idiomatic way of writing this code?

Tags:

go

Playing around with go, I threw together this code:

package main

import "fmt"

const N = 10

func main() {
    ch := make(chan int, N)
    done := make(chan bool)

    for i := 0; i < N; i++ {
        go (func(n int, ch chan int, done chan bool) {
            for i := 0; i < N; i++ {
                ch <- n*N + i
            }
            done <- true
        })(i, ch, done)
    }

    numDone := 0
    for numDone < N {
        select {
        case i := <-ch:
            fmt.Println(i)
        case <-done:
            numDone++
        }
    }

    for {
        select {
        case i := <-ch:
            fmt.Println(i)
        default:
            return
        }
    }
}

Basically I have N channels doing some work and reporting it on the same channel -- I want to know when all the channels are done. So I have this other done channel that each worker goroutine sends a message on (message doesn't matter), and this causes main to count that thread as done. When the count gets to N, we're actually done.

Is this "good" go? Is there a more go-idiomatic way of doing this?

edit: To clarify a bit, I'm doubtful because the done channel seems to be doing a job that channel closing seems to be for, but of course I can't actually close the channel in any goroutine because all the routines share the same channel. So I'm using done to simulate a channel that does some kind of "buffered closing".

edit2: Original code wasn't really working since sometimes the done signal from a routine was read before the int it just put on ch. Needs a "cleanup" loop.

like image 796
alecbz Avatar asked Dec 08 '22 16:12

alecbz


2 Answers

Here is an idiomatic use of sync.WaitGroup for you to study

(playground link)

package main

import (
    "fmt"
    "sync"
)

const N = 10

func main() {
    ch := make(chan int, N)
    var wg sync.WaitGroup
    for i := 0; i < N; i++ {
        wg.Add(1)
        go func(n int) {
            defer wg.Done()
            for i := 0; i < N; i++ {
                ch <- n*N + i
            }
        }(i)
    }
    go func() {
        wg.Wait()
        close(ch)
    }()
    for i := range ch {
        fmt.Println(i)
    }
}

Note the use of closures in the two go routine definitions and note the second go statement to wait for all the routines to finish, then close the channel, so range can be used.

like image 165
Nick Craig-Wood Avatar answered Dec 11 '22 07:12

Nick Craig-Wood


looks like you want a sync.WaitGroup (http://golang.org/pkg/sync/#WaitGroup)

like image 22
cthom06 Avatar answered Dec 11 '22 08:12

cthom06