64 lines
1.4 KiB
Go
64 lines
1.4 KiB
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
|
||
}
|
||
}
|