Stuck in infinite loop in for select in Golang

Question

The code given below is the sample code for my use case. I want to read data from ch1 and ch2 but got stuck into infinite loop.

package main
import "fmt"

func main() {
    ch1, ch2 := func() (<-chan int, <-chan int) {
        ch_1 := make(chan int)
        ch_2 := make(chan int)

        go worker_1(ch_1, ch_2)
        go worker_2(ch_1, ch_2)

        return ch_1, ch_2
    }()

    // trying to read this way but it is not working
    for {
        select {
        case a := <-ch1:
            fmt.Println("from ch1", a)
        case a := <-ch2:
            fmt.Println("from ch2", a)
        default:
            fmt.Println("done")
        }
    }
}

func worker_1(ch1, ch2 chan int) {
    for i := 0; i < 100; i++ {
        if i%2 == 0 {
            ch1 <- i
        } else {
            ch2 <- i
        }
    }
}

func worker_2(ch1, ch2 chan int) {
    for i := 101; i < 200; i++ {
        if i%2 == 0 {
            ch1 <- i
        } else {
            ch2 <- i
        }
    }
}

Answer 1

Here is one solution:

package main

import (
    "fmt"
    "sync"
)

func main() {
    // Create channels
    ch1, ch2 := make(chan int), make(chan int)

    // Create workers waitgroup with a counter of 2
    wgWorkers := sync.WaitGroup{}
    wgWorkers.Add(2)

    // Run workers
    go worker(&wgWorkers, ch1, ch2, 0, 100)   // Worker 1
    go worker(&wgWorkers, ch1, ch2, 101, 200) // Worker 2

    // Create readers waitgroup with a counter of 2
    wgReader := sync.WaitGroup{}
    wgReader.Add(2)

    // Run readers
    go reader(&wgReader, ch1, 1) // Reader 1
    go reader(&wgReader, ch2, 2) // Reader 2

    // Wait for workers to finish
    wgWorkers.Wait()

    // Close workers channels
    close(ch1) // Makes reader 1 exit after processing the last element in the channel
    close(ch2) // Makes reader 2 exit after processing the last element in the channel

    // Wait for both readers to finish processing before exiting the program
    wgReader.Wait()
}

// Worker function definition
func worker(wg *sync.WaitGroup, ch1, ch2 chan<- int, from, to int) {
    // Decrement worker waitgroup counter by one when function returns
    defer wg.Done()
    for i := from; i < to; i++ {
        if i%2 == 0 {
            ch1 <- i
        } else {
            ch2 <- i
        }
    }
}

// Reader function definition
func reader(wg *sync.WaitGroup, ch <-chan int, chNum int) {
    // Decrement reader waitgroup counter by one when function returns
    defer wg.Done()
    // Here we iterate on the channel fed by worker 1 or worker 2.
    // for-range on a channel exits when the channel is closed.
    for i := range ch {
        fmt.Printf("from ch%d: %d\n", chNum, i)
    }
}

Explainations are in the code comments.

Answer 2

Close the channels when the workers are done. Break out of the receive loop after both channels are closed.

package main

import (
    "fmt"
    "sync"
)

func main() {
    ch1, ch2 := func() (<-chan int, <-chan int) {
        ch_1 := make(chan int)
        ch_2 := make(chan int)

        var wg sync.WaitGroup
        wg.Add(2)

        go worker_1(&wg, ch_1, ch_2)
        go worker_2(&wg, ch_1, ch_2)

        // Close channels after goroutiens complete.
        go func() {
            wg.Wait()
            close(ch_1)
            close(ch_2)
        }()

        return ch_1, ch_2
    }()

    // While we still have open channels ...
    for ch1 != nil || ch2 != nil {
        select {
        case a, ok := <-ch1:
            if ok {
                fmt.Println("from ch1", a)
            } else {
                // note that channel is closed.
                ch1 = nil
            }
        case a, ok := <-ch2:
            if ok {
                fmt.Println("from ch2", a)
            } else {
                // note that channel is closed.
                ch2 = nil
            }
        }
    }
}

func worker_1(wg *sync.WaitGroup, ch1, ch2 chan int) {
    defer wg.Done()
    for i := 0; i < 100; i++ {
        if i%2 == 0 {
            ch1 <- i
        } else {
            ch2 <- i
        }
    }
}

func worker_2(wg *sync.WaitGroup, ch1, ch2 chan int) {
    defer wg.Done()
    for i := 101; i < 200; i++ {
        if i%2 == 0 {
            ch1 <- i
        } else {
            ch2 <- i
        }
    }
}