fox/safeChan/safeChan.go

package safeChan

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

/*
	原生chan主张谁申请谁负责关闭。但实际项目中权责往往没有这么清晰。
	有时候会出现多个协程对同一个chan写入的情况，此时chan被某一个协程关闭，另一个正在写入会导致panic。
	如果对原生chan关闭多次也会panic。原生chan更适合一写一读或多读的场景，生产者负责chan生命周期，消费者只负责读。
	SafeChan可以多写多读多关闭而不会panic。
*/

type ByteChan = SafeChan[[]byte]

type SafeChan[T any] struct {
	ch     chan T
	ctx    context.Context
	cancel context.CancelFunc
	once   sync.Once
}

// size为0会创建无缓冲队列，读chan时，没有数据会阻塞到有数据可读
func NewSafeChan[T any](size int) *SafeChan[T] {
	ch := &SafeChan[T]{}
	ch.ctx, ch.cancel = context.WithCancel(context.Background())
	if size < 1 {
		ch.ch = make(chan T)
	} else {
		ch.ch = make(chan T, size)
	}
	// ch.ch = make(chan T, size)
	return ch
}

func (s *SafeChan[T]) Close() {
	s.once.Do(func() {
		s.cancel()
		close(s.ch)
	})
}

// 管道中剩余数量
func (s *SafeChan[T]) Size() int {
	return len(s.ch)
}

func (s *SafeChan[T]) Reader() <-chan T {
	//log.Debug("reader channel")
	return s.ch
}

func (s *SafeChan[T]) Write(d T) error {
	select {
	case <-s.ctx.Done():
		return fmt.Errorf("chan was closed")
	default:
		s.ch <- d
		return nil
	}
}