feat(PubSub): Add SQS PubSub wrappers

Author: Neurostep

pkg/pubsub/sqs/publisher.go

@@ -0,0 +1,29 @@
+package sqs
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/service/sqs"
+)
+
+type (
+	Publisher struct {
+		client   *sqs.Client
+		queueURL string
+	}
+)
+
+func NewPublisher(sqsClient *sqs.Client, queueURL string) *Publisher {
+	return &Publisher{
+		client:   sqsClient,
+		queueURL: queueURL,
+	}
+}
+
+func (p *Publisher) Publish(ctx context.Context, msg *sqs.SendMessageInput) (*sqs.SendMessageOutput, error) {
+	if msg.QueueUrl == nil {
+		msg.QueueUrl = &p.queueURL
+	}
+
+	return p.client.SendMessage(ctx, msg)
+}

pkg/pubsub/sqs/subscriber.go

@@ -0,0 +1,119 @@
+package sqs
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/sqs"
+	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
+
+	"github.com/scribd/go-sdk/pkg/pubsub"
+	"github.com/scribd/go-sdk/pkg/pubsub/pool"
+)
+
+type (
+	SQSClient interface {
+		ReceiveMessage(
+			ctx context.Context,
+			params *sqs.ReceiveMessageInput, optFns ...func(*sqs.Options)) (*sqs.ReceiveMessageOutput, error)
+	}
+
+	Subscriber struct {
+		client      SQSClient
+		queueURL    string
+		handler     MsgHandler
+		maxMessages int
+		pool        *pool.Pool
+		waitTime    time.Duration
+		wg          sync.WaitGroup // tracks active handlers
+		stopCh      chan struct{}
+	}
+
+	SubscriberConfig struct {
+		SQSClient  *sqs.Client
+		MsgHandler MsgHandler
+		SQSConfig  pubsub.SQS
+	}
+
+	MsgHandler func(msg types.Message)
+)
+
+const (
+	defaultNumWorkers = 1
+	maxWaitTime       = 20 * time.Second
+)
+
+func NewSubscriber(c SubscriberConfig) *Subscriber {
+	workers := c.SQSConfig.Subscriber.Workers
+	if workers == 0 {
+		workers = defaultNumWorkers
+	}
+
+	waitTime := c.SQSConfig.Subscriber.WaitTime
+	if waitTime > maxWaitTime {
+		waitTime = maxWaitTime
+	}
+
+	return &Subscriber{
+		client:      c.SQSClient,
+		handler:     c.MsgHandler,
+		maxMessages: c.SQSConfig.Subscriber.MaxMessages,
+		queueURL:    c.SQSConfig.Subscriber.QueueURL,
+		pool:        pool.New(workers),
+		waitTime:    waitTime,
+		stopCh:      make(chan struct{}),
+	}
+}
+
+func (s *Subscriber) Subscribe(ctx context.Context) chan error {
+	ch := make(chan error)
+
+	req := &sqs.ReceiveMessageInput{
+		QueueUrl:              aws.String(s.queueURL),
+		MaxNumberOfMessages:   int32(s.maxMessages),
+		MessageAttributeNames: []string{"All"},
+		MessageSystemAttributeNames: []types.MessageSystemAttributeName{
+			types.MessageSystemAttributeNameAll,
+		},
+	}
+	if s.waitTime > 0 {
+		req.WaitTimeSeconds = int32(s.waitTime.Seconds())
+	}
+
+	go func() {
+		defer close(ch)
+
+		for {
+			select {
+			case <-s.stopCh:
+				return
+			default:
+				response, err := s.client.ReceiveMessage(ctx, req)
+				if err != nil {
+					ch <- err
+
+					return
+				}
+
+				for _, message := range response.Messages {
+					s.wg.Add(1)
+					s.pool.Schedule(func() {
+						s.handler(message)
+						s.wg.Done()
+					})
+				}
+			}
+		}
+	}()
+
+	return ch
+}
+
+func (s *Subscriber) Unsubscribe() error {
+	close(s.stopCh)
+	s.wg.Wait()
+
+	return nil
+}

pkg/pubsub/sqs/subscriber_test.go

@@ -0,0 +1,92 @@
+package sqs
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/sqs"
+	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
+
+	"github.com/scribd/go-sdk/pkg/pubsub/pool"
+)
+
+type (
+	mockSQSClient struct {
+		msgs []types.Message
+	}
+)
+
+func (m *mockSQSClient) ReceiveMessage(
+	ctx context.Context,
+	params *sqs.ReceiveMessageInput, optFns ...func(*sqs.Options)) (*sqs.ReceiveMessageOutput, error) {
+	return &sqs.ReceiveMessageOutput{
+		Messages: m.msgs,
+	}, nil
+}
+
+func Test_Subscriber_Subscribe(t *testing.T) {
+	t.Run("all subscribers finished", func(t *testing.T) {
+		var nHandlers int64 // atomic
+		var executedTimes int64
+
+		c := make(chan int, 6)
+
+		sub := &Subscriber{
+			pool:   pool.New(2),
+			stopCh: make(chan struct{}),
+			client: &mockSQSClient{
+				msgs: []types.Message{
+					{
+						Body: aws.String("1"),
+					},
+					{
+						Body: aws.String("2"),
+					},
+					{
+						Body: aws.String("3"),
+					},
+					{
+						Body: aws.String("4"),
+					},
+					{
+						Body: aws.String("5"),
+					},
+					{
+						Body: aws.String("6"),
+					},
+				},
+			},
+			handler: func(msg types.Message) {
+				c <- 0
+
+				atomic.AddInt64(&nHandlers, 1)
+				defer atomic.AddInt64(&nHandlers, -1)
+				atomic.AddInt64(&executedTimes, 1)
+
+				time.Sleep(time.Millisecond * 10)
+			},
+		}
+
+		_ = sub.Subscribe(context.Background())
+		// Make sure all goroutines have started.
+		for i := 0; i < cap(c); i++ {
+			<-c
+		}
+
+		err := sub.Unsubscribe()
+		if err != nil {
+			t.Errorf("expected nil, got %v", err)
+		}
+
+		if got := atomic.LoadInt64(&nHandlers); got != 0 {
+			t.Errorf("expected 0, got %d", got)
+		}
+
+		if got := atomic.LoadInt64(&executedTimes); got != 6 {
+			t.Errorf("expected 6, got %d", got)
+		}
+	})
+}