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  }