I’m learning about golang context package. I have a main go routine with the following code.

func main() {
    defer fmt.Println("main closing...")
    wg := sync.WaitGroup{}
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    applesStream := Generator("apple", ctx)
    orangesStream := Generator("orange", ctx)
    peachesStream := Generator("peach", ctx)

    wg.Add(1)
    go func1(ctx, &wg, applesStream)
    wg.Add(1)
    go genericFunc(ctx, &wg, orangesStream)
    wg.Add(1)
    go genericFunc(ctx, &wg, peachesStream)

    wg.Wait()

}

The Generator function will print the string value passed to it infinitely and will close after the ctx.Done() channel receives any value.
where the func1() and genericFunc() are as follows

func func1[T any](ctx context.Context, parentWg *sync.WaitGroup, stream <-chan T) {
    defer fmt.Println("func1 closing...")
    defer parentWg.Done()
    wg := sync.WaitGroup{}
    doWork := func(ctx context.Context) {
        defer wg.Done()
        for {
            select {
            case <-ctx.Done():
                return
            case val, ok := <-stream:
                if !ok {
                    fmt.Println("channel closed!")
                    return
                }
                fmt.Println(val)
            }
        }
    }

    newCtx, cancel := context.WithTimeout(ctx, time.Second*2)
    defer cancel()

    for i := 0; i < 3; i++ {
        wg.Add(1)
        go doWork(newCtx)
    }

    wg.Wait()

}

func genericFunc[T any](ctx context.Context, wg *sync.WaitGroup, stream <-chan T) {
    defer fmt.Println("genericFunc closing...")
    defer wg.Done()
    go func() {
        for {
            select {
            case <-ctx.Done():
                return
            case val, ok := <-stream:
                if !ok {
                    fmt.Println("channel closed!")
                    return
                }
                fmt.Println(val)
            }
        }
    }()
}

I expect this code to print apple, orange and peach for the first 3 seconds and after that it will print orange and peach infinitely.
But the main goroutine is shutting down after 3 seconds, which shouldn’t happen.

When I’m incrementing the waigGroup counter by only 1 in the main go routine then it’s behaving as expected. So just changing any of the wg.Add(1) to wg.Add(2) does the trick. so that means in main go routine the waigGroup counter somehow reaches zero after func1 finishes or somehow waitGroup doesn’t wait after func1 finishes, but how this is happening I don’t understand.