Files
upstream-watermill/message/router.go
2026-09-15 21:27:46 +05:00

889 lines
24 KiB
Go

package message
import (
"context"
"fmt"
"runtime/debug"
"sync"
"time"
"github.com/pkg/errors"
"github.com/ThreeDotsLabs/watermill"
"github.com/ThreeDotsLabs/watermill/internal"
sync_internal "github.com/ThreeDotsLabs/watermill/pubsub/sync"
)
var (
// ErrOutputInNoPublisherHandler happens when a handler func returned some messages in a no-publisher handler.
// todo: maybe change the handler func signature in no-publisher handler so that there's no possibility for this
ErrOutputInNoPublisherHandler = errors.New("returned output messages in a handler without publisher")
)
// HandlerFunc is function called when message is received.
//
// msg.Ack() is called automatically when HandlerFunc doesn't return error.
// When HandlerFunc returns error, msg.Nack() is called.
// When msg.Ack() was called in handler and HandlerFunc returns error,
// msg.Nack() will be not sent because Ack was already sent.
//
// HandlerFunc's are executed parallel when multiple messages was received
// (because msg.Ack() was sent in HandlerFunc or Subscriber supports multiple consumers).
type HandlerFunc func(msg *Message) ([]*Message, error)
// NoPublishHandlerFunc is HandlerFunc alternative, which doesn't produce any messages.
type NoPublishHandlerFunc func(msg *Message) error
// PassthroughHandler is a handler that passes the message unchanged from the subscriber to the publisher.
var PassthroughHandler HandlerFunc = func(msg *Message) ([]*Message, error) {
return []*Message{msg}, nil
}
// HandlerMiddleware allows us to write something like decorators to HandlerFunc.
// It can execute something before handler (for example: modify consumed message)
// or after (modify produced messages, ack/nack on consumed message, handle errors, logging, etc.).
//
// It can be attached to the router by using `AddMiddleware` method.
//
// Example:
//
// func ExampleMiddleware(h message.HandlerFunc) message.HandlerFunc {
// return func(message *message.Message) ([]*message.Message, error) {
// fmt.Println("executed before handler")
// producedMessages, err := h(message)
// fmt.Println("executed after handler")
//
// return producedMessages, err
// }
// }
type HandlerMiddleware func(h HandlerFunc) HandlerFunc
// RouterPlugin is function which is executed on Router start.
type RouterPlugin func(*Router) error
// PublisherDecorator wraps the underlying Publisher, adding some functionality.
type PublisherDecorator func(pub Publisher) (Publisher, error)
// SubscriberDecorator wraps the underlying Subscriber, adding some functionality.
type SubscriberDecorator func(sub Subscriber) (Subscriber, error)
// RouterConfig holds the Router's configuration options.
type RouterConfig struct {
// CloseTimeout determines how long router should work for handlers when closing.
CloseTimeout time.Duration
}
func (c *RouterConfig) setDefaults() {
if c.CloseTimeout == 0 {
c.CloseTimeout = time.Second * 30
}
}
// Validate returns Router configuration error, if any.
func (c RouterConfig) Validate() error {
return nil
}
// NewRouter creates a new Router with given configuration.
func NewRouter(config RouterConfig, logger watermill.LoggerAdapter) (*Router, error) {
config.setDefaults()
if err := config.Validate(); err != nil {
return nil, errors.Wrap(err, "invalid config")
}
return newRouter(config, logger), nil
}
// NewDefaultRouter creates a new Router with default configuration.
func NewDefaultRouter(logger watermill.LoggerAdapter) *Router {
config := RouterConfig{}
config.setDefaults()
return newRouter(config, logger)
}
func newRouter(config RouterConfig, logger watermill.LoggerAdapter) *Router {
if logger == nil {
logger = watermill.NopLogger{}
}
return &Router{
config: config,
handlers: map[string]*handler{},
handlersWg: &sync.WaitGroup{},
runningHandlersWg: &sync.WaitGroup{},
runningHandlersWgLock: &sync.Mutex{},
handlerAdded: make(chan struct{}),
middlewaresLock: &sync.RWMutex{},
handlersLock: &sync.RWMutex{},
closingInProgressCh: make(chan struct{}),
closedCh: make(chan struct{}),
logger: logger,
running: make(chan struct{}),
}
}
type middleware struct {
Handler HandlerMiddleware
HandlerName string
IsRouterLevel bool
}
// Router is responsible for handling messages from subscribers using provided handler functions.
//
// If the handler function returns a message, the message is published with the publisher.
// You can use middlewares to wrap handlers with common logic like logging, instrumentation, etc.
type Router struct {
config RouterConfig
middlewares []middleware
middlewaresLock *sync.RWMutex
plugins []RouterPlugin
handlers map[string]*handler
handlersLock *sync.RWMutex
handlersWg *sync.WaitGroup
runningHandlersWg *sync.WaitGroup
runningHandlersWgLock *sync.Mutex
handlerAdded chan struct{}
closingInProgressCh chan struct{}
closedCh chan struct{}
closed bool
closedLock sync.Mutex
logger watermill.LoggerAdapter
publisherDecorators []PublisherDecorator
subscriberDecorators []SubscriberDecorator
isRunning bool
running chan struct{}
}
// Logger returns the Router's logger.
func (r *Router) Logger() watermill.LoggerAdapter {
return r.logger
}
// AddMiddleware adds a new middleware to the router.
//
// The order of middleware matters. Middleware added at the beginning is executed first.
func (r *Router) AddMiddleware(m ...HandlerMiddleware) {
r.logger.Debug("Adding middleware", watermill.LogFields{"count": fmt.Sprintf("%d", len(m))})
r.addRouterLevelMiddleware(m...)
}
func (r *Router) addRouterLevelMiddleware(m ...HandlerMiddleware) {
for _, handlerMiddleware := range m {
middleware := middleware{
Handler: handlerMiddleware,
HandlerName: "",
IsRouterLevel: true,
}
r.middlewares = append(r.middlewares, middleware)
}
}
func (r *Router) addHandlerLevelMiddleware(handlerName string, m ...HandlerMiddleware) {
r.middlewaresLock.Lock()
defer r.middlewaresLock.Unlock()
for _, handlerMiddleware := range m {
middleware := middleware{
Handler: handlerMiddleware,
HandlerName: handlerName,
IsRouterLevel: false,
}
r.middlewares = append(r.middlewares, middleware)
}
}
// AddPlugin adds a new plugin to the router.
// Plugins are executed during startup of the router.
//
// A plugin can, for example, close the router after SIGINT or SIGTERM is sent to the process (SignalsHandler plugin).
func (r *Router) AddPlugin(p ...RouterPlugin) {
r.logger.Debug("Adding plugins", watermill.LogFields{"count": fmt.Sprintf("%d", len(p))})
r.plugins = append(r.plugins, p...)
}
// AddPublisherDecorators wraps the router's Publisher.
// The first decorator is the innermost, i.e. calls the original publisher.
func (r *Router) AddPublisherDecorators(dec ...PublisherDecorator) {
r.logger.Debug("Adding publisher decorators", watermill.LogFields{"count": fmt.Sprintf("%d", len(dec))})
r.publisherDecorators = append(r.publisherDecorators, dec...)
}
// AddSubscriberDecorators wraps the router's Subscriber.
// The first decorator is the innermost, i.e. calls the original subscriber.
func (r *Router) AddSubscriberDecorators(dec ...SubscriberDecorator) {
r.logger.Debug("Adding subscriber decorators", watermill.LogFields{"count": fmt.Sprintf("%d", len(dec))})
r.subscriberDecorators = append(r.subscriberDecorators, dec...)
}
// Handlers returns all registered handlers.
func (r *Router) Handlers() map[string]HandlerFunc {
handlers := map[string]HandlerFunc{}
for handlerName, handler := range r.handlers {
handlers[handlerName] = handler.handlerFunc
}
return handlers
}
// DuplicateHandlerNameError is sent in a panic when you try to add a second handler with the same name.
type DuplicateHandlerNameError struct {
HandlerName string
}
func (d DuplicateHandlerNameError) Error() string {
return fmt.Sprintf("handler with name %s already exists", d.HandlerName)
}
// AddHandler adds a new handler.
//
// handlerName must be unique. For now, it is used only for debugging.
//
// subscribeTopic is a topic from which handler will receive messages.
//
// publishTopic is a topic to which router will produce messages returned by handlerFunc.
// When handler needs to publish to multiple topics,
// it is recommended to use AddConsumerHandler and inject a Publisher or implement middleware
// which will catch messages and publish to topic based on metadata for example.
//
// If handler is added while router is already running, you need to explicitly call RunHandlers().
func (r *Router) AddHandler(
handlerName string,
subscribeTopic string,
subscriber Subscriber,
publishTopic string,
publisher Publisher,
handlerFunc HandlerFunc,
) *Handler {
r.logger.Info("Adding handler", watermill.LogFields{
"handler_name": handlerName,
"topic": subscribeTopic,
})
r.handlersLock.Lock()
defer r.handlersLock.Unlock()
if _, ok := r.handlers[handlerName]; ok {
panic(DuplicateHandlerNameError{handlerName})
}
publisherName, subscriberName := internal.StructName(publisher), internal.StructName(subscriber)
newHandler := &handler{
name: handlerName,
logger: r.logger,
subscriber: subscriber,
subscribeTopic: subscribeTopic,
subscriberName: subscriberName,
publisher: publisher,
publishTopic: publishTopic,
publisherName: publisherName,
handlerFunc: handlerFunc,
runningHandlersWg: r.runningHandlersWg,
runningHandlersWgLock: r.runningHandlersWgLock,
messagesCh: nil,
routersCloseCh: r.closingInProgressCh,
startedCh: make(chan struct{}),
}
r.handlersWg.Add(1)
r.handlers[handlerName] = newHandler
select {
case r.handlerAdded <- struct{}{}:
default:
// watchAllHandlersStopped is not always waiting for handlerAdded
}
return &Handler{
router: r,
handler: newHandler,
}
}
// AddConsumerHandler adds a new handler that does not return any messages.
// It can publish messages by directly using a Publisher.
//
// handlerName must be unique. For now, it is used only for debugging.
//
// subscribeTopic is a topic from which handler will receive messages.
//
// subscriber is Subscriber from which messages will be consumed.
//
// If handler is added while router is already running, you need to explicitly call RunHandlers().
func (r *Router) AddConsumerHandler(
handlerName string,
subscribeTopic string,
subscriber Subscriber,
handlerFunc NoPublishHandlerFunc,
) *Handler {
handlerFuncAdapter := func(msg *Message) ([]*Message, error) {
return nil, handlerFunc(msg)
}
return r.AddHandler(handlerName, subscribeTopic, subscriber, "", disabledPublisher{}, handlerFuncAdapter)
}
// AddNoPublisherHandler adds a new handler.
// This handler cannot return messages.
//
// handlerName must be unique. For now, it is used only for debugging.
//
// subscribeTopic is a topic from which handler will receive messages.
//
// subscriber is Subscriber from which messages will be consumed.
//
// If handler is added while router is already running, you need to explicitly call RunHandlers().
//
// Deprecated: use AddConsumerHandler instead.
func (r *Router) AddNoPublisherHandler(
handlerName string,
subscribeTopic string,
subscriber Subscriber,
handlerFunc NoPublishHandlerFunc,
) *Handler {
return r.AddConsumerHandler(handlerName, subscribeTopic, subscriber, handlerFunc)
}
// Run runs all plugins and handlers and starts subscribing to provided topics.
// This call is blocking while the router is running.
//
// When all handlers have stopped (for example, because subscriptions were closed), the router will also stop.
//
// To stop Run() you should call Close() on the router.
//
// ctx will be propagated to all subscribers.
//
// When all handlers are stopped (for example: because of closed connection), Run() will be also stopped.
func (r *Router) Run(ctx context.Context) (err error) {
if r.isRunning {
return errors.New("router is already running")
}
r.isRunning = true
ctx, cancel := context.WithCancel(ctx)
defer cancel()
r.logger.Debug("Loading plugins", nil)
for _, plugin := range r.plugins {
if err := plugin(r); err != nil {
return errors.Wrapf(err, "cannot initialize plugin %v", plugin)
}
}
r.watchAllHandlersStopped(ctx)
if err := r.RunHandlers(ctx); err != nil {
return err
}
close(r.running)
<-r.closingInProgressCh
cancel()
r.logger.Info("Waiting for messages", watermill.LogFields{
"timeout": r.config.CloseTimeout,
})
<-r.closedCh
r.logger.Info("All messages processed", nil)
return nil
}
// RunHandlers runs all handlers that were added after Run().
// RunHandlers is idempotent, so can be called multiple times safely.
func (r *Router) RunHandlers(ctx context.Context) error {
if !r.isRunning {
return errors.New("you can't call RunHandlers on non-running router")
}
r.handlersLock.Lock()
defer r.handlersLock.Unlock()
r.logger.Info("Running router handlers", watermill.LogFields{"count": len(r.handlers)})
for name, h := range r.handlers {
if h.started {
continue
}
if err := r.decorateHandlerPublisher(h); err != nil {
return errors.Wrapf(err, "could not decorate publisher of handler %s", name)
}
if err := r.decorateHandlerSubscriber(h); err != nil {
return errors.Wrapf(err, "could not decorate subscriber of handler %s", name)
}
logger := r.logger.With(watermill.LogFields{
"subscriber_name": h.name,
"topic": h.subscribeTopic,
})
logger.Debug("Subscribing to topic", nil)
ctx, cancel := context.WithCancel(ctx)
messages, err := h.subscriber.Subscribe(ctx, h.subscribeTopic)
if err != nil {
cancel()
return errors.Wrapf(err, "cannot subscribe topic %s", h.subscribeTopic)
}
h.messagesCh = messages
h.started = true
close(h.startedCh)
h.stopFn = cancel
h.stopped = make(chan struct{})
go func() {
defer cancel()
r.middlewaresLock.Lock()
middlewares := append([]middleware{}, r.middlewares...)
r.middlewaresLock.Unlock()
h.run(ctx, middlewares)
r.handlersWg.Done()
logger.Info("Subscriber stopped", nil)
r.handlersLock.Lock()
delete(r.handlers, name)
r.handlersLock.Unlock()
logger.Trace("Removed subscriber from r.handlers", nil)
close(h.stopped)
}()
}
return nil
}
// watchAllHandlersStopped closes router when all handlers have stopped,
// (for example, because for example all subscriptions are closed)
func (r *Router) watchAllHandlersStopped(ctx context.Context) {
r.handlersLock.RLock()
hasNoHandlersYet := len(r.handlers) == 0
r.handlersLock.RUnlock()
go func() {
if hasNoHandlersYet {
// we can start router without any handlers,
// in that situation router would be closed immediately (even if they are no routers)
// let's wait for
select {
case <-r.handlerAdded:
// it should be some handler to track
case <-r.closedCh:
// let's avoid goroutine leak
return
}
}
r.handlersWg.Wait()
if r.IsClosed() {
r.logger.Trace("watchAllHandlersStopped: already closed", nil)
// already closed
return
}
// Only log an error if the context was not canceled, but handlers were stopped.
select {
case <-ctx.Done():
default:
r.logger.Error("All handlers stopped, closing router", errors.New("all router handlers stopped"), nil)
}
if err := r.Close(); err != nil {
r.logger.Error("Cannot close router", err, nil)
}
}()
}
// Running is closed when router is running.
// In other words: you can wait till router is running using
//
// fmt.Println("Starting router")
// go r.Run(ctx)
// <- r.Running()
// fmt.Println("Router is running")
//
// Warning: for historical reasons, this channel is not aware of router closing - the channel will be closed if the router has been running and closed.
func (r *Router) Running() chan struct{} {
return r.running
}
// IsRunning returns true when router is running.
//
// Warning: for historical reasons, this method is not aware of router closing.
// If you want to know if the router was closed, use IsClosed.
func (r *Router) IsRunning() bool {
select {
case <-r.running:
return true
default:
return false
}
}
// Close gracefully closes the router with a timeout provided in the configuration.
func (r *Router) Close() error {
r.closedLock.Lock()
defer r.closedLock.Unlock()
r.handlersLock.Lock()
defer r.handlersLock.Unlock()
if r.closed {
r.logger.Debug("Already closed", nil)
return nil
}
r.logger.Debug("Running Close()", nil)
r.closed = true
r.logger.Info("Closing router", nil)
defer r.logger.Info("Router closed", nil)
close(r.closingInProgressCh)
defer close(r.closedCh)
timedout := r.waitForHandlers()
if timedout {
return errors.New("router close timeout")
}
return nil
}
func (r *Router) waitForHandlers() bool {
var waitGroup sync.WaitGroup
waitGroup.Add(1)
go func() {
defer waitGroup.Done()
r.handlersWg.Wait()
}()
waitGroup.Add(1)
go func() {
defer waitGroup.Done()
r.runningHandlersWgLock.Lock()
defer r.runningHandlersWgLock.Unlock()
r.runningHandlersWg.Wait()
}()
return sync_internal.WaitGroupTimeout(&waitGroup, r.config.CloseTimeout)
}
func (r *Router) IsClosed() bool {
r.closedLock.Lock()
defer r.closedLock.Unlock()
return r.closed
}
type handler struct {
name string
logger watermill.LoggerAdapter
subscriber Subscriber
subscribeTopic string
subscriberName string
publisher Publisher
publishTopic string
publisherName string
handlerFunc HandlerFunc
runningHandlersWg *sync.WaitGroup
runningHandlersWgLock *sync.Mutex
messagesCh <-chan *Message
started bool
startedCh chan struct{}
stopFn context.CancelFunc
stopped chan struct{}
routersCloseCh chan struct{}
}
func (h *handler) run(ctx context.Context, middlewares []middleware) {
h.logger.Info("Starting handler", watermill.LogFields{
"subscriber_name": h.name,
"topic": h.subscribeTopic,
})
middlewareHandler := h.handlerFunc
// first added middlewares should be executed first (so should be at the top of call stack)
for i := len(middlewares) - 1; i >= 0; i-- {
currentMiddleware := middlewares[i]
isValidHandlerLevelMiddleware := currentMiddleware.HandlerName == h.name
if currentMiddleware.IsRouterLevel || isValidHandlerLevelMiddleware {
middlewareHandler = currentMiddleware.Handler(middlewareHandler)
}
}
go h.handleClose(ctx)
for msg := range h.messagesCh {
h.runningHandlersWgLock.Lock()
h.runningHandlersWg.Add(1)
h.runningHandlersWgLock.Unlock()
go h.handleMessage(msg, middlewareHandler)
}
if h.publisher != nil {
h.logger.Debug("Waiting for publisher to close", nil)
if err := h.publisher.Close(); err != nil {
h.logger.Error("Failed to close publisher", err, nil)
}
h.logger.Debug("Publisher closed", nil)
}
h.logger.Debug("Router handler stopped", nil)
}
// Handler handles Messages.
type Handler struct {
router *Router
handler *handler
}
// AddMiddleware adds new middleware to the specified handler in the router.
//
// The order of middleware matters. Middleware added at the beginning is executed first.
func (h *Handler) AddMiddleware(m ...HandlerMiddleware) {
handler := h.handler
handler.logger.Debug("Adding middleware to handler", watermill.LogFields{
"count": fmt.Sprintf("%d", len(m)),
"handlerName": handler.name,
})
h.router.addHandlerLevelMiddleware(handler.name, m...)
}
// Started returns channel which is stopped when handler is running.
func (h *Handler) Started() chan struct{} {
return h.handler.startedCh
}
// Stop stops the handler.
// Stop is asynchronous.
// You can check if handler was stopped with Stopped() function.
func (h *Handler) Stop() {
if !h.handler.started {
panic("handler is not started")
}
h.handler.stopFn()
}
// Stopped returns channel which is stopped when handler did stop.
func (h *Handler) Stopped() chan struct{} {
return h.handler.stopped
}
// decorateHandlerPublisher applies the decorator chain to handler's publisher.
// They are applied in reverse order, so that the later decorators use the result of former ones.
func (r *Router) decorateHandlerPublisher(h *handler) error {
var err error
pub := h.publisher
for i := len(r.publisherDecorators) - 1; i >= 0; i-- {
decorator := r.publisherDecorators[i]
pub, err = decorator(pub)
if err != nil {
return errors.Wrap(err, "could not apply publisher decorator")
}
}
r.handlers[h.name].publisher = pub
return nil
}
// decorateHandlerSubscriber applies the decorator chain to handler's subscriber.
// They are applied in regular order, so that the later decorators use the result of former ones.
func (r *Router) decorateHandlerSubscriber(h *handler) error {
var err error
sub := h.subscriber
// add values to message context to subscriber
// it goes before other decorators, so that they may take advantage of these values
messageTransform := func(msg *Message) {
if msg != nil {
h.addHandlerContext(msg)
}
}
sub, err = MessageTransformSubscriberDecorator(messageTransform)(sub)
if err != nil {
return errors.Wrapf(err, "cannot wrap subscriber with context decorator")
}
for _, decorator := range r.subscriberDecorators {
sub, err = decorator(sub)
if err != nil {
return errors.Wrap(err, "could not apply subscriber decorator")
}
}
r.handlers[h.name].subscriber = sub
return nil
}
// addHandlerContext enriches the context with values that are relevant within this handler's context.
func (h *handler) addHandlerContext(messages ...*Message) {
for i, msg := range messages {
ctx := msg.Context()
if h.name != "" {
ctx = context.WithValue(ctx, handlerNameKey, h.name)
}
if h.publisherName != "" {
ctx = context.WithValue(ctx, publisherNameKey, h.publisherName)
}
if h.subscriberName != "" {
ctx = context.WithValue(ctx, subscriberNameKey, h.subscriberName)
}
if h.subscribeTopic != "" {
ctx = context.WithValue(ctx, subscribeTopicKey, h.subscribeTopic)
}
if h.publishTopic != "" {
ctx = context.WithValue(ctx, publishTopicKey, h.publishTopic)
}
messages[i].SetContext(ctx)
}
}
func (h *handler) handleClose(ctx context.Context) {
closeSubscriber := func() {
h.logger.Debug("Waiting for subscriber to close", nil)
if err := h.subscriber.Close(); err != nil {
h.logger.Error("Failed to close subscriber", err, nil)
}
h.logger.Debug("Subscriber closed", nil)
}
select {
case <-h.routersCloseCh:
closeSubscriber()
case <-ctx.Done():
// ctx.Done can race with routersCloseCh during a router-wide close:
// Router.Close closes routersCloseCh, then Router.Run cancels the parent ctx,
// and the select picks whichever fires first. If routersCloseCh is also closed
// here, this is a router-wide close and the subscriber must be closed —
// otherwise it's Handler.Stop() and the subscriber may be shared.
select {
case <-h.routersCloseCh:
closeSubscriber()
default:
}
}
h.stopFn()
}
func (h *handler) handleMessage(msg *Message, handler HandlerFunc) {
defer h.runningHandlersWg.Done()
msgFields := watermill.LogFields{"message_uuid": msg.UUID, "handler_name": h.name}
defer func() {
if recovered := recover(); recovered != nil {
h.logger.Error(
"Panic recovered in handler. Stack: "+string(debug.Stack()),
errors.Errorf("%s", recovered),
msgFields,
)
msg.Nack()
}
}()
h.logger.Trace("Received message", msgFields)
producedMessages, err := handler(msg)
if err != nil {
if !errors.Is(err, context.Canceled) {
h.logger.Error("Handler returned error", err, msgFields)
}
msg.Nack()
return
}
h.addHandlerContext(producedMessages...)
if err := h.publishProducedMessages(producedMessages, msgFields); err != nil {
h.logger.Error("Publishing produced messages failed", err, nil)
msg.Nack()
return
}
msg.Ack()
h.logger.Trace("Message acked", msgFields)
}
func (h *handler) publishProducedMessages(producedMessages Messages, msgFields watermill.LogFields) error {
if len(producedMessages) == 0 {
return nil
}
if h.publisher == nil {
return ErrOutputInNoPublisherHandler
}
h.logger.Trace("Sending produced messages", msgFields.Add(watermill.LogFields{
"produced_messages_count": len(producedMessages),
"publish_topic": h.publishTopic,
}))
if err := h.publisher.Publish(h.publishTopic, producedMessages...); err != nil {
h.logger.Error("Cannot publish messages", err, msgFields.Add(watermill.LogFields{
"not_sent_message": fmt.Sprintf("%#v", producedMessages),
}))
return err
}
return nil
}
type disabledPublisher struct{}
func (disabledPublisher) Publish(topic string, messages ...*Message) error {
return ErrOutputInNoPublisherHandler
}
func (disabledPublisher) Close() error {
return nil
}