github.com/benz9527/xboot@v0.0.0-20240504061247-c23f15593274/lib/infra/safe_chan.go (about) 1 package infra 2 3 import ( 4 "fmt" 5 "io" 6 "sync" 7 "sync/atomic" 8 ) 9 10 type ReadOnlyChannel[T comparable] interface { 11 Wait() <-chan T 12 } 13 14 type SendOnlyChannel[T comparable] interface { 15 Send(v T, nonBlocking ...bool) error 16 IsClosed() bool 17 } 18 19 type ClosableChannel[T comparable] interface { 20 io.Closer 21 ReadOnlyChannel[T] 22 SendOnlyChannel[T] 23 } 24 25 // safeClosableChannel is a generic channel wrapper. 26 // Why we need this wrapper? For the following reasons: 27 // 1. We need to make sure the channel is closed only once. 28 // 2. We need to make sure that we would not close the channel when there is still sending data. 29 // Let the related goroutines exit, then the channel auto-collected by GC. 30 type safeClosableChannel[T comparable] struct { 31 queueC chan T // Receive data to temporary queue. 32 isClosed atomic.Bool 33 once *sync.Once 34 } 35 36 func (c *safeClosableChannel[T]) IsClosed() bool { 37 return c.isClosed.Load() 38 } 39 40 // Close According to the Go memory model, a send operation on a channel happens before 41 // the corresponding to receive from that channel completes 42 // https://go.dev/doc/articles/race_detector 43 func (c *safeClosableChannel[T]) Close() error { 44 c.once.Do(func() { 45 // Note: forbid to call close(queueC) directly, 46 // because it will cause panic of "send on closed channel" 47 c.isClosed.Store(true) 48 }) 49 return nil 50 } 51 52 func (c *safeClosableChannel[T]) Wait() <-chan T { 53 return c.queueC 54 } 55 56 func (c *safeClosableChannel[T]) Send(v T, nonBlocking ...bool) error { 57 if c.isClosed.Load() { 58 return fmt.Errorf("channel has been closed") 59 } 60 61 if len(nonBlocking) <= 0 { 62 nonBlocking = []bool{false} 63 } 64 if !nonBlocking[0] { 65 c.queueC <- v 66 } else { 67 // non blocking send 68 select { 69 case c.queueC <- v: 70 default: 71 72 } 73 } 74 return nil 75 } 76 77 var ( 78 _ ReadOnlyChannel[struct{}] = &safeClosableChannel[struct{}]{} // type check assertion 79 _ SendOnlyChannel[struct{}] = &safeClosableChannel[struct{}]{} // type check assertion 80 ) 81 82 func NewSafeClosableChannel[T comparable](chSize ...int) ClosableChannel[T] { 83 isNoCacheCh := true 84 size := 1 85 if len(chSize) > 0 { 86 if chSize[0] > 0 { 87 size = chSize[0] 88 isNoCacheCh = false 89 } 90 } 91 if isNoCacheCh { 92 return &safeClosableChannel[T]{ 93 queueC: make(chan T), 94 once: &sync.Once{}, 95 } 96 } 97 return &safeClosableChannel[T]{ 98 queueC: make(chan T, size), 99 once: &sync.Once{}, 100 } 101 }