82 lines
1.9 KiB
Go
82 lines
1.9 KiB
Go
package message
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
)
|
|
|
|
// MessageTransformSubscriberDecorator creates a subscriber decorator that calls transform
|
|
// on each message that passes through the subscriber.
|
|
func MessageTransformSubscriberDecorator(transform func(*Message)) SubscriberDecorator {
|
|
if transform == nil {
|
|
panic("transform function is nil")
|
|
}
|
|
return func(sub Subscriber) (Subscriber, error) {
|
|
return &messageTransformSubscriberDecorator{
|
|
sub: sub,
|
|
transform: transform,
|
|
}, nil
|
|
}
|
|
}
|
|
|
|
// MessageTransformPublisherDecorator creates a publisher decorator that calls transform
|
|
// on each message that passes through the publisher.
|
|
func MessageTransformPublisherDecorator(transform func(*Message)) PublisherDecorator {
|
|
if transform == nil {
|
|
panic("transform function is nil")
|
|
}
|
|
return func(pub Publisher) (Publisher, error) {
|
|
return &messageTransformPublisherDecorator{
|
|
Publisher: pub,
|
|
transform: transform,
|
|
}, nil
|
|
}
|
|
}
|
|
|
|
type messageTransformSubscriberDecorator struct {
|
|
sub Subscriber
|
|
|
|
transform func(*Message)
|
|
subscribeWg sync.WaitGroup
|
|
}
|
|
|
|
func (t *messageTransformSubscriberDecorator) Subscribe(ctx context.Context, topic string) (<-chan *Message, error) {
|
|
in, err := t.sub.Subscribe(ctx, topic)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
out := make(chan *Message)
|
|
t.subscribeWg.Add(1)
|
|
go func() {
|
|
for msg := range in {
|
|
t.transform(msg)
|
|
out <- msg
|
|
}
|
|
close(out)
|
|
t.subscribeWg.Done()
|
|
}()
|
|
|
|
return out, nil
|
|
}
|
|
|
|
func (t *messageTransformSubscriberDecorator) Close() error {
|
|
err := t.sub.Close()
|
|
|
|
t.subscribeWg.Wait()
|
|
return err
|
|
}
|
|
|
|
type messageTransformPublisherDecorator struct {
|
|
Publisher
|
|
transform func(*Message)
|
|
}
|
|
|
|
// Publish applies the transform to each message and returns the underlying Publisher's result.
|
|
func (d messageTransformPublisherDecorator) Publish(topic string, messages ...*Message) error {
|
|
for i := range messages {
|
|
d.transform(messages[i])
|
|
}
|
|
return d.Publisher.Publish(topic, messages...)
|
|
}
|