resolve race-condition when using Subscribe

Author: mbretter

.gitignore

@@ -19,5 +19,7 @@ start.txt
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
 
+cmd/demo/queue
+
 # Dependency directories (remove the comment below to include it)
 # vendor/

README.md

@@ -13,6 +13,8 @@ polling too.
 The motivation was to build an easy-to-integrate queuing system without sophisticated features, without external 
 dependencies, and with direct integration into your application.
 
+:warning: There is no stable release yet, API is changing 
+
 ## Install
 
 ```
@@ -34,11 +36,9 @@ Along the task, any arbitrary data can be stored.
 
 Each task belongs to a topic, when publishing to a topic, the handler of this topic gets the first unprocessed task.
 
-When using the event based `Subscribe` function, only one handler/consumer for one topic should be used at the same time, 
-for avoiding race conditions.
+You can use either the event based `Subscribe` function, or the `GetNext` function which is needed for polling.
+It both cases it is totally safe to have multiple consumers running on the same topic.
 
-If you need to run multiple handlers/consumers on the same topic, you have to use the `GetNext` function, which supports 
-polling only.
 
 ```go
 ctx := context.TODO()
@@ -77,16 +77,17 @@ payload := Payload{
     Num:  73,
 }
 
-task, err := qu.Publish("some.topic", &payload, queue.DefaultMaxTries)
+task, err := qu.Publish("some.topic", &payload)
 if err != nil {
     log.Fatal(err)
 }
 ```
 
 ## Subscribe
 
-Any handler/application can subscribe to a certain topic. However, there can only be one handler for a certain topic 
-otherwise, you will run into race conditions.
+Any handler/application can subscribe to a certain topic, the callback function receives a copy of the task.
+
+After processing the task you have to `Ack` it, or mark it as `Err`.
 
 Here is a small snippet which demonstrates the usage of subscribe using goroutines.
 ```go
