github.com/Mericusta/go-stp@v0.6.8/channel.go (about)

     1  package stp
     2  
     3  import "sync/atomic"
     4  
     5  type SharedChannel[T any] struct {
     6  	noCopy
     7  	c  chan T
     8  	uc *int64
     9  }
    10  
    11  func NewSharedChannel[T any]() *SharedChannel[T] {
    12  	sc := &SharedChannel[T]{
    13  		c:  make(chan T),
    14  		uc: new(int64),
    15  	}
    16  	*sc.uc = 1
    17  	return sc
    18  }
    19  
    20  func NewSharedBufferChannel[T any](b int64) *SharedChannel[T] {
    21  	sc := &SharedChannel[T]{
    22  		c:  make(chan T, b),
    23  		uc: new(int64),
    24  	}
    25  	*sc.uc = 1
    26  	return sc
    27  }
    28  
    29  func (sc *SharedChannel[T]) Share() *SharedChannel[T] {
    30  	atomic.AddInt64(sc.uc, 1)
    31  	return &SharedChannel[T]{c: sc.c, uc: sc.uc}
    32  }
    33  
    34  func (s *SharedChannel[T]) Get() chan T {
    35  	return s.c
    36  }
    37  
    38  func (s *SharedChannel[T]) UseCount() int64 {
    39  	return atomic.LoadInt64(s.uc)
    40  }
    41  
    42  func (sc *SharedChannel[T]) Close() {
    43  	if atomic.AddInt64(sc.uc, -1) == 0 {
    44  		close(sc.c)
    45  	}
    46  }