cache size checks in mongo backend

graph/mongo/iterator.go

@@ -47,13 +47,6 @@ func NewIterator(qs *QuadStore, collection string, d quad.Direction, val graph.V
 
 	constraint := bson.M{d.String(): name}
 
-	size, err := qs.db.C(collection).Find(constraint).Count()
-	if err != nil {
-		// FIXME(kortschak) This should be passed back rather than just logging.
-		glog.Errorln("Trouble getting size for iterator! ", err)
-		return nil
-	}
-
 	return &Iterator{
 		uid:        iterator.NextUID(),
 		name:       name,
@@ -61,29 +54,29 @@ func NewIterator(qs *QuadStore, collection string, d quad.Direction, val graph.V
 		collection: collection,
 		qs:         qs,
 		dir:        d,
-		iter:       qs.db.C(collection).Find(constraint).Iter(),
-		size:       int64(size),
+		iter:       nil,
+		size:       -1,
 		hash:       val.(string),
 		isAll:      false,
 	}
 }
 
-func NewAllIterator(qs *QuadStore, collection string) *Iterator {
-	size, err := qs.db.C(collection).Count()
-	if err != nil {
-		// FIXME(kortschak) This should be passed back rather than just logging.
-		glog.Errorln("Trouble getting size for iterator! ", err)
-		return nil
+func (it *Iterator) makeMongoIterator() *mgo.Iter {
+	if it.isAll {
+		return it.qs.db.C(it.collection).Find(nil).Iter()
 	}
+	return it.qs.db.C(it.collection).Find(it.constraint).Iter()
+}
 
+func NewAllIterator(qs *QuadStore, collection string) *Iterator {
 	return &Iterator{
 		uid:        iterator.NextUID(),
 		qs:         qs,
 		dir:        quad.Any,
 		constraint: nil,
 		collection: collection,
-		iter:       qs.db.C(collection).Find(nil).Iter(),
-		size:       int64(size),
+		iter:       nil,
+		size:       -1,
 		hash:       "",
 		isAll:      true,
 	}
@@ -94,13 +87,16 @@ func (it *Iterator) UID() uint64 {
 }
 
 func (it *Iterator) Reset() {
-	it.iter.Close()
+	it.Close()
 	it.iter = it.qs.db.C(it.collection).Find(it.constraint).Iter()
 
 }
 
 func (it *Iterator) Close() error {
-	return it.iter.Close()
+	if it.iter != nil {
+		return it.iter.Close()
+	}
+	return nil
 }
 
 func (it *Iterator) Tagger() *graph.Tagger {
@@ -134,6 +130,9 @@ func (it *Iterator) Next() bool {
 		Added   []int64 `bson:"Added"`
 		Deleted []int64 `bson:"Deleted"`
 	}
+	if it.iter == nil {
+		it.iter = it.makeMongoIterator()
+	}
 	found := it.iter.Next(&result)
 	if !found {
 		err := it.iter.Err()
@@ -197,6 +196,13 @@ func (it *Iterator) Contains(v graph.Value) bool {
 }
 
 func (it *Iterator) Size() (int64, bool) {
+	if it.size == -1 {
+		var err error
+		it.size, err = it.qs.getSize(it.collection, &it.constraint)
+		if err != nil {
+			it.err = err
+		}
+	}
 	return it.size, true
 }
 

graph/mongo/lru.go

@@ -30,7 +30,7 @@ type cache struct {
 
 type kv struct {
 	key   string
-	value string
+	value interface{}
 }
 
 func newCache(size int) *cache {
@@ -41,7 +41,7 @@ func newCache(size int) *cache {
 	return &lru
 }
 
-func (lru *cache) Put(key string, value string) {
+func (lru *cache) Put(key string, value interface{}) {
 	if _, ok := lru.Get(key); ok {
 		return
 	}
@@ -55,7 +55,7 @@ func (lru *cache) Put(key string, value string) {
 	lru.cache[key] = lru.priority.Front()
 }
 
-func (lru *cache) Get(key string) (string, bool) {
+func (lru *cache) Get(key string) (interface{}, bool) {
 	if element, ok := lru.cache[key]; ok {
 		lru.priority.MoveToFront(element)
 		return element.Value.(kv).value, true

graph/mongo/quadstore.go

@@ -48,6 +48,7 @@ type QuadStore struct {
 	session *mgo.Session
 	db      *mgo.Database
 	ids     *cache
+	sizes   *cache
 }
 
 func createNewMongoGraph(addr string, options graph.Options) error {
@@ -106,6 +107,7 @@ func newQuadStore(addr string, options graph.Options) (graph.QuadStore, error) {
 	qs.db = conn.DB(dbName)
 	qs.session = conn
 	qs.ids = newCache(1 << 16)
+	qs.sizes = newCache(1 << 16)
 	return &qs, nil
 }
 
@@ -313,7 +315,7 @@ func (qs *QuadStore) ValueOf(s string) graph.Value {
 func (qs *QuadStore) NameOf(v graph.Value) string {
 	val, ok := qs.ids.Get(v.(string))
 	if ok {
-		return val
+		return val.(string)
 	}
 	var node MongoNode
 	err := qs.db.C("nodes").FindId(v.(string)).One(&node)
@@ -377,3 +379,28 @@ func (qs *QuadStore) QuadDirection(in graph.Value, d quad.Direction) graph.Value
 func (qs *QuadStore) Type() string {
 	return QuadStoreType
 }
+
+func (qs *QuadStore) getSize(collection string, constraint *bson.M) (int64, error) {
+	var size int
+	var err error
+	bytes, err := bson.Marshal(constraint)
+	if err != nil {
+		glog.Errorf("Couldn't marshal internal constraint")
+		return -1, err
+	}
+	key := collection + string(bytes)
+	if val, ok := qs.sizes.Get(key); ok {
+		return val.(int64), nil
+	}
+	if constraint == nil {
+		size, err = qs.db.C(collection).Count()
+	} else {
+		size, err = qs.db.C(collection).Find(constraint).Count()
+	}
+	if err != nil {
+		glog.Errorln("Trouble getting size for iterator! ", err)
+		return -1, err
+	}
+	qs.sizes.Put(key, int64(size))
+	return int64(size), nil
+}