github.com/eagleql/xray-core@v1.4.4/common/signal/pubsub/pubsub.go (about)

     1  package pubsub
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/eagleql/xray-core/common"
     9  	"github.com/eagleql/xray-core/common/signal/done"
    10  	"github.com/eagleql/xray-core/common/task"
    11  )
    12  
    13  type Subscriber struct {
    14  	buffer chan interface{}
    15  	done   *done.Instance
    16  }
    17  
    18  func (s *Subscriber) push(msg interface{}) {
    19  	select {
    20  	case s.buffer <- msg:
    21  	default:
    22  	}
    23  }
    24  
    25  func (s *Subscriber) Wait() <-chan interface{} {
    26  	return s.buffer
    27  }
    28  
    29  func (s *Subscriber) Close() error {
    30  	return s.done.Close()
    31  }
    32  
    33  func (s *Subscriber) IsClosed() bool {
    34  	return s.done.Done()
    35  }
    36  
    37  type Service struct {
    38  	sync.RWMutex
    39  	subs  map[string][]*Subscriber
    40  	ctask *task.Periodic
    41  }
    42  
    43  func NewService() *Service {
    44  	s := &Service{
    45  		subs: make(map[string][]*Subscriber),
    46  	}
    47  	s.ctask = &task.Periodic{
    48  		Execute:  s.Cleanup,
    49  		Interval: time.Second * 30,
    50  	}
    51  	return s
    52  }
    53  
    54  // Cleanup cleans up internal caches of subscribers.
    55  // Visible for testing only.
    56  func (s *Service) Cleanup() error {
    57  	s.Lock()
    58  	defer s.Unlock()
    59  
    60  	if len(s.subs) == 0 {
    61  		return errors.New("nothing to do")
    62  	}
    63  
    64  	for name, subs := range s.subs {
    65  		newSub := make([]*Subscriber, 0, len(s.subs))
    66  		for _, sub := range subs {
    67  			if !sub.IsClosed() {
    68  				newSub = append(newSub, sub)
    69  			}
    70  		}
    71  		if len(newSub) == 0 {
    72  			delete(s.subs, name)
    73  		} else {
    74  			s.subs[name] = newSub
    75  		}
    76  	}
    77  
    78  	if len(s.subs) == 0 {
    79  		s.subs = make(map[string][]*Subscriber)
    80  	}
    81  	return nil
    82  }
    83  
    84  func (s *Service) Subscribe(name string) *Subscriber {
    85  	sub := &Subscriber{
    86  		buffer: make(chan interface{}, 16),
    87  		done:   done.New(),
    88  	}
    89  	s.Lock()
    90  	subs := append(s.subs[name], sub)
    91  	s.subs[name] = subs
    92  	s.Unlock()
    93  	common.Must(s.ctask.Start())
    94  	return sub
    95  }
    96  
    97  func (s *Service) Publish(name string, message interface{}) {
    98  	s.RLock()
    99  	defer s.RUnlock()
   100  
   101  	for _, sub := range s.subs[name] {
   102  		if !sub.IsClosed() {
   103  			sub.push(message)
   104  		}
   105  	}
   106  }