mongo_conn_test.go

package main

import "flag"
import "runtime"
import "sync/atomic"
import "time"

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

func play(session *mgo.Session, iterations int, db string, collection string, numFinished *int32) {
	defer session.Close()

	for i := 0; i < iterations; i++ {
		_, err := session.DB(db).C(collection).Find(bson.D{{"xxx", "xxx"}}).Count()
		if err != nil {
			panic(err)
		}
	}

	atomic.AddInt32(numFinished, 1)
}

func checkCollection(session *mgo.Session, db string, collection string, numDocs int) {
	c := session.DB(db).C(collection)
	n, err := c.Count()
	if err != nil {
		panic(err)
	}
	if n == numDocs {
		return
	}

	c.DropCollection()

	for i := 0; i < numDocs; i++ {
		err = c.Insert(bson.D{{"x", i}})
		if err != nil {
			panic(err)
		}
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	
	host := flag.String("host", "127.0.0.1:27017", "mongo host:port")
	threads := flag.Int("threads", 20, "number of threads to use")

	iterations := flag.Int("n", 1000, "iterations")
	db := flag.String("db", "test", "which db to iterate over")
	collection := flag.String("collection", "test", "which collection to iterate over")

	numDocs := flag.Int("numDocs", 1000, "how many docs in collection")
	
	flag.Parse()
	
	session, err := mgo.Dial(*host)
	if err != nil {
		panic(err)
	}
	defer session.Close()

	checkCollection(session, *db, *collection, *numDocs)
	
	var numFinished int32 = 0
	
	for i := 0; i < *threads; i++ {
		go play(session.New(), *iterations, *db, *collection, &numFinished)
	}

	for numFinished < int32(*threads) {
		time.Sleep(time.Second)
	}
}