Defer close on multiple Goroutines?

Question

I have the following code:

package main

import (
    "fmt"
    "time"
)

func main() {
    t := time.Now()
    stuff := fanIn(
        generator(4, 5, 6, 7),
        generator(1, 2, 6, 3, 7),
        generator(12, 15, 33, 40, 10),
        generator(18, 13, 20, 40, 15),
        generator(100, 200, 64000, 3121, 1237),
    )

    for v := range stuff {
        fmt.Println(v)
    }

    fmt.Println(t.Sub(time.Now()))
}

func generator(nums ...int) <-chan int {
    out := make(chan int, 10)
    go func() {
        defer close(out)
        for _, v := range nums {
            out <- v
        }
    }()
    return out
}

func fanIn(in ...<-chan int) <-chan int {
    out := make(chan int, 10)

    for _, v := range in {
        go func(ch <-chan int) {
            for val := range ch {
                go func(c int) { out <- c }(val)
            }
        }(v)
    }

    return out
}

It results in a deadlock on line 18:

for v := range stuff {...}

The issue (I think) is that I'm not deferring the close on the fanIn function that returns a read-only channel. I don't know when to defer it since it's got to wait for the end of multiple goroutines to complete.

What's the idiomatic way to solve this deadlock? Is this code even idiomatic?

Thanks!

GoPlay

Answer 1

You are correct about the cause of error being un-closed fanIn's channel. You can use a sync.WaitGroup to resolve the issue:

func fanIn(in ...<-chan int) <-chan int {
    // use a WaitGroup here
    var wg sync.WaitGroup
    out := make(chan int, 10)

    for _, v := range in {
        wg.Add(1)
        go func(ch <-chan int) {
            defer wg.Done()
            for val := range ch {
                out <- val
            }
        }(v)
    }

    // wait for wait groups to finish in another goroutine
    go func() {
        wg.Wait()
        close(out)
    }()
    return out
}

Working code.