github.com/wfusion/gofusion@v1.1.14/common/infra/watermill/message/router/middleware/correlation.go (about)

     1  package middleware
     2  
     3  import (
     4  	"github.com/wfusion/gofusion/common/infra/watermill/message"
     5  )
     6  
     7  // CorrelationIDMetadataKey is used to store the correlation ID in metadata.
     8  const CorrelationIDMetadataKey = "correlation_id"
     9  
    10  // SetCorrelationID sets a correlation ID for the message.
    11  //
    12  // SetCorrelationID should be called when the message enters the system.
    13  // When message is produced in a request (for example HTTP),
    14  // message correlation ID should be the same as the request's correlation ID.
    15  func SetCorrelationID(id string, msg *message.Message) {
    16  	if MessageCorrelationID(msg) != "" {
    17  		return
    18  	}
    19  
    20  	msg.Metadata.Set(CorrelationIDMetadataKey, id)
    21  }
    22  
    23  // MessageCorrelationID returns correlation ID from the message.
    24  func MessageCorrelationID(message *message.Message) string {
    25  	return message.Metadata.Get(CorrelationIDMetadataKey)
    26  }
    27  
    28  // CorrelationID adds correlation ID to all messages produced by the handler.
    29  // ID is based on ID from message received by handler.
    30  //
    31  // To make CorrelationID working correctly, SetCorrelationID must be called to first message entering the system.
    32  func CorrelationID(h message.HandlerFunc) message.HandlerFunc {
    33  	return func(message *message.Message) ([]*message.Message, error) {
    34  		producedMessages, err := h(message)
    35  
    36  		correlationID := MessageCorrelationID(message)
    37  		for _, msg := range producedMessages {
    38  			SetCorrelationID(correlationID, msg)
    39  		}
    40  
    41  		return producedMessages, err
    42  	}
    43  }