for c := range ch{}節内でchに値を送信する場合のclose(ch)の方法について

Question

前提・実現したいこと

goでbufferedChannelを用いて繰り返し処理を書いています。
for rangeを使ってchannelから値を読み込み、節内で呼び出している別の関数hoge()で同じchannelに値を送信するようなソースコードになっています。
hoge()が呼ばれるたびに関数内でchannelに送信する回数は減っていき、最後にはchannelへの送信自体がされなくなるような挙動のイメージです。
channelに値が送信されなくなり、chのバッファが空になったときにcloseをしたいのですが、適切なハンドリングができずに困っています。
恐れ入りますが、どのようにclose処理を書けばいいか（もしくは、こういった書き方はそもそも文法としてふさわしくないか）ご教示いただければと思います。
期待している出力は下記です。

3
0
1
2
0
1
0
0
done

発生している問題・エラーメッセージ

該当のソースコード

package main

import (
	"fmt"
	"time"
)

func main() {
	ch := make(chan int, 10)
	ch <- 3
	go func() {
		time.Sleep(2 * time.Second) // ここを適切な形に修正したい
		close(ch)
	}()
	for c := range ch {
		fmt.Println(c)
		go hoge(ch, c)
	}
	fmt.Println("done")
}

func hoge(ch chan int, n int) {
	for i := 0; i < n; i++ {
		ch <- i
	}
}

試したこと

sync.WaitGroupを使って、下記のように修正してみました。
うまくいく場合もあるのですが、次の繰り返しによるwg.Add(1)より先にhoge()内のwg.Done()が評価されてしまうと、wg.Wait()以降に入ってしまい、次のhogeが実行される前に終了する場合もあります。

package main

import (
	"fmt"
	"sync"
)

func main() {
	ch := make(chan int, 10)
	ch <- 3
	var wg sync.WaitGroup
	go func() {
		wg.Wait()
		close(ch)
	}()
	for c := range ch {
		wg.Add(1)
		fmt.Println(c)
		go hoge(&wg, ch, c)
	}
	fmt.Println("done")
}

func hoge(wg *sync.WaitGroup, ch chan int, n int) {
	defer wg.Done()
	for i := 0; i < n; i++ {
		ch <- i
	}
}

補足情報（FW/ツールのバージョンなど）

Comments

[nobonobo]

「挙動のイメージ」から具体的な挙動が想像するのが難しいので、
期待している出力（質問者の思ってるうまくいった時の出力）を追記して欲しい。

[knr]

ありがとうございます。追記しました

Answer 1

後者のコードだけに言及してみます。
wg.Waitがcloseを待ち構えた構造ですが、
wg.Addのあとにwg.Waitにはいる動作順を保証できていません。

以下のようにwg.Waitの手前にスリープを入れると確実に動作します。
これが良いコードかというと良くないです。
終了条件を整理してちゃんとロジカルにcloseする方が良いと思います。

https://play.golang.org/p/qt26a7_KaxY

package main

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

func main() {
	ch := make(chan int, 10)
	ch <- 3
	var wg sync.WaitGroup
	go func() {
		time.Sleep(time.Second)
		wg.Wait()
		close(ch)
	}()
	for c := range ch {
		wg.Add(1)
		fmt.Println(c)
		go hoge(&wg, ch, c)
	}
	fmt.Println("done")
}

func hoge(wg *sync.WaitGroup, ch chan int, n int) {
	defer wg.Done()
	for i := 0; i < n; i++ {
		ch <- i
	}
}

作ろうとされている挙動の目的が質問文からは読み取れません。
なので全体のより良い仕組みを提案はできないので細かい点だけアドバイスします。

• chanを再帰的に利用するのはパズルの発生を招きます。chanの書き込みを行う責任と読み出しを行う責任はしっかり分離しましょう。（現状のコードは読み出し結果から書き込み内容が変化するという実装）
• chanとsyncパッケージの同期オブジェクトを併用しだすとパズルになってしまいます。そこの整理ができない場合は無理してchanやgoroutineを使わない方が良い結果になることが多いです。
• ほとんどのケースでWaitGroupを関数引数に渡して使う必要はなくて以下のコードのようにgoroutine起動側に書くことをおすすめします。

go func() {
    defer wg.Done()
    hoge(ch, c)
}()

Comments

[knr]

chanの使い方に関するアドバイスもいただき、ありがとうございます！

[nobonobo]

goroutine、chanを使わない回答例を書いてみました。参考に。
https://play.golang.org/p/4SydHsYWplv