generics channels merge

generics_channels_merge.go

func merge[T any, C chan T](cs ...C) C {
	out := make(C)
	var wg sync.WaitGroup
	wg.Add(len(cs))
	for i := range cs {
		go func(c C) {
			defer wg.Done()
			for value := range c {
				out <- value
			}
		}(cs[i])
	}
	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

usage example:

package main

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log"
	"sync"
	"time"
)

type Msg struct {
	WriterID int
	Payload  string
}

func (m Msg) String() string {
	return fmt.Sprintf("id=%d payload=%s", m.WriterID, m.Payload)
}

type workerFunc[T any] func(int) chan T

func main() {
	var workers = make([]workerFunc[Msg], 0)
	for i := 0; i < 3; i++ {
		workers = append(workers, writer)
	}
	ch := PoolWorkers[Msg](workers...)
	for val := range ch {
		log.Println(val)
	}
}

func PoolWorkers[T any](wfs ...workerFunc[T]) <-chan T {
	var chans = make([]chan T, len(wfs))
	for i, wf := range wfs {
		chans[i] = wf(i + 1)
	}
	return merge[T](chans...)
}

func merge[T any, C chan T](cs ...C) C {
	out := make(C)
	var wg sync.WaitGroup
	wg.Add(len(cs))
	for i := range cs {
		go func(c C, wg *sync.WaitGroup) {
			defer wg.Done()
			for value := range c {
				out <- value
			}
		}(cs[i], &wg)
	}
	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

func writer(id int) chan Msg {
	var ch = make(chan Msg)
	go func() {
		var buf = make([]byte, 16)
		var ticker = time.NewTicker(time.Millisecond * 100)
		defer ticker.Stop()
		defer close(ch)
		for i := 0; i < 10; i++ {
			<-ticker.C
			rand.Read(buf)
			ch <- Msg{
				WriterID: id,
				Payload:  hex.EncodeToString(buf),
			}
		}
	}()
	return ch
}