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

     1  package message
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // MessageTransformSubscriberDecorator creates a subscriber decorator that calls transform
     9  // on each message that passes through the subscriber.
    10  func MessageTransformSubscriberDecorator(transform func(*Message)) SubscriberDecorator {
    11  	if transform == nil {
    12  		panic("transform function is nil")
    13  	}
    14  	return func(sub Subscriber) (Subscriber, error) {
    15  		return &messageTransformSubscriberDecorator{
    16  			sub:       sub,
    17  			transform: transform,
    18  		}, nil
    19  	}
    20  }
    21  
    22  // MessageTransformPublisherDecorator creates a publisher decorator that calls transform
    23  // on each message that passes through the publisher.
    24  func MessageTransformPublisherDecorator(transform func(*Message)) PublisherDecorator {
    25  	if transform == nil {
    26  		panic("transform function is nil")
    27  	}
    28  	return func(pub Publisher) (Publisher, error) {
    29  		return &messageTransformPublisherDecorator{
    30  			Publisher: pub,
    31  			transform: transform,
    32  		}, nil
    33  	}
    34  }
    35  
    36  type messageTransformSubscriberDecorator struct {
    37  	sub Subscriber
    38  
    39  	transform   func(*Message)
    40  	subscribeWg sync.WaitGroup
    41  }
    42  
    43  func (t *messageTransformSubscriberDecorator) Subscribe(ctx context.Context, topic string) (<-chan *Message, error) {
    44  	in, err := t.sub.Subscribe(ctx, topic)
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  
    49  	out := make(chan *Message)
    50  	t.subscribeWg.Add(1)
    51  	go func() {
    52  		for msg := range in {
    53  			t.transform(msg)
    54  			out <- msg
    55  		}
    56  		close(out)
    57  		t.subscribeWg.Done()
    58  	}()
    59  
    60  	return out, nil
    61  }
    62  
    63  func (t *messageTransformSubscriberDecorator) Close() error {
    64  	err := t.sub.Close()
    65  
    66  	t.subscribeWg.Wait()
    67  	return err
    68  }
    69  
    70  type messageTransformPublisherDecorator struct {
    71  	Publisher
    72  	transform func(*Message)
    73  }
    74  
    75  // Publish applies the transform to each message and returns the underlying Publisher's result.
    76  func (d messageTransformPublisherDecorator) Publish(ctx context.Context, topic string, messages ...*Message) error {
    77  	for i := range messages {
    78  		d.transform(messages[i])
    79  	}
    80  	return d.Publisher.Publish(ctx, topic, messages...)
    81  }