How to handle multiple go-routines closing the same channel?

Issue

I am having 2 go-routines reading from a single channel. After 4 seconds, I cancel the context and terminate the select loop. Before terminating the loop I call close on the channel, since there are 2 go-routines the close gets called twice and causes a panic because one of the go-routines would have already closed the channel. Currently I am using a recover to recover from the panic, is there a better way of doing this?

package main

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

func numberGen(ctx context.Context, numChan chan int) {
    num := 0
    doneCh := ctx.Done()
    defer func() {
        if r := recover(); r != nil {
            fmt.Println("recovered from ", r)
        }
    }()
    for {
        select {
        case <-doneCh:
            fmt.Println("done generating...")
            close(numChan)
            return
        default:
            num++
            numChan <- num
        }
    }
}

func main() {
    ctx, cancelFn := context.WithCancel(context.Background())
    numChan := make(chan int)
    var wg sync.WaitGroup
    wg.Add(2)
    go numberGen(ctx, numChan)
    go numberGen(ctx, numChan)

    go func(cfn context.CancelFunc) {
        time.Sleep(10 * time.Millisecond)
        cfn()
    }(cancelFn)

    for n := range numChan {
        fmt.Println("received value ", n)
    }

    time.Sleep(2 * time.Second)
}

Solution

Close the channel after the goroutines are done sending values.

var wg sync.WaitGroup
wg.Add(2)
go numberGen(ctx, numChan, &wg)
go numberGen(ctx, numChan, &wg)

go func() {
    wg.Wait()
    close(numChan)
}()

Update numberGen to call Done() on the wait group. Also, remove the call to close.

func numberGen(ctx context.Context, numChan chan int, wg *sync.WaitGroup) {
    defer wg.Done()
    ...

Answered By – Bayta Darell

Answer Checked By – Robin (GoLangFix Admin)

Leave a Reply

Your email address will not be published.