Create a WaitGroup by yourself

If you had written any concurrency code in Go. I think you could seem sync.WaitGroup before. And today point is to focus on creating a waitgroup by channel trick.

How?

First, you need a channel without the buffer.

func main() {
    wait := make(chan struct{})
}

Then get something from it so if there was no value in wait, the process keeps going. But if just trying to get something from an empty channel. You will get blocking and cause deadlock. Then go will panic it. So we have to change the code a little bit but also be more extendable for next iteration.

// ...
n := 0
wait := make(chan struct{})

for i := 0; i < n; i++ {
    <-wait
}
// ...

Now let’s create works loop.

import (
    "time"
)
// ...
n := 10000
wait := make(chan struct{})

for i := 0; i < n; i++ {
    time.Sleep(time.Microsecond)
    wait <- struct{}{}
}

for i := 0; i < n; i++ {
    <-wait
}
// ...

ps. This code has a little bug(you can try to find it, or read the answer at the end)

Now we can see it work. The reason for this code can work is because size n is our expected amount of workers. After each worker done their jobs. They will send something(here is struct{}{}, but exactly it doesn’t matter thing) into our wait channel. We only read n things from wait.

So after n things are read. We won’t be blocked anymore even wait got the new thing. Else we have to wait wait.

Whole code dependent on this fact. Having this knowledge, we can create ours WaitGroup now.

type waitGroup struct {
    n    int
    wait chan struct{}
}

func WaitGroup() *waitGroup {
    return &waitGroup{
        n:    0,
        wait: make(chan struct{}),
    }
}

func (wg *waitGroup) Add(n int) {
    wg.n = wg.n + n
}

func (wg *waitGroup) Done() {
    wg.wait <- struct{}{}
}

func (wg *waitGroup) Wait() {
    defer close(wg.wait)
    for i := 0; i < wg.n; i++ {
        <-wg.wait
    }
}

As you can see, we use a type wrapping all the thing we need. (It’s a basic idiom, so I don’t want to say why)

Then method Add is preparing for n we talk before. Adding these things in a dynamic way.

Next Done do the thing as we manually do in previous code.

And Wait is read the number of things equal final n.

The end lets say what happened in the previous code. You should close the channel always. So the code will be:

// ...
wait := make(chan interface{})
defer close(wait)
// ...

Maybe you will feel confusing about this part. The reason is channel won’t be collected by GC automatic(if it can, it will be another hell). So always closing it is important.

ps. In productive code, please still using the sync.WaitGroup, I do a test, sync.WaitGroup is 25% faster than the version you see here.

References:

The Go programming language

• Author: Alan A. A. Donovan & Brian W. Kernighan
• ISBN: 978-986-476-133-3

Concurrency in Go

• Author: Katherine Cox-Buday
• ISBN: 978-1-491-94119-5