github.com/neatio-net/neatio@v1.7.3-0.20231114194659-f4d7a2226baa/utilities/event/feed.go (about)

     1  package event
     2  
     3  import (
     4  	"errors"
     5  	"reflect"
     6  	"sync"
     7  )
     8  
     9  var errBadChannel = errors.New("event: Subscribe argument does not have sendable channel type")
    10  
    11  type Feed struct {
    12  	once      sync.Once
    13  	sendLock  chan struct{}
    14  	removeSub chan interface{}
    15  	sendCases caseList
    16  
    17  	mu     sync.Mutex
    18  	inbox  caseList
    19  	etype  reflect.Type
    20  	closed bool
    21  }
    22  
    23  const firstSubSendCase = 1
    24  
    25  type feedTypeError struct {
    26  	got, want reflect.Type
    27  	op        string
    28  }
    29  
    30  func (e feedTypeError) Error() string {
    31  	return "event: wrong type in " + e.op + " got " + e.got.String() + ", want " + e.want.String()
    32  }
    33  
    34  func (f *Feed) init() {
    35  	f.removeSub = make(chan interface{})
    36  	f.sendLock = make(chan struct{}, 1)
    37  	f.sendLock <- struct{}{}
    38  	f.sendCases = caseList{{Chan: reflect.ValueOf(f.removeSub), Dir: reflect.SelectRecv}}
    39  }
    40  
    41  func (f *Feed) Subscribe(channel interface{}) Subscription {
    42  	f.once.Do(f.init)
    43  
    44  	chanval := reflect.ValueOf(channel)
    45  	chantyp := chanval.Type()
    46  	if chantyp.Kind() != reflect.Chan || chantyp.ChanDir()&reflect.SendDir == 0 {
    47  		panic(errBadChannel)
    48  	}
    49  	sub := &feedSub{feed: f, channel: chanval, err: make(chan error, 1)}
    50  
    51  	f.mu.Lock()
    52  	defer f.mu.Unlock()
    53  	if !f.typecheck(chantyp.Elem()) {
    54  		panic(feedTypeError{op: "Subscribe", got: chantyp, want: reflect.ChanOf(reflect.SendDir, f.etype)})
    55  	}
    56  
    57  	cas := reflect.SelectCase{Dir: reflect.SelectSend, Chan: chanval}
    58  	f.inbox = append(f.inbox, cas)
    59  	return sub
    60  }
    61  
    62  func (f *Feed) typecheck(typ reflect.Type) bool {
    63  	if f.etype == nil {
    64  		f.etype = typ
    65  		return true
    66  	}
    67  	return f.etype == typ
    68  }
    69  
    70  func (f *Feed) remove(sub *feedSub) {
    71  
    72  	ch := sub.channel.Interface()
    73  	f.mu.Lock()
    74  	index := f.inbox.find(ch)
    75  	if index != -1 {
    76  		f.inbox = f.inbox.delete(index)
    77  		f.mu.Unlock()
    78  		return
    79  	}
    80  	f.mu.Unlock()
    81  
    82  	select {
    83  	case f.removeSub <- ch:
    84  
    85  	case <-f.sendLock:
    86  
    87  		f.sendCases = f.sendCases.delete(f.sendCases.find(ch))
    88  		f.sendLock <- struct{}{}
    89  	}
    90  }
    91  
    92  func (f *Feed) Send(value interface{}) (nsent int) {
    93  	rvalue := reflect.ValueOf(value)
    94  
    95  	f.once.Do(f.init)
    96  	<-f.sendLock
    97  
    98  	f.mu.Lock()
    99  	f.sendCases = append(f.sendCases, f.inbox...)
   100  	f.inbox = nil
   101  
   102  	if !f.typecheck(rvalue.Type()) {
   103  		f.sendLock <- struct{}{}
   104  		panic(feedTypeError{op: "Send", got: rvalue.Type(), want: f.etype})
   105  	}
   106  	f.mu.Unlock()
   107  
   108  	for i := firstSubSendCase; i < len(f.sendCases); i++ {
   109  		f.sendCases[i].Send = rvalue
   110  	}
   111  
   112  	cases := f.sendCases
   113  	for {
   114  
   115  		for i := firstSubSendCase; i < len(cases); i++ {
   116  			if cases[i].Chan.TrySend(rvalue) {
   117  				nsent++
   118  				cases = cases.deactivate(i)
   119  				i--
   120  			}
   121  		}
   122  		if len(cases) == firstSubSendCase {
   123  			break
   124  		}
   125  
   126  		chosen, recv, _ := reflect.Select(cases)
   127  		if chosen == 0 {
   128  			index := f.sendCases.find(recv.Interface())
   129  			f.sendCases = f.sendCases.delete(index)
   130  			if index >= 0 && index < len(cases) {
   131  
   132  				cases = f.sendCases[:len(cases)-1]
   133  			}
   134  		} else {
   135  			cases = cases.deactivate(chosen)
   136  			nsent++
   137  		}
   138  	}
   139  
   140  	for i := firstSubSendCase; i < len(f.sendCases); i++ {
   141  		f.sendCases[i].Send = reflect.Value{}
   142  	}
   143  	f.sendLock <- struct{}{}
   144  	return nsent
   145  }
   146  
   147  type feedSub struct {
   148  	feed    *Feed
   149  	channel reflect.Value
   150  	errOnce sync.Once
   151  	err     chan error
   152  }
   153  
   154  func (sub *feedSub) Unsubscribe() {
   155  	sub.errOnce.Do(func() {
   156  		sub.feed.remove(sub)
   157  		close(sub.err)
   158  	})
   159  }
   160  
   161  func (sub *feedSub) Err() <-chan error {
   162  	return sub.err
   163  }
   164  
   165  type caseList []reflect.SelectCase
   166  
   167  func (cs caseList) find(channel interface{}) int {
   168  	for i, cas := range cs {
   169  		if cas.Chan.Interface() == channel {
   170  			return i
   171  		}
   172  	}
   173  	return -1
   174  }
   175  
   176  func (cs caseList) delete(index int) caseList {
   177  	return append(cs[:index], cs[index+1:]...)
   178  }
   179  
   180  func (cs caseList) deactivate(index int) caseList {
   181  	last := len(cs) - 1
   182  	cs[index], cs[last] = cs[last], cs[index]
   183  	return cs[:last]
   184  }