main.go

package main

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

// generator function creates a channel and generates numbers to send to the channel
// it listen on a 'done' to cancel generation if needed1
func generator(done <-chan struct{}, nums ...int) <-chan int {
	out := make(chan int, len(nums))
	go func() {
		for _, n := range nums {
			fmt.Printf("generator: sending %d\n", n)
			select {
			case out <- n:		
			case <-done:		
				fmt.Println("generator: canceled")
				return
			}
		}
		close(out)
	}()
	return out
}

// square function receives integers,square them, and sends them to an output channel
// it also listens for a cancellation signal via 'done' channel
func square(done <-chan struct{}, in <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range in {
			select {
			case out <- n * n:
			case <-done:
				return
			}
		}
	}()
	return out
}

// merge function merges multiple channels into one output channel
// it listens on each input channel and sends data to 'out', with cancellation support
func merge(done <-chan struct{}, cs ...<-chan int) <-chan int {
	out := make(chan int)
	var wg sync.WaitGroup

	// 'output' reads data from each input channel and sends it to 'out' channel
	output := func(c <-chan int) {
		defer wg.Done()
		for n := range c {
			select {
			case out <- n:
			case <-done:
				return
			}
		}
	}

	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

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

	return out
}

// Flow work:
// 
// [ Generator ] -----> 	[ In Channel ] 
//       |                           |
//       v                           v
//   Send 2, 3               Receive 2, 3 
//       |                           |
//       v                           v
// [ Square (c1) ]         [ Square (c2) ]
//       |                           |
//       v                           v
//   Square 2, 3              Square 2, 3
//       |                           |
//       v                           v
// [ Merge ] <--------------------->
//       |
//       v
//   Final output
func main() {
	// 'done' channel for cancellation
	done := make(chan struct{})
	 
	// start generator with input numbers 2 and 3
	in := generator(done, 2, 3)

	c1 := square(done, in)
	c2 := square(done, in)
	out := merge(done, c1, c2)

	for v := range out {
		fmt.Println("Result:", v)
	}
	close(done)

	time.Sleep(time.Second)
	g := runtime.NumGoroutine()
	fmt.Printf("=> number of goroutine active = %d\n", g)
}