queues-go/app/storage/mongo/collections/tasks.go

package tasks

import (
	"context"
	client "queues-go/app/storage/mongo"
	errors "queues-go/app/utils"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type Task struct {
	Queue            string             `bson:"queue"`
	Payload          bson.M             `bson:"payload"`
	PutAt            time.Time          `bson:"put_at"`
	AvailableFrom    time.Time          `bson:"available_from"`
	SecondsToExecute int                `bson:"seconds_to_execute"`
	Id               primitive.ObjectID `bson:"_id"`
	TakenAt          *time.Time         `bson:"taken_at"`
	Attempts         int                `bson:"attempts"`
}

type InsertedTask struct {
	Queue            string     `bson:"queue"`
	Payload          bson.M     `bson:"payload"`
	PutAt            time.Time  `bson:"put_at"`
	AvailableFrom    time.Time  `bson:"available_from"`
	SecondsToExecute int        `bson:"seconds_to_execute"`
	TakenAt          *time.Time `bson:"taken_at"`
	Attempts         int        `bson:"attempts"`
}

func Add(task InsertedTask) error {
	_, err := collection().InsertOne(context.TODO(), task)
	if err != nil {
		return errors.ErrInternalError
	}
	return nil
}

func Finish(id string) error {
	objectId, err := primitive.ObjectIDFromHex(id)
	if err != nil {
		return errors.ErrIncorrectFormat
	}

	_, err = collection().DeleteOne(context.TODO(), bson.M{"_id": objectId})
	if err != nil {
		return errors.ErrInternalError
	}
	return nil
}

func Take(queue string) (*Task, error) {
	now := time.Now()
	task, err := findTask(queue, now)
	if err != nil {
		return nil, errors.ErrInternalError
	}
	if task == nil {
		return nil, nil
	}
	_, err = collection().UpdateByID(
		context.TODO(),
		task.Id,
		bson.M{
			"$set": bson.M{
				"taken_at": now,
				"attempts": task.Attempts + 1,
				"available_from": now.Add(
					time.Duration(task.SecondsToExecute) * time.Second,
				),
			},
		},
	)
	if err != nil {
		println("Error updaing")
		println(err.Error())
		return nil, errors.ErrInternalError
	}
	return task, nil
}

func findTask(queue string, now time.Time) (*Task, error) {
	cursor, err := collection().Find(
		context.TODO(),
		bson.M{
			"queue":          queue,
			"available_from": bson.M{"$lte": now},
		},
		options.Find().SetLimit(1),
	)
	if err != nil {
		println("Error find")
		println(err.Error())
		return nil, errors.ErrInternalError
	}

	var results []Task
	err = cursor.All(context.TODO(), &results)
	if err != nil {
		println("Error all")
		println(err.Error())
		return nil, errors.ErrInternalError
	}

	for _, task := range results {
		return &task, nil
	}
	return nil, nil
}

func collection() *mongo.Collection {
	return client.Database.Collection("tasks")
}