Fix race in impl/memory.

impl/memory/datastore_data.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"
         "sync"
@@ skipping 189 matching lines @@
         if amt <= 0 {
                 panic(fmt.Errorf("incrementLocked called with bad `amt`: %d", amt))
         }
         ret := curVersion(ents, key) + 1
         ents.Set(key, serialize.ToBytes(ds.PropertyMap{
                 "__version__": {ds.MkPropertyNI(ret + int64(amt-1))},
         }))
         return ret
 }
 
-func (d *dataStoreData) mutableEnts(ns string) *memCollection {
-        d.Lock()
-        defer d.Unlock()
-
+func (d *dataStoreData) mutableEntsLocked(ns string) *memCollection {
         coll := "ents:" + ns
         ents := d.head.GetCollection(coll)
         if ents == nil {
                 ents = d.head.SetCollection(coll, nil)
         }
         return ents
 }
 
 func (d *dataStoreData) allocateIDs(incomplete *ds.Key, n int) (int64, error) {
-        ents := d.mutableEnts(incomplete.Namespace())

[iannucci, 2015/11/06 23:25:03]

apparently getting a collection locked, releasing the lock, then writing to it
with the lock again is not safe according to the race detector. Now we always
re-acquire the collection when we want to write to it.

-
         d.Lock()
         defer d.Unlock()
+
+        ents := d.mutableEntsLocked(incomplete.Namespace())
         return d.allocateIDsLocked(ents, incomplete, n)
 }
 
 func (d *dataStoreData) allocateIDsLocked(ents *memCollection, incomplete *ds.Key, n int) (int64, error) {
         if d.disableSpecialEntities {
                 return 0, errors.New("disableSpecialEntities is true so allocateIDs is disabled")
         }
 
         idKey := []byte(nil)
         if incomplete.Parent() == nil {
@@ skipping 10 matching lines @@
                 if err != nil {
                         return key, err
                 }
                 key = ds.NewKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
         }
         return key, nil
 }
 
 func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
         ns := keys[0].Namespace()
-        ents := d.mutableEnts(ns)
 
         for i, k := range keys {
                 pmap, _ := vals[i].Save(false)
                 dataBytes := serialize.ToBytes(pmap)
 
                 k, err := func() (ret *ds.Key, err error) {
                         d.Lock()
                         defer d.Unlock()
 
+                        ents := d.mutableEntsLocked(ns)
+
                         ret, err = d.fixKeyLocked(ents, k)
                         if err != nil {
                                 return
                         }
                         if !d.disableSpecialEntities {
                                 incrementLocked(ents, groupMetaKey(ret), 1)
                         }
 
                         old := ents.Get(keyBytes(ret))
                         oldPM := ds.PropertyMap(nil)
@@ skipping 38 matching lines @@
 func (d *dataStoreData) getMulti(keys []*ds.Key, cb ds.GetMultiCB) error {
         return getMultiInner(keys, cb, func() (*memCollection, error) {
                 s := d.takeSnapshot()
 
                 return s.GetCollection("ents:" + keys[0].Namespace()), nil
         })
 }
 
 func (d *dataStoreData) delMulti(keys []*ds.Key, cb ds.DeleteMultiCB) {
         ns := keys[0].Namespace()
-        ents := d.mutableEnts(ns)
 
-        if ents != nil {
+        hasEntsInNS := func() bool {
+                d.Lock()
+                defer d.Unlock()
+                return d.mutableEntsLocked(ns) != nil
+        }()
+
+        if hasEntsInNS {
                 for _, k := range keys {
                         err := func() error {
                                 kb := keyBytes(k)
 
                                 d.Lock()
                                 defer d.Unlock()
 
+                                ents := d.mutableEntsLocked(ns)
+
                                 if !d.disableSpecialEntities {
                                         incrementLocked(ents, groupMetaKey(k), 1)
                                 }
                                 if old := ents.Get(kb); old != nil {
                                         oldPM, err := rpm(old)
                                         if err != nil {
                                                 return err
                                         }
                                         updateIndexes(d.head, k, oldPM, nil)
                                         ents.Delete(kb)
@@ skipping 156 matching lines @@
                 td.muts[rk] = []txnMutation{}
         }
         if !getOnly {
                 td.muts[rk] = append(td.muts[rk], txnMutation{key, data})
         }
 
         return nil
 }
 
 func (td *txnDataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
-        ents := td.parent.mutableEnts(keys[0].Namespace())
+        ns := keys[0].Namespace()
 
         for i, k := range keys {
                 err := func() (err error) {
                         td.parent.Lock()
                         defer td.parent.Unlock()
+                        ents := td.parent.mutableEntsLocked(ns)
+
                         k, err = td.parent.fixKeyLocked(ents, k)
                         return
                 }()
                 if err == nil {
                         err = td.writeMutation(false, k, vals[i])
                 }
                 if cb != nil {
                         cb(k, err)
                 }
         }
@@ skipping 23 matching lines @@
 }
 
 func keyBytes(key *ds.Key) []byte {
         return serialize.ToBytes(ds.MkProperty(key))
 }
 
 func rpm(data []byte) (ds.PropertyMap, error) {
         return serialize.ReadPropertyMap(bytes.NewBuffer(data),
                 serialize.WithContext, "", "")
 }