@@ -6,33 +6,36 @@ import (
66 "fmt"
77 "reflect"
88 "strings"
9+ "sync"
910 "time"
1011
12+ "go.mongodb.org/mongo-driver/bson"
13+ "go.mongodb.org/mongo-driver/mongo"
14+ "go.mongodb.org/mongo-driver/mongo/options"
15+
1116 "github.com/RichardKnop/machinery/v1/backends/iface"
1217 "github.com/RichardKnop/machinery/v1/common"
1318 "github.com/RichardKnop/machinery/v1/config"
1419 "github.com/RichardKnop/machinery/v1/log"
1520 "github.com/RichardKnop/machinery/v1/tasks"
16- "go.mongodb.org/mongo-driver/bson"
17- "go.mongodb.org/mongo-driver/mongo"
18- "go.mongodb.org/mongo-driver/mongo/options"
1921)
2022
2123// Backend represents a MongoDB result backend
2224type Backend struct {
2325 common.Backend
24- client * mongo.Client
25- tasksCollection * mongo.Collection
26- groupMetasCollection * mongo.Collection
26+ client * mongo.Client
27+ tc * mongo.Collection
28+ gmc * mongo.Collection
29+ once sync.Once
2730}
2831
2932// New creates Backend instance
3033func New (cnf * config.Config ) (iface.Backend , error ) {
31- backend := & Backend {Backend : common .NewBackend (cnf )}
32- err := backend .connect ()
33- if err != nil {
34- return nil , err
34+ backend := & Backend {
35+ Backend : common .NewBackend (cnf ),
36+ once : sync.Once {},
3537 }
38+
3639 return backend , nil
3740}
3841
@@ -43,7 +46,7 @@ func (b *Backend) InitGroup(groupUUID string, taskUUIDs []string) error {
4346 TaskUUIDs : taskUUIDs ,
4447 CreatedAt : time .Now ().UTC (),
4548 }
46- _ , err := b .groupMetasCollection .InsertOne (context .Background (), groupMeta )
49+ _ , err := b .groupMetasCollection () .InsertOne (context .Background (), groupMeta )
4750 return err
4851}
4952
@@ -94,7 +97,7 @@ func (b *Backend) TriggerChord(groupUUID string) (bool, error) {
9497 },
9598 }
9699
97- _ , err := b .groupMetasCollection .UpdateOne (context .Background (), query , change , options .Update ())
100+ _ , err := b .groupMetasCollection () .UpdateOne (context .Background (), query , change , options .Update ())
98101
99102 if err != nil {
100103 if err == mongo .ErrNoDocuments {
@@ -175,7 +178,7 @@ func (b *Backend) SetStateFailure(signature *tasks.Signature, err string) error
175178// GetState returns the latest task state
176179func (b * Backend ) GetState (taskUUID string ) (* tasks.TaskState , error ) {
177180 state := & tasks.TaskState {}
178- err := b .tasksCollection .FindOne (context .Background (), bson.M {"_id" : taskUUID }).Decode (state )
181+ err := b .tasksCollection () .FindOne (context .Background (), bson.M {"_id" : taskUUID }).Decode (state )
179182
180183 if err != nil {
181184 return nil , err
@@ -185,13 +188,13 @@ func (b *Backend) GetState(taskUUID string) (*tasks.TaskState, error) {
185188
186189// PurgeState deletes stored task state
187190func (b * Backend ) PurgeState (taskUUID string ) error {
188- _ , err := b .tasksCollection .DeleteOne (context .Background (), bson.M {"_id" : taskUUID })
191+ _ , err := b .tasksCollection () .DeleteOne (context .Background (), bson.M {"_id" : taskUUID })
189192 return err
190193}
191194
192195// PurgeGroupMeta deletes stored group meta data
193196func (b * Backend ) PurgeGroupMeta (groupUUID string ) error {
194- _ , err := b .groupMetasCollection .DeleteOne (context .Background (), bson.M {"_id" : groupUUID })
197+ _ , err := b .groupMetasCollection () .DeleteOne (context .Background (), bson.M {"_id" : groupUUID })
195198 return err
196199}
197200
@@ -207,15 +210,15 @@ func (b *Backend) lockGroupMeta(groupUUID string) error {
207210 },
208211 }
209212
210- _ , err := b .groupMetasCollection .UpdateOne (context .Background (), query , change , options .Update ().SetUpsert (true ))
213+ _ , err := b .groupMetasCollection () .UpdateOne (context .Background (), query , change , options .Update ().SetUpsert (true ))
211214
212215 return err
213216}
214217
215218// unlockGroupMeta releases lock on groupUUID document
216219func (b * Backend ) unlockGroupMeta (groupUUID string ) error {
217220 update := bson.M {"$set" : bson.M {"lock" : false }}
218- _ , err := b .groupMetasCollection .UpdateOne (context .Background (), bson.M {"_id" : groupUUID }, update , options .Update ())
221+ _ , err := b .groupMetasCollection () .UpdateOne (context .Background (), bson.M {"_id" : groupUUID }, update , options .Update ())
219222 return err
220223}
221224
@@ -224,7 +227,7 @@ func (b *Backend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
224227 groupMeta := & tasks.GroupMeta {}
225228 query := bson.M {"_id" : groupUUID }
226229
227- err := b .groupMetasCollection .FindOne (context .Background (), query ).Decode (groupMeta )
230+ err := b .groupMetasCollection () .FindOne (context .Background (), query ).Decode (groupMeta )
228231 if err != nil {
229232 return nil , err
230233 }
@@ -234,7 +237,7 @@ func (b *Backend) getGroupMeta(groupUUID string) (*tasks.GroupMeta, error) {
234237// getStates returns multiple task states
235238func (b * Backend ) getStates (taskUUIDs ... string ) ([]* tasks.TaskState , error ) {
236239 states := make ([]* tasks.TaskState , 0 , len (taskUUIDs ))
237- cur , err := b .tasksCollection .Find (context .Background (), bson.M {"_id" : bson.M {"$in" : taskUUIDs }})
240+ cur , err := b .tasksCollection () .Find (context .Background (), bson.M {"_id" : bson.M {"$in" : taskUUIDs }})
238241 if err != nil {
239242 return nil , err
240243 }
@@ -256,10 +259,26 @@ func (b *Backend) getStates(taskUUIDs ...string) ([]*tasks.TaskState, error) {
256259// updateState saves current task state
257260func (b * Backend ) updateState (signature * tasks.Signature , update bson.M ) error {
258261 update = bson.M {"$set" : update }
259- _ , err := b .tasksCollection .UpdateOne (context .Background (), bson.M {"_id" : signature .UUID }, update , options .Update ().SetUpsert (true ))
262+ _ , err := b .tasksCollection () .UpdateOne (context .Background (), bson.M {"_id" : signature .UUID }, update , options .Update ().SetUpsert (true ))
260263 return err
261264}
262265
266+ func (b * Backend ) tasksCollection () * mongo.Collection {
267+ b .once .Do (func () {
268+ b .connect ()
269+ })
270+
271+ return b .tc
272+ }
273+
274+ func (b * Backend ) groupMetasCollection () * mongo.Collection {
275+ b .once .Do (func () {
276+ b .connect ()
277+ })
278+
279+ return b .gmc
280+ }
281+
263282// connect creates the underlying mgo connection if it doesn't exist
264283// creates required indexes for our collections
265284func (b * Backend ) connect () error {
@@ -275,8 +294,8 @@ func (b *Backend) connect() error {
275294 database = b .GetConfig ().MongoDB .Database
276295 }
277296
278- b .tasksCollection = b .client .Database (database ).Collection ("tasks" )
279- b .groupMetasCollection = b .client .Database (database ).Collection ("group_metas" )
297+ b .tc = b .client .Database (database ).Collection ("tasks" )
298+ b .gmc = b .client .Database (database ).Collection ("group_metas" )
280299
281300 err = b .createMongoIndexes (database )
282301 if err != nil {
0 commit comments