github.com/wfusion/gofusion@v1.1.14/common/infra/watermill/pubsub/gochannel/fanout.go (about)

     1  package gochannel
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"sync"
     8  
     9  	"go.uber.org/multierr"
    10  
    11  	"github.com/wfusion/gofusion/common/infra/watermill"
    12  	"github.com/wfusion/gofusion/common/infra/watermill/message"
    13  )
    14  
    15  // FanOut is a component that receives messages from the subscriber and passes them
    16  // to all publishers. In effect, messages are "multiplied".
    17  //
    18  // A typical use case for using FanOut is having one external subscription and multiple workers
    19  // inside the process.
    20  //
    21  // You need to call AddSubscription method for all topics that you want to listen to.
    22  // This needs to be done *before* starting the FanOut.
    23  //
    24  // FanOut exposes the standard Subscriber interface.
    25  type FanOut struct {
    26  	internalPubSub *GoChannel
    27  	internalRouter *message.Router
    28  
    29  	subscriber message.Subscriber
    30  
    31  	logger watermill.LoggerAdapter
    32  
    33  	subscribedTopics map[string]struct{}
    34  	subscribedLock   sync.Mutex
    35  }
    36  
    37  // NewFanOut creates a new FanOut.
    38  func NewFanOut(
    39  	subscriber message.Subscriber,
    40  	logger watermill.LoggerAdapter,
    41  ) (*FanOut, error) {
    42  	if subscriber == nil {
    43  		return nil, errors.New("missing subscriber")
    44  	}
    45  	if logger == nil {
    46  		logger = watermill.NopLogger{}
    47  	}
    48  
    49  	router, err := message.NewRouter(message.RouterConfig{}, logger)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	return &FanOut{
    55  		internalPubSub: NewGoChannel(Config{}, logger),
    56  		internalRouter: router,
    57  
    58  		subscriber: subscriber,
    59  
    60  		logger: logger,
    61  
    62  		subscribedTopics: map[string]struct{}{},
    63  	}, nil
    64  }
    65  
    66  // AddSubscription add an internal subscription for the given topic.
    67  // You need to call this method with all topics that you want to listen to, before the FanOut is started.
    68  // AddSubscription is idempotent.
    69  func (f *FanOut) AddSubscription(topic string) {
    70  	f.subscribedLock.Lock()
    71  	defer f.subscribedLock.Unlock()
    72  
    73  	_, ok := f.subscribedTopics[topic]
    74  	if ok {
    75  		// Subscription already exists
    76  		return
    77  	}
    78  
    79  	f.logger.Trace("Adding fan-out subscription for topic", watermill.LogFields{
    80  		"topic": topic,
    81  	})
    82  
    83  	f.internalRouter.AddHandler(
    84  		fmt.Sprintf("fanout-%s", topic),
    85  		topic,
    86  		f.subscriber,
    87  		topic,
    88  		f.internalPubSub,
    89  		message.PassthroughHandler,
    90  	)
    91  
    92  	f.subscribedTopics[topic] = struct{}{}
    93  }
    94  
    95  // Run runs the FanOut.
    96  func (f *FanOut) Run(ctx context.Context) error {
    97  	return f.internalRouter.Run(ctx)
    98  }
    99  
   100  // Running is closed when FanOut is running.
   101  func (f *FanOut) Running() chan struct{} {
   102  	return f.internalRouter.Running()
   103  }
   104  
   105  func (f *FanOut) IsClosed() bool {
   106  	return f.internalRouter.IsClosed()
   107  }
   108  
   109  // Subscribe starts subscription to the FanOut's internal Pub/Sub.
   110  func (f *FanOut) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
   111  	return f.internalPubSub.Subscribe(ctx, topic)
   112  }
   113  
   114  // Close closes the FanOut's internal Pub/Sub.
   115  func (f *FanOut) Close() error {
   116  	var err error
   117  
   118  	if routerCloseErr := f.internalRouter.Close(); routerCloseErr != nil {
   119  		err = multierr.Append(err, routerCloseErr)
   120  	}
   121  	if internalPubSubCloseErr := f.internalPubSub.Close(); internalPubSubCloseErr != nil {
   122  		err = multierr.Append(err, internalPubSubCloseErr)
   123  	}
   124  
   125  	return err
   126  }