@@ -153,12 +154,23 @@ for {
 }
 ```
 
+## Reschedule
+
+If a task had an error, and you want to process this task again, you can use `Reschedule`.
+
+```go
+newTask, err := Reschedule(task)
+```
+
+When rescheduling, the original task remains untouched, there will create a new task with the same payload and the 
+initial tries value will be increased. 
+
 ## Selfcare
 
 The selfcare function re-schedules long-running tasks, this might happen, if the application could not acknowledge 
 the task, and it sets the task to the error state, if the maximum number of tries have been exceeded.
 
-The selfcare function might be run per topic, if no topic was given, the selfcare runs over all tasks.
+The selfcare function might be run per topic, if no topic was given, the selfcare runs over all topics.
 As second argument, the timeout for long-running tasks can be given, if no timeout was given it defaults to 5 minutes.
 
 ## CreateIndexes

cmd/demo/Makefile

@@ -0,0 +1,4 @@
+.PHONY: build
+
+build:
+	CGO_ENABLED=0 go build -v -o queue main.go

cmd/demo/main.go

@@ -8,6 +8,7 @@ import (
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"log"
+	"os"
 	"sync"
 )
 
@@ -18,8 +19,6 @@ type Payload struct {
 }
 
 func main() {
-	var mongodbUri = flag.String("u", "", "mongodb url")
-	var dbName = flag.String("d", "", "mongodb database name")
 	var collName = flag.String("c", "queue", "mongodb collection name")
 	var publish = flag.String("p", "", "publish topic")
 	var getnext = flag.String("g", "", "next topic")
@@ -29,23 +28,25 @@ func main() {
 	var subscribe = flag.String("s", "", "subscribe on topic")
 	flag.Parse()
 
-	if len(*mongodbUri) == 0 {
+	mongodbUri := os.Getenv("MONGODB_URI")
+	dbName := os.Getenv("MONGODB_DB")
+	if len(mongodbUri) == 0 {
 		log.Fatal("mongodb uri missing")
 	}
 
-	if len(*dbName) == 0 {
+	if len(dbName) == 0 {
 		log.Fatal("mongodb database name missing")
 	}
 
 	ctx := context.TODO()
-	client, err := mongo.Connect(ctx, options.Client().ApplyURI(*mongodbUri))
+	client, err := mongo.Connect(ctx, options.Client().ApplyURI(mongodbUri))
 	if err != nil {
 		log.Fatal(err)
 	}
 	//goland:noinspection ALL
 	defer client.Disconnect(ctx)
 
-	collection := client.Database(*dbName).Collection(*collName)
+	collection := client.Database(dbName).Collection(*collName)
 
 	queueDb := queue.NewStdDb(collection, ctx)
 	qu := queue.NewQueue(queueDb)
@@ -80,7 +81,8 @@ func main() {
 	}
 
 	if *publish != "" {
-		task, err := qu.Publish(*publish, &payload, queue.DefaultMaxTries)
+		opts := queue.NewPublishOptions().SetMaxTries(1)
+		task, err := qu.Publish(*publish, &payload, opts)
 		if err != nil {
 			log.Fatal(err)
 		}

queue.go

@@ -61,18 +61,55 @@ func setNowFunc(n func() time.Time) {
 	nowFunc = n
 }
 
+type PublishOptions struct {
+	MaxTries uint
+	Tries    int
+}
+
+func NewPublishOptions() *PublishOptions {
+	return &PublishOptions{
+		MaxTries: 0,
+		Tries:    -1,
+	}
+}
+
+func (p *PublishOptions) SetMaxTries(maxTries uint) *PublishOptions {
+	p.MaxTries = maxTries
+	return p
+}
+
+func (p *PublishOptions) setTries(tries uint) *PublishOptions {
+	p.Tries = int(tries)
+	return p
+}
+
 // Publish inserts a new task into the queue with the given topic, payload, and maxTries.
 // If maxTries is zero, it defaults to DefaultMaxTries.
-func (q *Queue) Publish(topic string, payload any, maxTries uint) (*Task, error) {
-	if maxTries == 0 {
-		maxTries = DefaultMaxTries
+func (q *Queue) Publish(topic string, payload any, opts ...*PublishOptions) (*Task, error) {
+
+	o := PublishOptions{
+		MaxTries: DefaultMaxTries,
+		Tries:    0,
+	}
+
+	for _, opt := range opts {
+		if opt == nil {
+			continue
+		}
+		if opt.MaxTries > 0 {
+			o.MaxTries = opt.MaxTries
+		}
+
+		if opt.Tries >= 0 {
+			o.Tries = opt.Tries
+		}
 	}
 
 	t := Task{
 		Topic:    topic,
 		Payload:  payload,
-		Tries:    0,
-		MaxTries: maxTries,
+		Tries:    uint(o.Tries),
+		MaxTries: o.MaxTries,
 		Meta: Meta{
 			Created:    nowFunc(),
 			Dispatched: nil,
@@ -102,7 +139,7 @@ func (q *Queue) GetNext(topic string) (*Task, error) {
 			"$set": bson.M{"state": StateRunning, "meta.dispatched": nowFunc()},
 			"$inc": bson.M{"tries": 1},
 		},
-		options.FindOneAndUpdate().SetSort(bson.D{{"meta.scheduled", 1}}),
+		options.FindOneAndUpdate().SetSort(bson.D{{"meta.scheduled", 1}}).SetReturnDocument(options.After),
 	)
 
 	if errors.Is(res.Err(), mongo.ErrNoDocuments) {
@@ -116,11 +153,42 @@ func (q *Queue) GetNext(topic string) (*Task, error) {
 	return &t, nil
 }
 
+func (q *Queue) GetNextById(id primitive.ObjectID) (*Task, error) {
+	t := Task{}
+	res := q.db.FindOneAndUpdate(bson.M{
+		"_id":   id,
+		"state": StatePending,
+		"$expr": bson.M{"$lt": bson.A{"$tries", "$maxtries"}},
+	},
+		bson.M{
+			"$set": bson.M{"state": StateRunning, "meta.dispatched": nowFunc()},
+			"$inc": bson.M{"tries": 1},
+		},
+		options.FindOneAndUpdate().SetReturnDocument(options.After),
+	)
+
+	if errors.Is(res.Err(), mongo.ErrNoDocuments) {
+		return nil, nil
+	}
+
+	if err := res.Decode(&t); err != nil {
+		return nil, err
+	}
+
+	return &t, nil
+}
+
+func (q *Queue) Reschedule(task *Task) (*Task, error) {
+	return q.Publish(task.Topic, task.Payload, NewPublishOptions().setTries(task.Tries).SetMaxTries(task.MaxTries))
+}
+
 type Callback func(t Task)
 
 func (q *Queue) Subscribe(topic string, cb Callback) error {
-	pipeline := bson.D{
-		{"$match", bson.D{{"operationType", "insert"}, {"fullDocument.topic", topic}, {"fullDocument.state", StatePending}}},
+	pipeline := bson.D{{"$match", bson.D{
+		{"operationType", "insert"},
+		{"fullDocument.topic", topic},
+		{"fullDocument.state", StatePending}}},
 	}
 
 	stream, err := q.db.Watch(mongo.Pipeline{pipeline})
@@ -153,30 +221,20 @@ func (q *Queue) Subscribe(topic string, cb Callback) error {
 			continue
 		}
 
-		task := evt.Task
-
 		// already processed
-		if task.Meta.Created.Before(processedUntil) {
+		if evt.Task.Meta.Created.Before(processedUntil) {
 			continue
 		}
 
-		task.State = StateRunning
-		now := nowFunc()
-		task.Meta.Dispatched = &now
-
-		err := q.db.UpdateOne(
-			bson.M{"_id": task.Id},
-			bson.M{"$set": bson.M{
-				"state":           task.State,
-				"meta.dispatched": task.Meta.Dispatched,
-			}})
-
+		task, err := q.GetNextById(evt.Task.Id)
 		if err != nil {
-			_ = q.Err(task.Id.Hex(), err)
+			_ = q.Err(evt.Task.Id.Hex(), err)
 			continue
 		}
 
-		cb(task)
+		if task != nil {
+			cb(*task)
+		}
 	}
 
 	return nil

queue_test.go

@@ -58,7 +58,8 @@ func TestQueue_Publish(t *testing.T) {
 			}
 			dbMock.EXPECT().InsertOne(taskExpected).Return(oId, tt.error)
 
-			task, err := q.Publish(tt.topic, tt.payload, tt.maxTries)
+			opts := NewPublishOptions().SetMaxTries(tt.maxTries)
+			task, err := q.Publish(tt.topic, tt.payload, opts)
 
 			if tt.error == nil {
 				taskExpected.Id = oId
@@ -183,7 +184,7 @@ func TestQueue_Subscribe(t *testing.T) {
 			}, bson.M{
 				"$set": bson.M{"state": StateRunning, "meta.dispatched": now},
 				"$inc": bson.M{"tries": 1},
-			}, options.FindOneAndUpdate().SetSort(bson.D{{"meta.scheduled", 1}})).Return(res)
+			}, options.FindOneAndUpdate().SetSort(bson.D{{"meta.scheduled", 1}}).SetReturnDocument(options.After)).Return(res)
 
 			if tt.task != nil {
 				changeStream.EXPECT().Next(context.TODO()).Once().Return(true)
@@ -202,11 +203,18 @@ func TestQueue_Subscribe(t *testing.T) {
 					goto runTest
 				}
 
-				dbMock.EXPECT().UpdateOne(bson.M{"_id": tt.task.Id},
-					bson.M{"$set": bson.M{
-						"state":           StateRunning,
-						"meta.dispatched": &now,
-					}}).Return(tt.updateError)
+				retTask := *tt.task
+				retTask.State = StateRunning
+				res = mongo.NewSingleResultFromDocument(retTask, tt.updateError, nil)
+
+				dbMock.EXPECT().FindOneAndUpdate(bson.M{
+					"_id":   tt.task.Id,
+					"state": StatePending,
+					"$expr": bson.M{"$lt": bson.A{"$tries", "$maxtries"}},
+				}, bson.M{
+					"$set": bson.M{"state": StateRunning, "meta.dispatched": now},
+					"$inc": bson.M{"tries": 1},
+				}, options.FindOneAndUpdate().SetReturnDocument(options.After)).Return(res)
 
 				if tt.updateError != nil {
 					dbMock.EXPECT().UpdateOne(
@@ -300,7 +308,7 @@ func TestQueue_SubscribeUnprocessedTasks(t *testing.T) {
 				"$inc": bson.M{"tries": 1},
 			}
 
-			opts := options.FindOneAndUpdate().SetSort(bson.D{{"meta.scheduled", 1}})
+			opts := options.FindOneAndUpdate().SetSort(bson.D{{"meta.scheduled", 1}}).SetReturnDocument(options.After)
 			dbMock.EXPECT().FindOneAndUpdate(filter, update, opts).Once().Return(res)
 
 			if tt.error == nil {