Fix memory corruption bug in impl/memory

impl/memory/gkvlite_utils.go

Index: impl/memory/gkvlite_utils.go
diff --git a/impl/memory/gkvlite_utils.go b/impl/memory/gkvlite_utils.go
index ed7aa5cd5b0f0cacdfa650769fe26f8e897092d4..adbc45df0d69c3de9ad0c22d3c8b0b9930e1ed3d 100644
--- a/impl/memory/gkvlite_utils.go
+++ b/impl/memory/gkvlite_utils.go
@@ -9,13 +9,21 @@ import (
 	"runtime"
 	"sync"
 
+	"github.com/luci/gae/service/datastore"
 	"github.com/luci/gkvlite"
 )
 
-func gkvCollide(o, n *memCollection, f func(k, ov, nv []byte)) {
+func gkvCollide(o, n memCollection, f func(k, ov, nv []byte)) {
+	if o != nil && !o.IsReadOnly() {
+		panic("old collection is r/w")
+	}
+	if n != nil && !n.IsReadOnly() {
+		panic("new collection is r/w")
+	}
+
 	// TODO(riannucci): reimplement in terms of *iterator.
 	oldItems, newItems := make(chan *gkvlite.Item), make(chan *gkvlite.Item)
-	walker := func(c *memCollection, ch chan<- *gkvlite.Item, wg *sync.WaitGroup) {
+	walker := func(c memCollection, ch chan<- *gkvlite.Item, wg *sync.WaitGroup) {
 		defer close(ch)
 		defer wg.Done()
 		if c != nil {
@@ -66,78 +74,128 @@ func gkvCollide(o, n *memCollection, f func(k, ov, nv []byte)) {
 // This is reasonable for in-memory Store objects, since the only errors that
 // should occur happen with file IO on the underlying file (which of course
 // doesn't exist).
-type memStore gkvlite.Store
+type memStore interface {
+	datastore.TestingSnapshot
 
-func (*memStore) ImATestingSnapshot() {}
+	GetCollection(name string) memCollection
+	GetCollectionNames() []string
+	GetOrCreateCollection(name string) memCollection
+	Snapshot() memStore
 
-func newMemStore() *memStore {
-	ret, err := gkvlite.NewStore(nil)
-	memoryCorruption(err)
-	return (*memStore)(ret)
+	IsReadOnly() bool
+}
+
+// memCollection is a gkvlite.Collection which will panic for anything which
+// might otherwise return an error.
+//
+// This is reasonable for in-memory Store objects, since the only errors that
+// should occur happen with file IO on the underlying file (which of course
+// doesn't exist.
+type memCollection interface {
+	Name() string
+	Delete(k []byte) bool
+	Get(k []byte) []byte
+	GetTotals() (numItems, numBytes uint64)
+	MinItem(withValue bool) *gkvlite.Item
+	Set(k, v []byte)
+	VisitItemsAscend(target []byte, withValue bool, visitor gkvlite.ItemVisitor)
+
+	IsReadOnly() bool
 }
 
-func (ms *memStore) Snapshot() *memStore {
-	ret := (*memStore)((*gkvlite.Store)(ms).Snapshot())
-	runtime.SetFinalizer((*gkvlite.Store)(ret), func(s *gkvlite.Store) {
-		go s.Close()
-	})
+type memStoreImpl struct {
+	s  *gkvlite.Store
+	ro bool
+}
+
+var _ memStore = (*memStoreImpl)(nil)
+
+func (*memStoreImpl) ImATestingSnapshot() {}
+
+func (ms *memStoreImpl) IsReadOnly() bool { return ms.ro }
+
+func newMemStore() memStore {
+	store, err := gkvlite.NewStore(nil)
+	memoryCorruption(err)
+	ret := memStore(&memStoreImpl{store, false})
+	if *logMemCollectionFolder != "" {
+		ret = wrapTracingMemStore(ret)
+	}
 	return ret
 }
 
-func (ms *memStore) MakePrivateCollection(cmp gkvlite.KeyCompare) *memCollection {
-	return (*memCollection)((*gkvlite.Store)(ms).MakePrivateCollection(cmp))
+func (ms *memStoreImpl) Snapshot() memStore {
+	if ms.ro {
+		return ms
+	}
+	ret := ms.s.Snapshot()
+	runtime.SetFinalizer(ret, func(s *gkvlite.Store) { go s.Close() })
+	return &memStoreImpl{ret, true}
 }
 
-func (ms *memStore) GetCollection(name string) *memCollection {
-	return (*memCollection)((*gkvlite.Store)(ms).GetCollection(name))
+func (ms *memStoreImpl) GetCollection(name string) memCollection {
+	coll := ms.s.GetCollection(name)
+	if coll == nil {
+		return nil
+	}
+	return &memCollectionImpl{coll, ms.ro}
 }
 
-func (ms *memStore) SetCollection(name string, cmp gkvlite.KeyCompare) *memCollection {
-	return (*memCollection)((*gkvlite.Store)(ms).SetCollection(name, cmp))
+func (ms *memStoreImpl) GetOrCreateCollection(name string) memCollection {
+	coll := ms.GetCollection(name)
+	if coll == nil {
+		coll = &memCollectionImpl{(ms.s.SetCollection(name, nil)), ms.ro}
+	}
+	return coll
 }
 
-func (ms *memStore) GetCollectionNames() []string {
-	return (*gkvlite.Store)(ms).GetCollectionNames()
+func (ms *memStoreImpl) GetCollectionNames() []string {
+	return ms.s.GetCollectionNames()
 }
 
-// memCollection is a gkvlite.Collection which will panic for anything which
-// might otherwise return an error.
-//
-// This is reasonable for in-memory Store objects, since the only errors that
-// should occur happen with file IO on the underlying file (which of course
-// doesn't exist.
-type memCollection gkvlite.Collection
+type memCollectionImpl struct {
+	c  *gkvlite.Collection
+	ro bool
+}
+
+var _ memCollection = (*memCollectionImpl)(nil)
 
-func (mc *memCollection) Get(k []byte) []byte {
-	ret, err := (*gkvlite.Collection)(mc).Get(k)
+func (mc *memCollectionImpl) Name() string     { return mc.c.Name() }
+func (mc *memCollectionImpl) IsReadOnly() bool { return mc.ro }
+
+func (mc *memCollectionImpl) Get(k []byte) []byte {
+	ret, err := mc.c.Get(k)
 	memoryCorruption(err)
 	return ret
 }
 
-func (mc *memCollection) MinItem(withValue bool) *gkvlite.Item {
-	ret, err := (*gkvlite.Collection)(mc).MinItem(withValue)
+func (mc *memCollectionImpl) MinItem(withValue bool) *gkvlite.Item {
+	ret, err := mc.c.MinItem(withValue)
 	memoryCorruption(err)
 	return ret
 }
 
-func (mc *memCollection) Set(k, v []byte) {
-	err := (*gkvlite.Collection)(mc).Set(k, v)
+func (mc *memCollectionImpl) Set(k, v []byte) {
+	err := mc.c.Set(k, v)
 	memoryCorruption(err)
 }
 
-func (mc *memCollection) Delete(k []byte) bool {
-	ret, err := (*gkvlite.Collection)(mc).Delete(k)
+func (mc *memCollectionImpl) Delete(k []byte) bool {
+	ret, err := mc.c.Delete(k)
 	memoryCorruption(err)
 	return ret
 }
 
-func (mc *memCollection) VisitItemsAscend(target []byte, withValue bool, visitor gkvlite.ItemVisitor) {
-	err := (*gkvlite.Collection)(mc).VisitItemsAscend(target, withValue, visitor)
+func (mc *memCollectionImpl) VisitItemsAscend(target []byte, withValue bool, visitor gkvlite.ItemVisitor) {
+	if !mc.ro {
+		panic("attempting to VisitItemsAscend from r/w memCollection")
+	}
+	err := mc.c.VisitItemsAscend(target, withValue, visitor)
 	memoryCorruption(err)
 }
 
-func (mc *memCollection) GetTotals() (numItems, numBytes uint64) {
-	numItems, numBytes, err := (*gkvlite.Collection)(mc).GetTotals()
+func (mc *memCollectionImpl) GetTotals() (numItems, numBytes uint64) {
+	numItems, numBytes, err := mc.c.GetTotals()
 	memoryCorruption(err)
 	return
 }