api/internal/counter/counter.go

package counter

import (
	"git.aiterp.net/rpdata/api/internal/store"
	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
)

var collection *mgo.Collection

type counter struct {
	ID    string `bson:"_id"`
	Value int    `bson:"value"`
}

// Next gets the next value of a counter, or an error if it hasn't.
func Next(category, name string) (int, error) {
	id := category + "." + name
	doc := counter{}

	_, err := collection.Find(bson.M{"_id": id}).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"value": 1}},
		Upsert:    true,
		ReturnNew: true,
	}, &doc)
	if err != nil {
		return -1, err
	}

	return doc.Value, nil
}

// NextMany gets the next value of a counter, or an error if it hasn't, and increments by a specified value.
// Any value `returned` to `returned+(increment-1)` should then be safe to use.
func NextMany(category, name string, increment int) (int, error) {
	id := category + "." + name
	doc := counter{}

	_, err := collection.Find(bson.M{"_id": id}).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"value": increment}},
		Upsert:    true,
		ReturnNew: true,
	}, &doc)
	if err != nil {
		return -1, err
	}

	return doc.Value, nil
}

func init() {
	store.HandleInit(func(db *mgo.Database) {
		collection = db.C("core.counters")
	})
}