Fix memory corruption bug in impl/memory

impl/memory/datastore_index.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
         "bytes"
         "fmt"
         "sort"
@@ skipping 25 matching lines @@
                 }
                 ret = append(ret, &ds.IndexDefinition{Kind: kind, SortBy: []ds.IndexColumn{{Property: name}}})
                 ret = append(ret, &ds.IndexDefinition{Kind: kind, SortBy: []ds.IndexColumn{{Property: name, Descending: true}}})
         }
         if serializationDeterministic {
                 sort.Sort(ret)
         }
         return ret
 }
 
-func indexEntriesWithBuiltins(k *ds.Key, pm ds.PropertyMap, complexIdxs []*ds.IndexDefinition) *memStore {
+func indexEntriesWithBuiltins(k *ds.Key, pm ds.PropertyMap, complexIdxs []*ds.IndexDefinition) memStore {
         sip := serialize.PropertyMapPartially(k, pm)
         return indexEntries(sip, k.Namespace(), append(defaultIndexes(k.Kind(), pm), complexIdxs...))
 }
 
 // indexRowGen contains enough information to generate all of the index rows which
 // correspond with a propertyList and a ds.IndexDefinition.
 type indexRowGen struct {
         propVec   []serialize.SerializedPslice
         decending []bool
 }
@@ skipping 68 matching lines @@
                 if pv, ok := sip[sb.Property]; ok {
                         m.buf.propVec = append(m.buf.propVec, pv)
                         m.buf.decending = append(m.buf.decending, sb.Descending)
                 } else {
                         return indexRowGen{}, false
                 }
         }
         return m.buf, true
 }
 
-func indexEntries(sip serialize.SerializedPmap, ns string, idxs []*ds.IndexDefinition) *memStore {
+func indexEntries(sip serialize.SerializedPmap, ns string, idxs []*ds.IndexDefinition) memStore {
         ret := newMemStore()
-        idxColl := ret.SetCollection("idx", nil)
+        idxColl := ret.GetOrCreateCollection("idx")
 
         mtch := matcher{}
         for _, idx := range idxs {
                 idx = idx.Normalize()
                 if irg, ok := mtch.match(idx.GetFullSortOrder(), sip); ok {
                         idxBin := serialize.ToBytes(*idx.PrepForIdxTable())
                         idxColl.Set(idxBin, []byte{})
-                        coll := ret.SetCollection(fmt.Sprintf("idx:%s:%s", ns, idxBin), nil)
+                        coll := ret.GetOrCreateCollection(fmt.Sprintf("idx:%s:%s", ns, idxBin))
                         irg.permute(coll.Set)
                 }
         }
 
         return ret
 }
 
 // walkCompIdxs walks the table of compound indexes in the store. If `endsWith`
 // is provided, this will only walk over compound indexes which match
 // Kind, Ancestor, and whose SortBy has `endsWith.SortBy` as a suffix.
-func walkCompIdxs(store *memStore, endsWith *ds.IndexDefinition, cb func(*ds.IndexDefinition) bool) {
+func walkCompIdxs(store memStore, endsWith *ds.IndexDefinition, cb func(*ds.IndexDefinition) bool) {
         idxColl := store.GetCollection("idx")
         if idxColl == nil {
                 return
         }
         itrDef := iterDefinition{c: idxColl}
 
         if endsWith != nil {
                 full := serialize.ToBytes(*endsWith.Flip())
                 // chop off the null terminating byte
                 itrDef.prefix = full[:len(full)-1]
         }
 
         it := itrDef.mkIter()
         defer it.stop()
         for !it.stopped {
                 it.next(nil, func(i *gkvlite.Item) {
                         if i == nil {
                                 return
                         }
                         qi, err := serialize.ReadIndexDefinition(bytes.NewBuffer(i.Key))
                         memoryCorruption(err)
                         if !cb(qi.Flip()) {
                                 it.stop()
                         }
                 })
         }
 }
 
-func mergeIndexes(ns string, store, oldIdx, newIdx *memStore) {
+func mergeIndexes(ns string, store, oldIdx, newIdx memStore) {
         prefixBuf := []byte("idx:" + ns + ":")
         origPrefixBufLen := len(prefixBuf)
+
+        oldIdx = oldIdx.Snapshot()
+        newIdx = newIdx.Snapshot()
+
         gkvCollide(oldIdx.GetCollection("idx"), newIdx.GetCollection("idx"), func(k, ov, nv []byte) {
                 prefixBuf = append(prefixBuf[:origPrefixBufLen], k...)
                 ks := string(prefixBuf)
 
-                coll := store.GetCollection(ks)
-                if coll == nil {
-                        coll = store.SetCollection(ks, nil)
-                }
+                coll := store.GetOrCreateCollection(ks)
 
                 oldColl := oldIdx.GetCollection(ks)
                 newColl := newIdx.GetCollection(ks)
 
                 switch {
                 case ov == nil && nv != nil: // all additions
                         newColl.VisitItemsAscend(nil, false, func(i *gkvlite.Item) bool {
                                 coll.Set(i.Key, []byte{})
                                 return true
                         })
@@ skipping 11 matching lines @@
                                 }
                         })
                 default:
                         impossible(fmt.Errorf("both values from gkvCollide were nil?"))
                 }
                 // TODO(riannucci): remove entries from idxColl and remove index collections
                 // when there are no index entries for that index any more.
         })
 }
 
-func addIndexes(store *memStore, aid string, compIdx []*ds.IndexDefinition) {
+func addIndexes(store memStore, aid string, compIdx []*ds.IndexDefinition) {
         normalized := make([]*ds.IndexDefinition, len(compIdx))
-        idxColl := store.SetCollection("idx", nil)
+        idxColl := store.GetOrCreateCollection("idx")
         for i, idx := range compIdx {
                 normalized[i] = idx.Normalize()
                 idxColl.Set(serialize.ToBytes(*normalized[i].PrepForIdxTable()), []byte{})
         }
 
         for _, ns := range namespaces(store) {
-                if allEnts := store.GetCollection("ents:" + ns); allEnts != nil {
+                if allEnts := store.Snapshot().GetCollection("ents:" + ns); allEnts != nil {
                         allEnts.VisitItemsAscend(nil, true, func(i *gkvlite.Item) bool {
                                 pm, err := rpm(i.Val)
                                 memoryCorruption(err)
 
                                 prop, err := serialize.ReadProperty(bytes.NewBuffer(i.Key), serialize.WithoutContext, aid, ns)
                                 memoryCorruption(err)
 
                                 k := prop.Value().(*ds.Key)
 
                                 sip := serialize.PropertyMapPartially(k, pm)
 
                                 mergeIndexes(ns, store,
                                         newMemStore(),
                                         indexEntries(sip, ns, normalized))
                                 return true
                         })
                 }
         }
 }
 
-func updateIndexes(store *memStore, key *ds.Key, oldEnt, newEnt ds.PropertyMap) {
+func updateIndexes(store memStore, key *ds.Key, oldEnt, newEnt ds.PropertyMap) {
         // load all current complex query index definitions.
         compIdx := []*ds.IndexDefinition{}
-        walkCompIdxs(store, nil, func(i *ds.IndexDefinition) bool {
+        walkCompIdxs(store.Snapshot(), nil, func(i *ds.IndexDefinition) bool {
                 compIdx = append(compIdx, i)
                 return true
         })
 
         mergeIndexes(key.Namespace(), store,
                 indexEntriesWithBuiltins(key, oldEnt, compIdx),
                 indexEntriesWithBuiltins(key, newEnt, compIdx))
 }