Script to find or create docs in mongo

test.go

package main

import (
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type User struct {
	Id            bson.ObjectId `bson:"_id,omitempty"`
	Created time.Time `bson:"created"`
	Modified time.Time `bson:"modified"`
	AccountId string        `bson:"account_id"`
	RealId        int           `bson:"real_id"`
	Attributes map[string]interface{} `bson:"attributes"`
}

var numWorkers int

func main() {
	flag.IntVar(&numWorkers, "workers", 10, "number of workers to run")
	flag.Parse()

	rand.Seed(time.Now().UTC().UnixNano())

	acctIds := make([]string, 50)
	for x := 0; x < 50; x++ {
		acctIds[x] = fmt.Sprintf("e%v", x)
	}

	sess, _ := mgo.DialWithTimeout(fmt.Sprintf("localhost:27017?maxPoolSize=%d", 4096*3), 60*time.Second)
	sess.SetMode(mgo.Eventual, true)
	sess.SetSafe(&mgo.Safe{
		W: 1,
	})

	stopSignal := make(chan os.Signal)
	signal.Notify(stopSignal, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)

	var wg sync.WaitGroup

	type ui struct {
		AcctId  string
		RealId int
	}
	users := make(chan ui, 100000)

	var stop bool
	for x := 0; x < numWorkers; x++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for u := range users {
				var user User

				query := bson.M{"account_id": u.AcctId, "real_id": u.RealId}

				// this simulate the FindOrCreateUser logic
				c := sess.Copy()
				e := c.DB("users").C("users").Find(query).One(&user)
				if NotFound(e) {
					user = User{
						Id:            bson.NewObjectId(),
						Created: time.Now().UTC(),
						Modified: time.Now().UTC(),
						AccountId: u.AcctId,
						RealId:        u.RealId,
						Attributes: map[string]interface{}{
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
							bson.NewObjectId().Hex(): bson.NewObjectId(),
						},
					}
					e = c.DB("users").C("users").Insert(user)
					if Dupe(e) {
						e = c.DB("users").C("users").Find(query).One(&user)
					}
				}

				if e != nil {
					fmt.Println("no user", e, user)
				}
				c.Close()
			}
		}()
	}

	for x := 0; x < 10; x++ {
		go func() {
			for {
				if stop {
					break
				}

				aid := acctIds[rand.Intn(len(acctIds))]
				rid := rand.Intn(2500000)
				users <- ui{aid, rid}
			}
		}()
	}

	select {
	// signals from the system come in on this channel. anything recieved here
	// gets forwarded to the manager.
	case s := <-stopSignal:
		fmt.Println("stopping", s)
		stop = true
		close(users)
	}

	wg.Wait()
}

func NotFound(err error) bool {
	return err != nil && err == mgo.ErrNotFound
}

func Dupe(err error) bool {
	return err != nil && mgo.IsDup(err)
}