You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
55 lines
1.2 KiB
55 lines
1.2 KiB
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")
|
|
})
|
|
}
|