Add Azure service bus as a broker

README.md

@@ -31,6 +31,7 @@ Machinery is an asynchronous task queue/job queue based on distributed message p
   * [DynamoDB](#dynamodb)
   * [Redis](#redis-2)
   * [GCPPubSub](#gcppubsub)
+  * [ServiceBus](#service-bus)
 * [Custom Logger](#custom-logger)
 * [Server](#server)
 * [Workers](#workers)
@@ -245,6 +246,16 @@ cnf := &config.Config{
 }
 ```
 
+##### Service Bus
+
+Use Service bus connection string in the format:
+
+```
+Endpoint=sb://...
+```
+
+See [Azure Service Bus docs](https://docs.microsoft.com/en-us/azure/service-bus-messaging/) for more information.
+
 #### DefaultQueue
 
 Default queue name, e.g. `machinery_tasks`.

v1/brokers/servicebus/servicebus.go

@@ -0,0 +1,183 @@
+package servicebus
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	servicebus "github.com/Azure/azure-service-bus-go"
+	"github.com/RichardKnop/machinery/v1/brokers/iface"
+	"github.com/RichardKnop/machinery/v1/common"
+	"github.com/RichardKnop/machinery/v1/config"
+	"github.com/RichardKnop/machinery/v1/log"
+	"github.com/RichardKnop/machinery/v1/tasks"
+)
+
+// Broker struct to hold all service bus related stuff
+type Broker struct {
+	common.Broker
+	service      *servicebus.Namespace
+	publishQueue *servicebus.Queue
+	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
+
+	stopReceiving chan struct{}
+}
+
+// New creates a new broker
+func New(cnf *config.Config) (iface.Broker, error) {
+	b := &Broker{Broker: common.NewBroker(cnf), stopReceiving: make(chan struct{})}
+	if cnf.ServiceBus != nil && cnf.ServiceBus.Client != nil {
+		b.service = cnf.ServiceBus.Client
+	} else {
+		ns, err := servicebus.NewNamespace(servicebus.NamespaceWithConnectionString(cnf.Broker))
+		if err != nil {
+			return nil, err
+		}
+		b.service = ns
+	}
+	ctx := context.Background()
+	_, err := b.service.NewQueueManager().Get(ctx, cnf.DefaultQueue)
+	if err != nil {
+		if _, ok := err.(servicebus.ErrNotFound); ok {
+			return nil, fmt.Errorf("queue %s does not exist", cnf.DefaultQueue)
+		}
+		return nil, err
+	}
+	queue, err := b.service.NewQueue(b.GetConfig().DefaultQueue)
+	if err != nil {
+		return nil, err
+	}
+	b.publishQueue = queue
+	return b, nil
+}
+
+// StartConsuming ...
+func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
+	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	queue := b.publishQueue
+	var err error
+	// we need a new queue connection with prefetch count
+	if concurrency > 1 {
+		queue, err = b.service.NewQueue(b.GetConfig().DefaultQueue, servicebus.QueueWithPrefetchCount(uint32(concurrency)))
+		if err != nil {
+			return false, err
+		}
+	}
+
+	// Define msg chan
+	msgChan := make(chan *servicebus.Message, concurrency)
+	// Define a function that should be executed when a message is received.
+	var concurrentHandler servicebus.HandlerFunc = func(ctx context.Context, msg *servicebus.Message) error {
+		msgChan <- msg
+		return nil
+	}
+
+	// Define msg workers
+	for i := 0; i < concurrency; i++ {
+		go func() {
+			for msg := range msgChan {
+				b.processingWG.Add(1)
+				b.consumeOne(context.Background(), msg, taskProcessor)
+				b.processingWG.Done()
+			}
+		}()
+	}
+
+	go func() {
+		<-b.GetStopChan()
+		cancel()
+	}()
+
+	for {
+		err := queue.Receive(ctx, concurrentHandler)
+		if err == nil {
+			break
+		}
+
+		log.ERROR.Printf("Error when receiving messages. Error: %v", err)
+		continue
+	}
+
+	close(b.stopReceiving)
+
+	close(msgChan)
+
+	return b.GetRetry(), nil
+}
+
+// StopConsuming ...
+func (b *Broker) StopConsuming() {
+	b.Broker.StopConsuming()
+
+	<-b.stopReceiving
+
+	// Wait for all processing tasks to finish
+	b.processingWG.Wait()
+
+}
+
+// Publish message to queue
+func (b *Broker) Publish(ctx context.Context, sig *tasks.Signature) error {
+	// Adjust routing key (this decides which queue the message will be published to)
+	b.AdjustRoutingKey(sig)
+	sigMarshaled, err := json.Marshal(sig)
+	if err != nil {
+		return fmt.Errorf("JSON marshal error: %s", err)
+	}
+
+	msg := servicebus.NewMessage(sigMarshaled)
+	// Set message id to machinery task UUID
+	msg.ID = sig.UUID
+	// Check the ETA signature field, if it is set and it is in the future,
+	// delay the task
+	if sig.ETA != nil {
+		now := time.Now().UTC()
+		if sig.ETA.After(now) {
+			msg.ScheduleAt(*sig.ETA)
+		}
+	}
+
+	err = b.publishQueue.Send(ctx, msg)
+	if err != nil {
+		log.ERROR.Printf("Error when sending a message: %v", err)
+		return err
+	}
+	return nil
+}
+
+func (b *Broker) consumeOne(ctx context.Context, msg *servicebus.Message, taskProcessor iface.TaskProcessor) error {
+	if len(msg.Data) == 0 {
+		log.ERROR.Printf("received an empty message, the msg was %v", msg)
+		return msg.DeadLetter(ctx, fmt.Errorf("empty message data"))
+	}
+	sig := new(tasks.Signature)
+	decoder := json.NewDecoder(bytes.NewBuffer(msg.Data))
+	decoder.UseNumber()
+	if err := decoder.Decode(sig); err != nil {
+		log.ERROR.Printf("unmarshal error. the message is %v", msg)
+		return msg.DeadLetter(ctx, fmt.Errorf("unmarshal msg data error"))
+	}
+	// If the task is not registered return an error
+	// and leave the message in the queue
+	if !b.IsTaskRegistered(sig.Name) {
+		log.ERROR.Printf("task %s is not registered", sig.Name)
+		if sig.IgnoreWhenTaskNotRegistered {
+			return msg.DeadLetter(ctx, fmt.Errorf("task %s is not registered", sig.Name))
+		}
+		return msg.Abandon(ctx)
+	}
+
+	err := taskProcessor.Process(sig)
+	if err != nil {
+		log.ERROR.Printf("failed process of task %v", err)
+		return msg.Abandon(ctx)
+	}
+	// Call Complete() after successfully consuming and processing the message
+	return msg.Complete(ctx)
+}

v1/config/config.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/pubsub"
+	servicebus "github.com/Azure/azure-service-bus-go"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/sqs"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -53,17 +54,18 @@ var (
 
 // Config holds all configuration for our program
 type Config struct {
-	Broker                  string           `yaml:"broker" envconfig:"BROKER"`
-	Lock                    string           `yaml:"lock" envconfig:"LOCK"`
-	MultipleBrokerSeparator string           `yaml:"multiple_broker_separator" envconfig:"MULTIPLE_BROKEN_SEPARATOR"`
-	DefaultQueue            string           `yaml:"default_queue" envconfig:"DEFAULT_QUEUE"`
-	ResultBackend           string           `yaml:"result_backend" envconfig:"RESULT_BACKEND"`
-	ResultsExpireIn         int              `yaml:"results_expire_in" envconfig:"RESULTS_EXPIRE_IN"`
-	AMQP                    *AMQPConfig      `yaml:"amqp"`
-	SQS                     *SQSConfig       `yaml:"sqs"`
-	Redis                   *RedisConfig     `yaml:"redis"`
-	GCPPubSub               *GCPPubSubConfig `yaml:"-" ignored:"true"`
-	MongoDB                 *MongoDBConfig   `yaml:"-" ignored:"true"`
+	Broker                  string            `yaml:"broker" envconfig:"BROKER"`
+	Lock                    string            `yaml:"lock" envconfig:"LOCK"`
+	MultipleBrokerSeparator string            `yaml:"multiple_broker_separator" envconfig:"MULTIPLE_BROKEN_SEPARATOR"`
+	DefaultQueue            string            `yaml:"default_queue" envconfig:"DEFAULT_QUEUE"`
+	ResultBackend           string            `yaml:"result_backend" envconfig:"RESULT_BACKEND"`
+	ResultsExpireIn         int               `yaml:"results_expire_in" envconfig:"RESULTS_EXPIRE_IN"`
+	AMQP                    *AMQPConfig       `yaml:"amqp"`
+	SQS                     *SQSConfig        `yaml:"sqs"`
+	Redis                   *RedisConfig      `yaml:"redis"`
+	GCPPubSub               *GCPPubSubConfig  `yaml:"-" ignored:"true"`
+	ServiceBus              *ServiceBusConfig `yaml:"servicebus"`
+	MongoDB                 *MongoDBConfig    `yaml:"-" ignored:"true"`
 	TLSConfig               *tls.Config
 	// NoUnixSignals - when set disables signal handling in machinery
 	NoUnixSignals bool            `yaml:"no_unix_signals" envconfig:"NO_UNIX_SIGNALS"`
@@ -157,6 +159,11 @@ type GCPPubSubConfig struct {
 	MaxExtension time.Duration
 }
 
+// ServiceBusConfig wraps service bus related configuration
+type ServiceBusConfig struct {
+	Client *servicebus.Namespace
+}
+
 // MongoDBConfig ...
 type MongoDBConfig struct {
 	Client   *mongo.Client

v1/factories.go

@@ -3,6 +3,7 @@ package machinery
 import (
 	"errors"
 	"fmt"
+	"github.com/RichardKnop/machinery/v1/brokers/servicebus"
 	"os"
 	"strconv"
 	"strings"
@@ -103,7 +104,9 @@ func BrokerFactory(cnf *config.Config) (brokeriface.Broker, error) {
 		}
 		return gcppubsubbroker.New(cnf, projectID, subscriptionName)
 	}
-
+	if strings.HasPrefix(cnf.Broker, "Endpoint=sb://") {
+		return servicebus.New(cnf)
+	}
 	return nil, fmt.Errorf("Factory failed with broker URL: %v", cnf.Broker)
 }
 

v1/go.mod

@@ -3,10 +3,10 @@ module github.com/RichardKnop/machinery/v1
 require (
 	cloud.google.com/go v0.72.0 // indirect
 	cloud.google.com/go/pubsub v1.8.3
+	github.com/Azure/azure-service-bus-go v0.10.12
 	github.com/RichardKnop/logging v0.0.0-20190827224416-1a693bdd4fae
 	github.com/aws/aws-sdk-go v1.35.35
 	github.com/bradfitz/gomemcache v0.0.0-20190913173617-a41fca850d0b
-	github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
 	github.com/go-redis/redis v6.15.9+incompatible
 	github.com/go-redis/redis/v8 v8.4.0
 	github.com/go-redsync/redsync/v4 v4.0.4
@@ -20,10 +20,8 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0
 	github.com/pkg/errors v0.9.1
 	github.com/robfig/cron/v3 v3.0.1
-	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/streadway/amqp v1.0.0
 	github.com/stretchr/testify v1.6.1
-	github.com/urfave/cli v1.22.5
 	github.com/xdg/stringprep v1.0.0 // indirect
 	go.mongodb.org/mongo-driver v1.4.3
 	golang.org/x/crypto v0.0.0-20201124201722-c8d3bf9c5392 // indirect