How to receive data from multiple channels using for/select syntax?

July 27, 2022

I have case where I need to receive data on multiple channels and somehow quit infinite for loop.

Here is simple example:

package main

import (
    "fmt"
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup
    sumChannel := make(chan int)
    productChannel := make(chan int)

    for i := 1; i <= 5; i++ {
        wg.Add(1)
        go compute(i, sumChannel, productChannel, &wg)
    }

    go func() {
        wg.Wait()
        close(sumChannel)
        close(productChannel)
    }()

    for {
        select {
        case <-sumChannel:
            fmt.Println("sum")
        case <-productChannel:
            fmt.Println("prod")
        }
    }
}

func compute(i int, sumChannel chan int, productChannel chan int, wg *sync.WaitGroup) {
    time.Sleep(2 * time.Second)
    sumChannel <- i + i
    productChannel <- i * i
    wg.Done()
}

The problem is I get infinite for loop. I know that I should create new channel (for instance,quit) for quitting from loop. But, I really don’t understand where to put signal to quit, because I don’t know exact place where I can be sure that all goroutines finished.

>Solution :

If you run the for loop in a go routine and then wait for the compute go routines to finish, you can cancel the loop from there:

package main

import (
    "fmt"
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup
    sumChannel := make(chan int)
    productChannel := make(chan int)
    closeChannel := make(chan bool)

    for i := 1; i <= 5; i++ {
        wg.Add(1)
        go compute(i, sumChannel, productChannel, &wg)
    }

    go func() {
        for {
            select {
            case <-sumChannel:
                fmt.Println("sum")
            case <-productChannel:
                fmt.Println("prod")
            case <-closeChannel:
                fmt.Println("break")
                break
            }
        }
    }()

    wg.Wait()
    closeChannel <- true

    fmt.Println("done")
}

func compute(i int, sumChannel chan int, productChannel chan int, wg *sync.WaitGroup) {
    time.Sleep(2 * time.Second)
    sumChannel <- i + i
    productChannel <- i * i
    wg.Done()
}

Run it