github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/zutil/chan.go (about)

     1  //go:build go1.18
     2  // +build go1.18
     3  
     4  package zutil
     5  
     6  type (
     7  	chanType int
     8  	conf     struct {
     9  		len *Uint32
    10  		typ chanType
    11  		cap int64
    12  	}
    13  	Opt         func(*conf)
    14  	Chan[T any] struct {
    15  		in, out chan T
    16  		close   chan struct{}
    17  		conf    conf
    18  		q       []T
    19  	}
    20  
    21  	Options struct {
    22  		Cap int
    23  	}
    24  )
    25  
    26  const (
    27  	unbuffered chanType = iota
    28  	buffered
    29  	unbounded
    30  )
    31  
    32  func NewChan[T any](cap ...int) *Chan[T] {
    33  	o := conf{
    34  		typ: unbounded,
    35  		cap: -1,
    36  		len: NewUint32(0),
    37  	}
    38  
    39  	if len(cap) > 0 {
    40  		if cap[0] == 0 {
    41  			o.cap = int64(0)
    42  			o.typ = unbuffered
    43  		} else if cap[0] > 0 {
    44  			o.cap = int64(cap[0])
    45  			o.typ = buffered
    46  		} else {
    47  			o.cap = int64(-1)
    48  			o.typ = unbounded
    49  		}
    50  	}
    51  
    52  	ch := &Chan[T]{conf: o, close: make(chan struct{})}
    53  	switch ch.conf.typ {
    54  	case unbuffered:
    55  		ch.in = make(chan T)
    56  		ch.out = ch.in
    57  	case buffered:
    58  		ch.in = make(chan T, ch.conf.cap)
    59  		ch.out = ch.in
    60  	case unbounded:
    61  		ch.in = make(chan T, 16)
    62  		ch.out = make(chan T, 16)
    63  		go ch.process()
    64  	}
    65  	return ch
    66  }
    67  
    68  func (ch *Chan[T]) In() chan<- T { return ch.in }
    69  
    70  func (ch *Chan[T]) Out() <-chan T { return ch.out }
    71  
    72  func (ch *Chan[T]) Close() {
    73  	switch ch.conf.typ {
    74  	case buffered, unbuffered:
    75  		close(ch.in)
    76  		close(ch.close)
    77  	default:
    78  		ch.close <- struct{}{}
    79  	}
    80  }
    81  
    82  func (ch *Chan[T]) Len() int {
    83  	switch ch.conf.typ {
    84  	case buffered, unbuffered:
    85  		return len(ch.in)
    86  	default:
    87  		return int(ch.conf.len.Load()) + len(ch.in) + len(ch.out)
    88  	}
    89  }
    90  
    91  func (ch *Chan[T]) process() {
    92  	var nilT T
    93  
    94  	ch.q = make([]T, 0, 1<<10)
    95  	for {
    96  		select {
    97  		case e, ok := <-ch.in:
    98  			if !ok {
    99  				return
   100  			}
   101  			ch.conf.len.Add(1)
   102  			ch.q = append(ch.q, e)
   103  		case <-ch.close:
   104  			ch.closeUnbounded()
   105  			return
   106  		}
   107  
   108  		for len(ch.q) > 0 {
   109  			select {
   110  			case ch.out <- ch.q[0]:
   111  				ch.conf.len.Sub(1)
   112  				ch.q[0] = nilT
   113  				ch.q = ch.q[1:]
   114  			case e, ok := <-ch.in:
   115  				if !ok {
   116  					return
   117  				}
   118  				ch.conf.len.Add(1)
   119  				ch.q = append(ch.q, e)
   120  			case <-ch.close:
   121  				ch.closeUnbounded()
   122  				return
   123  			}
   124  		}
   125  		if cap(ch.q) < 1<<5 {
   126  			ch.q = make([]T, 0, 1<<10)
   127  		}
   128  	}
   129  }
   130  
   131  func (ch *Chan[T]) closeUnbounded() {
   132  	var nilT T
   133  
   134  	close(ch.in)
   135  
   136  	for e := range ch.in {
   137  		ch.q = append(ch.q, e)
   138  	}
   139  
   140  	for len(ch.q) > 0 {
   141  		ch.out <- ch.q[0]
   142  		ch.q[0] = nilT
   143  		ch.q = ch.q[1:]
   144  	}
   145  
   146  	close(ch.out)
   147  	close(ch.close)
   148  }