github.com/iron-io/functions@v0.0.0-20180820112432-d59d7d1c40b2/api/mqs/ironmq.go (about)

     1  package mqs
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"net/url"
     8  	"strconv"
     9  	"strings"
    10  	"sync"
    11  
    12  	"github.com/Sirupsen/logrus"
    13  	"github.com/iron-io/functions/api/models"
    14  	mq_config "github.com/iron-io/iron_go3/config"
    15  	ironmq "github.com/iron-io/iron_go3/mq"
    16  )
    17  
    18  type assoc struct {
    19  	msgId         string
    20  	reservationId string
    21  }
    22  
    23  type IronMQ struct {
    24  	queues []ironmq.Queue
    25  	// Protects the map
    26  	sync.Mutex
    27  	// job id to {msgid, reservationid}
    28  	msgAssoc map[string]*assoc
    29  }
    30  
    31  type IronMQConfig struct {
    32  	Token       string `mapstructure:"token"`
    33  	ProjectId   string `mapstructure:"project_id"`
    34  	Host        string `mapstructure:"host"`
    35  	Scheme      string `mapstructure:"scheme"`
    36  	Port        uint16 `mapstructure:"port"`
    37  	QueuePrefix string `mapstructure:"queue_prefix"`
    38  }
    39  
    40  func NewIronMQ(url *url.URL) *IronMQ {
    41  
    42  	if url.User == nil || url.User.Username() == "" {
    43  		logrus.Fatal("IronMQ requires PROJECT_ID and TOKEN")
    44  	}
    45  	p, ok := url.User.Password()
    46  	if !ok {
    47  		logrus.Fatal("IronMQ requires PROJECT_ID and TOKEN")
    48  	}
    49  	settings := &mq_config.Settings{
    50  		Token:     p,
    51  		ProjectId: url.User.Username(),
    52  		Host:      url.Host,
    53  		Scheme:    "https",
    54  	}
    55  
    56  	if url.Scheme == "ironmq+http" {
    57  		settings.Scheme = "http"
    58  	}
    59  
    60  	parts := strings.Split(url.Host, ":")
    61  	if len(parts) > 1 {
    62  		settings.Host = parts[0]
    63  		p, err := strconv.Atoi(parts[1])
    64  		if err != nil {
    65  			logrus.WithFields(logrus.Fields{"host_port": url.Host}).Fatal("Invalid host+port combination")
    66  		}
    67  		settings.Port = uint16(p)
    68  	}
    69  
    70  	var queueName string
    71  	if url.Path != "" {
    72  		queueName = url.Path
    73  	} else {
    74  		queueName = "titan"
    75  	}
    76  	mq := &IronMQ{
    77  		queues:   make([]ironmq.Queue, 3),
    78  		msgAssoc: make(map[string]*assoc),
    79  	}
    80  
    81  	// Check we can connect by trying to create one of the queues. Create is
    82  	// idempotent, so this is fine.
    83  	_, err := ironmq.ConfigCreateQueue(ironmq.QueueInfo{Name: fmt.Sprintf("%s_%d", queueName, 0)}, settings)
    84  	if err != nil {
    85  		logrus.WithError(err).Fatal("Could not connect to IronMQ")
    86  	}
    87  
    88  	for i := 0; i < 3; i++ {
    89  		mq.queues[i] = ironmq.ConfigNew(fmt.Sprintf("%s_%d", queueName, i), settings)
    90  	}
    91  
    92  	logrus.WithFields(logrus.Fields{"base_queue": queueName}).Info("IronMQ initialized")
    93  	return mq
    94  }
    95  
    96  func (mq *IronMQ) Push(ctx context.Context, job *models.Task) (*models.Task, error) {
    97  	if job.Priority == nil || *job.Priority < 0 || *job.Priority > 2 {
    98  		return nil, fmt.Errorf("IronMQ Push job %s: Bad priority", job.ID)
    99  	}
   100  
   101  	// Push the work onto the queue.
   102  	buf, err := json.Marshal(job)
   103  	if err != nil {
   104  		return nil, err
   105  	}
   106  	_, err = mq.queues[*job.Priority].PushMessage(ironmq.Message{Body: string(buf), Delay: int64(job.Delay)})
   107  	return job, err
   108  }
   109  
   110  func (mq *IronMQ) Reserve(ctx context.Context) (*models.Task, error) {
   111  	var job models.Task
   112  
   113  	var messages []ironmq.Message
   114  	var err error
   115  	for i := 2; i >= 0; i-- {
   116  		messages, err = mq.queues[i].LongPoll(1, 60, 0 /* wait */, false /* delete */)
   117  		if err != nil {
   118  			// It is OK if the queue does not exist, it will be created when a message is queued.
   119  			if !strings.Contains(err.Error(), "404 Not Found") {
   120  				return nil, err
   121  			}
   122  		}
   123  
   124  		if len(messages) == 0 {
   125  			// Try next priority.
   126  			if i == 0 {
   127  				return nil, nil
   128  			}
   129  			continue
   130  		}
   131  
   132  		// Found a message!
   133  		break
   134  	}
   135  
   136  	message := messages[0]
   137  	if message.Body == "" {
   138  		return nil, nil
   139  	}
   140  
   141  	err = json.Unmarshal([]byte(message.Body), &job)
   142  	if err != nil {
   143  		return nil, err
   144  	}
   145  	mq.Lock()
   146  	mq.msgAssoc[job.ID] = &assoc{message.Id, message.ReservationId}
   147  	mq.Unlock()
   148  	return &job, nil
   149  }
   150  
   151  func (mq *IronMQ) Delete(ctx context.Context, job *models.Task) error {
   152  	if job.Priority == nil || *job.Priority < 0 || *job.Priority > 2 {
   153  		return fmt.Errorf("IronMQ Delete job %s: Bad priority", job.ID)
   154  	}
   155  	mq.Lock()
   156  	assoc, exists := mq.msgAssoc[job.ID]
   157  	delete(mq.msgAssoc, job.ID)
   158  	mq.Unlock()
   159  
   160  	if exists {
   161  		return mq.queues[*job.Priority].DeleteMessage(assoc.msgId, assoc.reservationId)
   162  	}
   163  	return nil
   164  }