Merge pull request #62 from w-h-a/pubsub

feat: pub sub
w-h-a · Aug 29, 2024 · ef15e57 · ef15e57
2 parents 0e818a0 + fd92f47
commit ef15e57
Show file tree

Hide file tree

Showing 10 changed files with 425 additions and 229 deletions.
diff --git a/broker/memory/memory.go b/broker/memory/memory.go
@@ -27,9 +27,16 @@ func (b *memory) Publish(data interface{}, options broker.PublishOptions) error
 	}
 	b.mtx.RUnlock()
 
-	bs, err := json.Marshal(data)
-	if err != nil {
-		return err
+	var bs []byte
+
+	if p, ok := data.([]byte); ok {
+		bs = p
+	} else {
+		p, err := json.Marshal(data)
+		if err != nil {
+			return err
+		}
+		bs = p
 	}
 
 	for _, sub := range subsOfThisTopic {

diff --git a/broker/options.go b/broker/options.go
@@ -6,8 +6,8 @@ type BrokerOption func(o *BrokerOptions)
 
 type BrokerOptions struct {
 	Nodes            []string
-	PublishOptions   PublishOptions
-	SubscribeOptions SubscribeOptions
+	PublishOptions   *PublishOptions
+	SubscribeOptions *SubscribeOptions
 	Context          context.Context
 }
 
@@ -17,13 +17,13 @@ func BrokerWithNodes(addrs ...string) BrokerOption {
 	}
 }
 
-func BrokerWithPublishOptions(options PublishOptions) BrokerOption {
+func BrokerWithPublishOptions(options *PublishOptions) BrokerOption {
 	return func(o *BrokerOptions) {
 		o.PublishOptions = options
 	}
 }
 
-func BrokerWithSubscribeOptions(options SubscribeOptions) BrokerOption {
+func BrokerWithSubscribeOptions(options *SubscribeOptions) BrokerOption {
 	return func(o *BrokerOptions) {
 		o.SubscribeOptions = options
 	}

diff --git a/broker/snssqs/clients.go b/broker/snssqs/clients.go
@@ -2,6 +2,7 @@ package snssqs
 
 import (
 	"context"
+	"encoding/json"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/sns"
@@ -37,13 +38,18 @@ type SqsClient interface {
 
 type sqsClient struct {
 	*sqs.Client
+	queueUrl          *string
 	visibilityTimeout int32
 	waitTimeSeconds   int32
 }
 
+type sqsMsg struct {
+	Message string `json:"message"`
+}
+
 func (c *sqsClient) ConsumeFromGroup(sub broker.Subscriber) {
 	result, err := c.ReceiveMessage(context.Background(), &sqs.ReceiveMessageInput{
-		QueueUrl:              aws.String(sub.Options().Group),
+		QueueUrl:              c.queueUrl,
 		MaxNumberOfMessages:   1,
 		VisibilityTimeout:     c.visibilityTimeout,
 		WaitTimeSeconds:       c.waitTimeSeconds,
@@ -60,14 +66,23 @@ func (c *sqsClient) ConsumeFromGroup(sub broker.Subscriber) {
 
 	for _, msg := range result.Messages {
 		body := msg.Body
-		if err := sub.Handler([]byte(*body)); err != nil {
+
+		var sqsMsg sqsMsg
+
+		if err := json.Unmarshal([]byte(*body), &sqsMsg); err != nil {
+			log.Errorf("failed to unmarshal message from group %s: %s", sub.Options().Group, err)
+			continue
+		}
+
+		if err := sub.Handler([]byte(sqsMsg.Message)); err != nil {
 			log.Errorf("failed to handle message from group %s: %s", sub.Options().Group, err)
-		} else {
-			msgHandle := msg.ReceiptHandle
-			c.DeleteMessage(context.Background(), &sqs.DeleteMessageInput{
-				QueueUrl:      aws.String(sub.Options().Group),
-				ReceiptHandle: msgHandle,
-			})
+			continue
 		}
+
+		msgHandle := msg.ReceiptHandle
+		c.DeleteMessage(context.Background(), &sqs.DeleteMessageInput{
+			QueueUrl:      aws.String(sub.Options().Group),
+			ReceiptHandle: msgHandle,
+		})
 	}
 }
diff --git a/broker/snssqs/snssqs.go b/broker/snssqs/snssqs.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	awsconfig "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/sns"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
@@ -14,7 +15,7 @@ import (
 )
 
 const (
-	defaultVisibilityTimeout int32 = 4
+	defaultVisibilityTimeout int32 = 8
 	defaultWaitSeconds       int32 = 8
 )
 
@@ -29,9 +30,16 @@ func (b *snssqs) Options() broker.BrokerOptions {
 }
 
 func (b *snssqs) Publish(data interface{}, options broker.PublishOptions) error {
-	bs, err := json.Marshal(data)
-	if err != nil {
-		return err
+	var bs []byte
+
+	if p, ok := data.([]byte); ok {
+		bs = p
+	} else {
+		p, err := json.Marshal(data)
+		if err != nil {
+			return err
+		}
+		bs = p
 	}
 
 	if err := b.snsClient.ProduceToTopic(bs, options.Topic); err != nil {
@@ -77,7 +85,7 @@ func (b *snssqs) configure() error {
 		b.sqsClient = sqs
 	}
 
-	if b.snsClient != nil && b.sqsClient != nil {
+	if b.snsClient != nil || b.sqsClient != nil {
 		return nil
 	}
 
@@ -86,21 +94,34 @@ func (b *snssqs) configure() error {
 		return err
 	}
 
-	b.snsClient = &snsClient{sns.NewFromConfig(cfg)}
+	if b.options.PublishOptions != nil {
+		b.snsClient = &snsClient{sns.NewFromConfig(cfg)}
+	}
 
-	visibilityTimeout := defaultVisibilityTimeout
+	if b.options.SubscribeOptions != nil {
+		visibilityTimeout := defaultVisibilityTimeout
 
-	waitTimeSeconds := defaultWaitSeconds
+		waitTimeSeconds := defaultWaitSeconds
 
-	if timeout, ok := GetVisibilityTimeoutFromContext(b.options.SubscribeOptions.Context); ok {
-		visibilityTimeout = timeout
-	}
+		if timeout, ok := GetVisibilityTimeoutFromContext(b.options.SubscribeOptions.Context); ok {
+			visibilityTimeout = timeout
+		}
 
-	if waitTime, ok := GetWaitTimeSecondsFromContext(b.options.SubscribeOptions.Context); ok {
-		waitTimeSeconds = waitTime
-	}
+		if waitTime, ok := GetWaitTimeSecondsFromContext(b.options.SubscribeOptions.Context); ok {
+			waitTimeSeconds = waitTime
+		}
+
+		client := sqs.NewFromConfig(cfg)
 
-	b.sqsClient = &sqsClient{sqs.NewFromConfig(cfg), visibilityTimeout, waitTimeSeconds}
+		url, err := client.GetQueueUrl(context.Background(), &sqs.GetQueueUrlInput{
+			QueueName: aws.String(b.options.SubscribeOptions.Group),
+		})
+		if err != nil {
+			return err
+		}
+
+		b.sqsClient = &sqsClient{client, url.QueueUrl, visibilityTimeout, waitTimeSeconds}
+	}
 
 	return nil
 }