mxpv-podsync/pkg/queue/sender.go

package queue

import (
	"context"
	"encoding/json"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"

	"github.com/mxpv/podsync/pkg/api"
)

var (
	sess  = session.Must(session.NewSession())
	queue = sqs.New(sess)
)

const (
	chanSize           = 1024
	maxElementPerBatch = 10 // SQS Batch limit is 10 items per request
)

type Item struct {
	ID       string       `json:"id"`
	URL      string       `json:"url"`
	Start    int          `json:"start"`
	Count    int          `json:"count"`
	LastID   string       `json:"last_id"`
	LinkType api.LinkType `json:"link_type"`
	Format   string       `json:"format"`
	Quality  string       `json:"quality"`
}

type Sender struct {
	url    *string
	items  chan *Item
	cancel context.CancelFunc
}

func New(ctx context.Context, url string) *Sender {
	ctx, cancel := context.WithCancel(ctx)
	items := make(chan *Item, chanSize)

	sender := &Sender{
		url:    aws.String(url),
		items:  items,
		cancel: cancel,
	}

	go sender.transmit(ctx)

	return sender
}

func (s *Sender) Add(item *Item) {
	s.items <- item
}

func (s *Sender) Close() {
	s.cancel()
	close(s.items)
}

func (s *Sender) transmit(ctx context.Context) error {
	var list = make([]*Item, 0, maxElementPerBatch)

	flush := func(ctx context.Context) {
		if len(list) == 0 {
			return
		}

		if err := s.send(ctx, list); err != nil {
			log.WithError(err).Error("failed to send batch")
		}

		list = make([]*Item, 0, maxElementPerBatch)
	}

	for {
		select {
		case <-time.After(5 * time.Second):
			// Flush list if not filled up entirely within 5 seconds
			flush(ctx)

		case item := <-s.items:
			// Append an item to list and flush if filled up
			list = append(list, item)
			if len(list) == maxElementPerBatch {
				flush(ctx)
			}

		case <-ctx.Done():
			// Exiting, flush leftovers
			flush(context.Background())
			return ctx.Err()
		}
	}
}

func (s *Sender) send(ctx context.Context, list []*Item) error {
	if len(list) == 0 {
		return nil
	}

	log.Debugf("sending a new batch")

	sendInput := &sqs.SendMessageBatchInput{
		QueueUrl: s.url,
	}

	for _, item := range list {

		data, err := json.Marshal(item)
		if err != nil {
			return errors.Wrapf(err, "failed to marshal item %q", item.ID)
		}

		body := string(data)

		sendInput.Entries = append(sendInput.Entries, &sqs.SendMessageBatchRequestEntry{
			Id:          aws.String(item.ID),
			MessageBody: aws.String(body),
		})

		log.Debugf("sending batch: %+v", sendInput)
	}

	_, err := queue.SendMessageBatchWithContext(ctx, sendInput)
	if err != nil {
		return errors.Wrap(err, "failed to send message batch")
	}

	log.Infof("sent %d item(s) to SQS", len(list))
	return nil
}