Merge pull request globalsign#17 from globalsign/feature/fmpwizard-co…

…unt-maxtime-hint

Support index hints & deadlines for Count

.travis.yml

@@ -46,5 +46,6 @@ script:
     - (cd bson && go test -check.v)
     - go test -check.v -fast
     - (cd txn && go test -check.v)
+    - make stopdb
 
 # vim:sw=4:ts=4:et

README.md

@@ -18,15 +18,17 @@ Further PR's (with tests) are welcome, but please maintain backwards compatibili
 * Fixes timezone handling ([details](https://github.com/go-mgo/mgo/pull/464)) 
 * Improved multi-document transaction performance ([details](https://github.com/globalsign/mgo/pull/10), [more](https://github.com/globalsign/mgo/pull/11), [more](https://github.com/globalsign/mgo/pull/16))
 * Fixes cursor timeouts ([detials](https://jira.mongodb.org/browse/SERVER-24899))
+* Support index hints and timeouts for count queries ([details](https://github.com/globalsign/mgo/pull/17))
 
 ---
 
 ### Thanks to
 * @BenLubar
 * @carter2000
 * @cezarsa
-* @eaglerayp
 * @drichelson
+* @eaglerayp
+* @fmpwizard
 * @jameinel
 * @Reenjii
 * @smoya

session.go

@@ -4086,10 +4086,12 @@ func (iter *Iter) getMoreCmd() *queryOp {
 }
 
 type countCmd struct {
-	Count string
-	Query interface{}
-	Limit int32 ",omitempty"
-	Skip  int32 ",omitempty"
+	Count     string
+	Query     interface{}
+	Limit     int32  ",omitempty"
+	Skip      int32  ",omitempty"
+	Hint      bson.D `bson:"hint,omitempty"`
+	MaxTimeMS int    `bson:"maxTimeMS,omitempty"`
 }
 
 // Count returns the total number of documents in the result set.
@@ -4111,8 +4113,12 @@ func (q *Query) Count() (n int, err error) {
 	if query == nil {
 		query = bson.D{}
 	}
+	// not checking the error because if type assertion fails, we
+	// simply want a Zero bson.D
+	hint, _ := q.op.options.Hint.(bson.D)
 	result := struct{ N int }{}
-	err = session.DB(dbname).Run(countCmd{cname, query, limit, op.skip}, &result)
+	err = session.DB(dbname).Run(countCmd{cname, query, limit, op.skip, hint, op.options.MaxTimeMS}, &result)
+
 	return result.N, err
 }
 

session_test.go

@@ -1275,6 +1275,49 @@ func (s *S) TestCountSkipLimit(c *C) {
 	c.Assert(n, Equals, 4)
 }
 
+func (s *S) TestCountMaxTimeMS(c *C) {
+	if !s.versionAtLeast(2, 6) {
+		c.Skip("SetMaxTime only supported in 2.6+")
+	}
+
+	session, err := mgo.Dial("localhost:40001")
+	c.Assert(err, IsNil)
+	defer session.Close()
+
+	coll := session.DB("mydb").C("mycoll")
+
+	ns := make([]int, 100000)
+	for _, n := range ns {
+		err := coll.Insert(M{"n": n})
+		c.Assert(err, IsNil)
+	}
+	_, err = coll.Find(M{"n": M{"$gt": 1}}).SetMaxTime(1 * time.Millisecond).Count()
+	e := err.(*mgo.QueryError)
+	// We hope this query took longer than 1 ms, which triggers an error code 50
+	c.Assert(e.Code, Equals, 50)
+
+}
+
+func (s *S) TestCountHint(c *C) {
+	if !s.versionAtLeast(2, 6) {
+		c.Skip("Not implemented until mongo 2.5.5 https://jira.mongodb.org/browse/SERVER-2677")
+	}
+
+	session, err := mgo.Dial("localhost:40001")
+	c.Assert(err, IsNil)
+	defer session.Close()
+
+	coll := session.DB("mydb").C("mycoll")
+	err = coll.Insert(M{"n": 1})
+	c.Assert(err, IsNil)
+
+	_, err = coll.Find(M{"n": M{"$gt": 1}}).Hint("does_not_exists").Count()
+	e := err.(*mgo.QueryError)
+	// If Hint wasn't doing anything, then Count would ignore the non existent index hint
+	// and return the normal ount. But we instead get an error code 2: bad hint
+	c.Assert(e.Code, Equals, 2)
+}
+
 func (s *S) TestQueryExplain(c *C) {
 	session, err := mgo.Dial("localhost:40001")
 	c.Assert(err, IsNil)