github.com/searKing/golang/go@v1.2.74/sync/subject.go (about)

     1  // Copyright 2022 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package sync
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"sync"
    11  
    12  	"github.com/searKing/golang/go/errors"
    13  	"github.com/searKing/golang/go/pragma"
    14  	"github.com/searKing/golang/go/sync/atomic"
    15  )
    16  
    17  // Subject implements a condition variable like with channel, a rendezvous point
    18  // for goroutines waiting for or announcing the occurrence
    19  // of an event.
    20  //
    21  // The caller typically cannot assume that the condition is true when
    22  // Subscribe chan returns. Instead, the caller should Wait in a loop:
    23  //
    24  //    time.After(timeout, c.PublishBroadcast()) // for timeout or periodic event
    25  //    c.PublishBroadcast() // for async notify event directly
    26  //    eventC, cancel := c.Subscribe()
    27  //    for !condition() {
    28  //        select{
    29  //        case event, closed := <- eventC:
    30  //            ... make use of event ...
    31  //        }
    32  //    }
    33  //    ... make use of condition ...
    34  //
    35  type Subject struct {
    36  	noCopy pragma.DoNotCopy
    37  
    38  	mu          sync.Mutex
    39  	subscribers map[*subscriber]struct{}
    40  
    41  	inShutdown atomic.Bool // true when when server is in shutdown
    42  }
    43  
    44  type subscriber struct {
    45  	mu   sync.Mutex // guard close of channel msgC
    46  	msgC chan interface{}
    47  
    48  	once  sync.Once
    49  	doneC chan struct{} // closed when subscriber is in shutdown, like removed.
    50  }
    51  
    52  func (s *subscriber) Shutdown() {
    53  	if s == nil {
    54  		return
    55  	}
    56  	s.once.Do(func() {
    57  		close(s.doneC)
    58  		s.mu.Lock()
    59  		defer s.mu.Unlock()
    60  		close(s.msgC)
    61  	})
    62  }
    63  
    64  // publish wakes a listener waiting on c to consume the event.
    65  // event will be dropped if ctx is Done before event is received.
    66  func (s *subscriber) publish(ctx context.Context, event interface{}) error {
    67  	// guard of msgC's close
    68  	s.mu.Lock()
    69  	defer s.mu.Unlock()
    70  	select {
    71  	case <-ctx.Done():
    72  		// event dropped because of publisher
    73  		return ctx.Err()
    74  	case <-s.doneC:
    75  		// event dropped because of subscriber
    76  		return fmt.Errorf("event dropped because of subscriber unsubscribed")
    77  	default:
    78  	}
    79  
    80  	select {
    81  	case <-ctx.Done():
    82  		// event dropped because of publisher
    83  		return ctx.Err()
    84  	case <-s.doneC:
    85  		// event dropped because of subscriber
    86  		return fmt.Errorf("event dropped because of subscriber unsubscribed")
    87  	case s.msgC <- event:
    88  		// event consumed
    89  		return nil
    90  	}
    91  }
    92  
    93  // Subscribe returns a channel that's closed when awoken by PublishSignal or PublishBroadcast.
    94  // never be canceled. Successive calls to Subscribe return different values.
    95  // The close of the Subscribe channel may happen asynchronously,
    96  // after the cancel function returns.
    97  func (s *Subject) Subscribe() (<-chan interface{}, context.CancelFunc) {
    98  	listener := &subscriber{
    99  		msgC:  make(chan interface{}),
   100  		doneC: make(chan struct{}),
   101  	}
   102  	s.trackChannel(listener, true)
   103  	return listener.msgC, func() {
   104  		s.trackChannel(listener, false)
   105  	}
   106  }
   107  
   108  // PublishSignal wakes one listener waiting on c, if there is any.
   109  // PublishSignal blocks until event is received or dropped.
   110  func (s *Subject) PublishSignal(ctx context.Context, event interface{}) error {
   111  	s.mu.Lock()
   112  	defer s.mu.Unlock()
   113  	var wg sync.WaitGroup
   114  	var errs []error
   115  	for listener := range s.subscribers {
   116  		wg.Add(1)
   117  		go func(listener *subscriber) {
   118  			defer wg.Done()
   119  			err := listener.publish(ctx, event)
   120  			if err != nil {
   121  				errs = append(errs, err)
   122  			}
   123  		}(listener)
   124  		break
   125  	}
   126  	wg.Wait()
   127  	return errors.Multi(errs...)
   128  }
   129  
   130  // PublishBroadcast wakes all listeners waiting on c.
   131  // PublishBroadcast blocks until event is received or dropped.
   132  // event will be dropped if ctx is Done before event is received.
   133  func (s *Subject) PublishBroadcast(ctx context.Context, event interface{}) error {
   134  	var wg sync.WaitGroup
   135  	var errs []error
   136  	func() {
   137  		s.mu.Lock()
   138  		defer s.mu.Unlock()
   139  		for listener := range s.subscribers {
   140  			wg.Add(1)
   141  			go func(listener *subscriber) {
   142  				defer wg.Done()
   143  				err := listener.publish(ctx, event)
   144  				if err != nil {
   145  					errs = append(errs, err)
   146  				}
   147  			}(listener)
   148  		}
   149  	}()
   150  	wg.Wait()
   151  	return errors.Multi(errs...)
   152  }
   153  
   154  func (s *Subject) trackChannel(c *subscriber, add bool) {
   155  	func() {
   156  		s.mu.Lock()
   157  		defer s.mu.Unlock()
   158  		if s.subscribers == nil {
   159  			s.subscribers = make(map[*subscriber]struct{})
   160  		}
   161  		_, has := s.subscribers[c]
   162  		if has {
   163  			if add {
   164  				return
   165  			}
   166  			delete(s.subscribers, c)
   167  			return
   168  		}
   169  		if add {
   170  			s.subscribers[c] = struct{}{}
   171  		}
   172  	}()
   173  	if !add {
   174  		c.Shutdown()
   175  	}
   176  }