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") }) }