Merge pull request #154 from hpidcock/fix-perf

#154

collstats appears to be quite a slow way to determine if the collection is capped.

Author: jujubot

go.mod

@@ -4,6 +4,7 @@ go 1.14
 
 require (
 	github.com/juju/clock v0.0.0-20190205081909-9c5c9712527c
+	github.com/juju/collections v0.0.0-20200605021417-0d0ec82b7271
 	github.com/juju/errors v0.0.0-20200330140219-3fe23663418f
 	github.com/juju/loggo v0.0.0-20200526014432-9ce3a2e09b5e
 	github.com/juju/retry v0.0.0-20180821225755-9058e192b216

go.sum

@@ -3,6 +3,7 @@ github.com/juju/ansiterm v0.0.0-20160907234532-b99631de12cf/go.mod h1:UJSiEoRfvx
 github.com/juju/clock v0.0.0-20190205081909-9c5c9712527c h1:3UvYABOQRhJAApj9MdCN+Ydv841ETSoy6xLzdmmr/9A=
 github.com/juju/clock v0.0.0-20190205081909-9c5c9712527c/go.mod h1:nD0vlnrUjcjJhqN5WuCWZyzfd5AHZAC9/ajvbSx69xA=
 github.com/juju/cmd v0.0.0-20171107070456-e74f39857ca0/go.mod h1:yWJQHl73rdSX4DHVKGqkAip+huBslxRwS8m9CrOLq18=
+github.com/juju/collections v0.0.0-20200605021417-0d0ec82b7271 h1:4R626WTwa7pRYQFiIRLVPepMhm05eZMEx+wIurRnMLc=
 github.com/juju/collections v0.0.0-20200605021417-0d0ec82b7271/go.mod h1:5XgO71dV1JClcOJE+4dzdn4HrI5LiyKd7PlVG6eZYhY=
 github.com/juju/errors v0.0.0-20150916125642-1b5e39b83d18 h1:Sem5Flzxj8ZdAgY2wfHBUlOYyP4wrpIfM8IZgANNGh8=
 github.com/juju/errors v0.0.0-20150916125642-1b5e39b83d18/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=

mgo.go

@@ -27,6 +27,7 @@ import (
 	"time"
 
 	"github.com/juju/clock"
+	"github.com/juju/collections/set"
 	"github.com/juju/errors"
 	"github.com/juju/loggo"
 	"github.com/juju/retry"
@@ -750,6 +751,10 @@ func clearDatabases(session *mgo.Session) error {
 }
 
 func clearCollections(db *mgo.Database) error {
+	capped, err := listCappedCollections(db)
+	if err != nil {
+		return errors.Annotatef(err, "getting capped collection list")
+	}
 	collectionNames, err := db.CollectionNames()
 	if err != nil {
 		return errors.Trace(err)
@@ -760,11 +765,7 @@ func clearCollections(db *mgo.Database) error {
 		}
 		collection := db.C(name)
 		clearFunc := clearNormalCollection
-		capped, err := isCapped(collection)
-		if err != nil {
-			return errors.Trace(err)
-		}
-		if capped {
+		if capped.Contains(name) {
 			clearFunc = clearCappedCollection
 		}
 		err = clearFunc(collection)
@@ -775,21 +776,48 @@ func clearCollections(db *mgo.Database) error {
 	return nil
 }
 
-func isCapped(collection *mgo.Collection) (bool, error) {
-	result := bson.M{}
-	err := collection.Database.Run(bson.D{{"collstats", collection.Name}}, &result)
+func listCappedCollections(db *mgo.Database) (set.Strings, error) {
+	// Mostly pulled from mgo.DB.ListCollections
+	names := set.NewStrings()
+	var result struct {
+		Collections []bson.Raw
+		Cursor      struct {
+			FirstBatch []bson.Raw `bson:"firstBatch"`
+			NextBatch  []bson.Raw `bson:"nextBatch"`
+			NS         string     `bson:"ns"`
+			ID         int64      `bson:"id"`
+		}
+	}
+	err := db.Run(bson.D{{"listCollections", 1}, {"cursor", bson.D{{"batchSize", 10}}}}, &result)
 	if err != nil {
-		return false, errors.Trace(err)
+		return nil, errors.Trace(err)
 	}
-	value, found := result["capped"]
-	if !found {
-		return false, nil
+	firstBatch := result.Collections
+	if firstBatch == nil {
+		firstBatch = result.Cursor.FirstBatch
 	}
-	capped, ok := value.(bool)
-	if !ok {
-		return false, errors.Errorf("unexpected type for capped: %v", value)
+	var iter *mgo.Iter
+	ns := strings.SplitN(result.Cursor.NS, ".", 2)
+	if len(ns) < 2 {
+		iter = db.C("").NewIter(nil, firstBatch, result.Cursor.ID, nil)
+	} else {
+		iter = db.Session.DB(ns[0]).C(ns[1]).NewIter(nil, firstBatch, result.Cursor.ID, nil)
+	}
+	var coll struct {
+		Name    string `bson:"name"`
+		Options struct {
+			Capped bool `bson:"capped"`
+		} `bson:"options"`
+	}
+	for iter.Next(&coll) {
+		if coll.Options.Capped {
+			names.Add(coll.Name)
+		}
+	}
+	if err := iter.Close(); err != nil {
+		return nil, errors.Trace(err)
 	}
-	return capped, nil
+	return names, nil
 }
 
 func clearNormalCollection(collection *mgo.Collection) error {