Add Hint and MaxTimeMS support to Count()

.travis.yml

@@ -41,5 +41,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

session.go

@@ -3941,10 +3941,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.
@@ -3966,8 +3968,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

@@ -1187,6 +1187,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)
@@ -4159,11 +4202,11 @@ func (s *S) TestBypassValidation(c *C) {
 
 func (s *S) TestVersionAtLeast(c *C) {
 	tests := [][][]int{
-		{{3,2,1}, {3,2,0}},
-		{{3,2,1}, {3,2}},
-		{{3,2,1}, {2,5,5,5}},
-		{{3,2,1}, {2,5,5}},
-		{{3,2,1}, {2,5}},
+		{{3, 2, 1}, {3, 2, 0}},
+		{{3, 2, 1}, {3, 2}},
+		{{3, 2, 1}, {2, 5, 5, 5}},
+		{{3, 2, 1}, {2, 5, 5}},
+		{{3, 2, 1}, {2, 5}},
 	}
 	for _, pair := range tests {
 		bi := mgo.BuildInfo{VersionArray: pair[0]}