ajbeach2/sqsworker

package sqsworker

import (
    "context"
    "fmt"
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/awserr"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/sns"
    "github.com/aws/aws-sdk-go/service/sns/snsiface"
    "github.com/aws/aws-sdk-go/service/sqs"
    "github.com/aws/aws-sdk-go/service/sqs/sqsiface"
    "go.uber.org/zap"
    "os"
    "runtime"
    "sync"
)

// DefaultWorkers Number of worker goroutines to spawn, each runs the handler function
const DefaultWorkers = 1

// DefaultMaxNumberOfMessages amount of messages received by each SQS request
const DefaultMaxNumberOfMessages = 10

// DefaultVisibilityTimeout SQS visibility Timeout
const DefaultVisibilityTimeout = 60

// DefaultWaitTimeSeconds Long-polling interval for SQS
const DefaultWaitTimeSeconds = 20

// Handler interface for SQS consumers
type Processor interface {
    Process(context.Context, *sqs.Message, *sns.PublishInput) error
}

// Callback which is passed result from handler on success
type Callback func(*string, error)

// Worker encapsulates the SQS consumer
type Worker struct {
    QueueURL  string
    TopicArn  string
    Queue     sqsiface.SQSAPI
    Topic     snsiface.SNSAPI
    Session   *session.Session
    Consumers int
    Logger    *zap.Logger
    Processor Processor
    Callback  Callback
    Name      string
    done      chan error
}

// WorkerConfig settings for Worker to be passed in NewWorker Contstuctor
type WorkerConfig struct {
    QueueURL string
    TopicArn string
    // If the number of workers is 0, the number of workers defaults to runtime.NumCPU()
    Workers   int
    Processor Processor
    Callback  Callback
    Name      string
    Logger    *zap.Logger
}

func (w *Worker) logError(msg string, err error) {
    if w.Logger != nil {
        w.Logger.Error(err.Error(),
            zap.String("app", w.Name),
            zap.String("msg", msg),
            zap.Error(err),
        )
    }
}

func (w *Worker) logInfo(msg string) {
    if w.Logger != nil {
        w.Logger.Info(msg,
            zap.String("app", w.Name),
        )
    }
}

func (w *Worker) deleteMessage(m *sqs.DeleteMessageInput) error {
    _, err := w.Queue.DeleteMessage(m)
    if err != nil {
        return err
    }
    return nil
}

func (w *Worker) sendMessage(msg *sns.PublishInput) error {
    if w.TopicArn == "" {
        return nil
    }

    if msg.Message == nil {
        return nil
    }

    _, err := w.Topic.Publish(msg)
    return err
}

func (w *Worker) consumer(ctx context.Context, in chan *sqs.Message) {
    var msgString string
    deleteInput := &sqs.DeleteMessageInput{QueueUrl: &w.QueueURL}
    var sendInput *sns.PublishInput
    var err error
    for {
        select {
        case <-ctx.Done():
            return
        case msg := <-in:
            if w.Callback != nil || w.TopicArn != "" {
                sendInput = &sns.PublishInput{TopicArn: &w.TopicArn, Message: &msgString}
            }
            err = w.Processor.Process(ctx, msg, sendInput)
            if err == nil {
                err = w.sendMessage(sendInput)
                if err != nil {
                    w.logError("send message failed!", err)
                }
                deleteInput.ReceiptHandle = msg.ReceiptHandle
                err = w.deleteMessage(deleteInput)
                if err != nil {
                    w.logError("delete message failed!", err)
                }
            } else {
                w.logError("handler failed!", err)
            }

            if w.Callback != nil {
                w.Callback(sendInput.Message, err)
            }
        }
    }
}

func (w *Worker) producer(ctx context.Context, out chan *sqs.Message) {
    params := &sqs.ReceiveMessageInput{
        QueueUrl:            aws.String(w.QueueURL),
        MaxNumberOfMessages: aws.Int64(DefaultMaxNumberOfMessages),
        VisibilityTimeout:   aws.Int64(DefaultVisibilityTimeout),
        WaitTimeSeconds:     aws.Int64(DefaultWaitTimeSeconds),
    }

    for {
        select {
        case <-ctx.Done():
            return
        default:
            req, resp := w.Queue.ReceiveMessageRequest(params)
            err := req.Send()
            if err != nil {
                w.logError("receive messages failed!", err)
            } else {
                messages := resp.Messages
                if len(messages) > 0 {
                    for _, message := range messages {
                        out <- message
                    }
                }
            }
        }
    }
}

// Close function will send a signal to all workers to exit
func (w *Worker) Close() {
    close(w.done)
}

// Run does the main consumer/producer loop
func (w *Worker) Run() {
    ctx, cancel := context.WithCancel(context.Background())
    messages := make(chan *sqs.Message, w.Consumers)

    w.logInfo(fmt.Sprint("Staring producer"))
    go func() {
        w.producer(ctx, messages)
        close(messages)
    }()

    go func() {
        <-w.done
        cancel()
    }()

    w.logInfo(fmt.Sprint("Staring consumer with ", w.Consumers, " consumers"))
    // Consume messages
    var wg sync.WaitGroup
    for x := 0; x < w.Consumers; x++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            w.consumer(ctx, messages)
        }()
    }
    wg.Wait()
}

// CreateQueue Create queue by name.
func CreateQueue(name string, sqsc sqsiface.SQSAPI) (string, error) {
    result, err := sqsc.CreateQueue(&sqs.CreateQueueInput{
        QueueName: aws.String(name),
    })
    if err != nil {
        return "", err
    }

    return *result.QueueUrl, nil
}

// GetOrCreateQueue an SQS Queue by name.
func GetOrCreateQueue(name string, sqsc sqsiface.SQSAPI) (string, error) {
    queueOut, err := sqsc.GetQueueUrl(&sqs.GetQueueUrlInput{
        QueueName: aws.String(name),
    })

    if aerr, ok := err.(awserr.Error); ok {
        switch aerr.Code() {
        case sqs.ErrCodeQueueDoesNotExist:
            return CreateQueue(name, sqsc)
        }
    }

    return *queueOut.QueueUrl, err
}

// GetOrCreateTopic Create SNS topic by name.
func GetOrCreateTopic(name string, snsc snsiface.SNSAPI) (string, error) {
    if name == "" {
        return "", nil
    }

    snsOut, err := snsc.CreateTopic(&sns.CreateTopicInput{
        Name: aws.String(name),
    })

    return *snsOut.TopicArn, err
}

// NewWorker constructor for SQS Worker
func NewWorker(sess *session.Session, wc WorkerConfig) *Worker {
    var logger *zap.Logger
    workers := runtime.NumCPU()
    var queueURL, topicARN = wc.QueueURL, wc.TopicArn

    if wc.Workers != 0 {
        workers = wc.Workers
    }

    if wc.Logger == nil {
        logger, _ = zap.NewProduction()
    } else {
        logger = wc.Logger
    }

    if queueURL == "" {
        queueURL = os.Getenv("QUEUE_URL")
    }

    if topicARN == "" {
        topicARN = os.Getenv("TOPIC_ARN")
    }

    return &Worker{
        queueURL,
        topicARN,
        sqs.New(sess),
        sns.New(sess),
        sess,
        workers,
        logger,
        wc.Processor,
        wc.Callback,
        wc.Name,
        make(chan error),
    }
}