snapshot of https://github.com/ThreeDotsLabs/watermill.git (master, 2026-09-15; blobs > 131072B omitted)

message/router.go

@@ -0,0 +1,888 @@
+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
+}