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

     1  package amqp
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/pkg/errors"
     7  	"github.com/wfusion/gofusion/common/infra/watermill"
     8  
     9  	"github.com/wfusion/gofusion/common/infra/watermill/message"
    10  
    11  	amqp "github.com/rabbitmq/amqp091-go"
    12  )
    13  
    14  const DefaultMessageUUIDHeaderKey = "_watermill_message_uuid"
    15  
    16  // deprecated, please use DefaultMessageUUIDHeaderKey instead
    17  const MessageUUIDHeaderKey = DefaultMessageUUIDHeaderKey
    18  
    19  // Marshaler marshals Watermill's message to amqp.Publishing and unmarshals amqp.Delivery to Watermill's message.
    20  type Marshaler interface {
    21  	Marshal(msg *message.Message) (amqp.Publishing, error)
    22  	Unmarshal(amqpMsg amqp.Delivery) (*message.Message, error)
    23  }
    24  
    25  type DefaultMarshaler struct {
    26  	// PostprocessPublishing can be used to make some extra processing with amqp.Publishing,
    27  	// for example add CorrelationId and ContentType:
    28  	//
    29  	//  amqp.DefaultMarshaler{
    30  	//		PostprocessPublishing: func(publishing stdAmqp.Publishing) stdAmqp.Publishing {
    31  	//			publishing.CorrelationId = "correlation"
    32  	//			publishing.ContentType = "application/json"
    33  	//
    34  	//			return publishing
    35  	//		},
    36  	//	}
    37  	PostprocessPublishing func(amqp.Publishing) amqp.Publishing
    38  
    39  	// When true, DeliveryMode will be not set to Persistent.
    40  	//
    41  	// DeliveryMode Transient means higher throughput, but messages will not be
    42  	// restored on broker restart. The delivery mode of publishings is unrelated
    43  	// to the durability of the queues they reside on. Transient messages will
    44  	// not be restored to durable queues, persistent messages will be restored to
    45  	// durable queues and lost on non-durable queues during server restart.
    46  	NotPersistentDeliveryMode bool
    47  
    48  	// Header used to store and read message UUID.
    49  	//
    50  	// If value is empty, DefaultMessageUUIDHeaderKey value is used.
    51  	// If header doesn't exist, empty value is passed as message UUID.
    52  	MessageUUIDHeaderKey string
    53  
    54  	AppID string
    55  }
    56  
    57  func (d DefaultMarshaler) Marshal(msg *message.Message) (amqp.Publishing, error) {
    58  	headers := make(amqp.Table, len(msg.Metadata)+1) // metadata + plus uuid
    59  
    60  	for key, value := range msg.Metadata {
    61  		headers[key] = value
    62  	}
    63  	headers[d.computeMessageUUIDHeaderKey()] = msg.UUID
    64  	headers[watermill.MessageHeaderAppID] = d.AppID
    65  
    66  	publishing := amqp.Publishing{
    67  		Body:         msg.Payload,
    68  		Headers:      headers,
    69  		AppId:        d.AppID,
    70  		MessageId:    msg.UUID,
    71  		Timestamp:    time.Now(),
    72  		DeliveryMode: amqp.Transient,
    73  	}
    74  	if !d.NotPersistentDeliveryMode {
    75  		publishing.DeliveryMode = amqp.Persistent
    76  	}
    77  
    78  	if d.PostprocessPublishing != nil {
    79  		publishing = d.PostprocessPublishing(publishing)
    80  	}
    81  
    82  	return publishing, nil
    83  }
    84  
    85  func (d DefaultMarshaler) Unmarshal(amqpMsg amqp.Delivery) (*message.Message, error) {
    86  	msgUUIDStr, err := d.unmarshalMessageUUID(amqpMsg)
    87  	if err != nil {
    88  		return nil, err
    89  	}
    90  
    91  	msg := message.NewMessage(msgUUIDStr, amqpMsg.Body)
    92  	msg.Metadata = make(message.Metadata, len(amqpMsg.Headers)-1) // headers - minus uuid
    93  
    94  	for key, value := range amqpMsg.Headers {
    95  		if key == d.computeMessageUUIDHeaderKey() {
    96  			continue
    97  		}
    98  
    99  		var ok bool
   100  		msg.Metadata[key], ok = value.(string)
   101  		if !ok {
   102  			return nil, errors.Errorf("metadata %s is not a string, but %#v", key, value)
   103  		}
   104  	}
   105  
   106  	return msg, nil
   107  }
   108  
   109  func (d DefaultMarshaler) unmarshalMessageUUID(amqpMsg amqp.Delivery) (string, error) {
   110  	var msgUUIDStr string
   111  
   112  	msgUUID, hasMsgUUID := amqpMsg.Headers[d.computeMessageUUIDHeaderKey()]
   113  	if !hasMsgUUID {
   114  		return "", nil
   115  	}
   116  
   117  	msgUUIDStr, hasMsgUUID = msgUUID.(string)
   118  	if !hasMsgUUID {
   119  		return "", errors.Errorf("message UUID is not a string, but: %#v", msgUUID)
   120  	}
   121  
   122  	return msgUUIDStr, nil
   123  }
   124  
   125  func (d DefaultMarshaler) computeMessageUUIDHeaderKey() string {
   126  	if d.MessageUUIDHeaderKey != "" {
   127  		return d.MessageUUIDHeaderKey
   128  	}
   129  
   130  	return DefaultMessageUUIDHeaderKey
   131  }