put.go

package nds

import (
	"reflect"
	"sync"

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/memcache"
)

// putMultiLimit is the App Engine datastore limit for the maximum number
// of entities that can be put by datastore.PutMulti at once.
const putMultiLimit = 500

// PutMulti is a batch version of Put. It works just like datastore.PutMulti
// except it interacts appropriately with NDS's caching strategy. It also
// removes the API limit of 500 entities per request by calling the datastore as
// many times as required to put all the keys. It does this efficiently and
// concurrently.
func PutMulti(c context.Context,
	keys []*datastore.Key, vals interface{}) ([]*datastore.Key, error) {

	if len(keys) == 0 {
		return nil, nil
	}

	v := reflect.ValueOf(vals)
	if err := checkKeysValues(keys, v); err != nil {
		return nil, err
	}

	callCount := (len(keys)-1)/putMultiLimit + 1
	putKeys := make([][]*datastore.Key, callCount)
	errs := make([]error, callCount)

	var wg sync.WaitGroup
	wg.Add(callCount)
	for i := 0; i < callCount; i++ {
		lo := i * putMultiLimit
		hi := (i + 1) * putMultiLimit
		if hi > len(keys) {
			hi = len(keys)
		}

		go func(i int, keys []*datastore.Key, vals reflect.Value) {
			putKeys[i], errs[i] = putMulti(c, keys, vals.Interface())
			wg.Done()
		}(i, keys[lo:hi], v.Slice(lo, hi))
	}
	wg.Wait()

	if isErrorsNil(errs) {
		groupedKeys := make([]*datastore.Key, len(keys))
		for i, k := range putKeys {
			lo := i * putMultiLimit
			hi := (i + 1) * putMultiLimit
			if hi > len(keys) {
				hi = len(keys)
			}
			copy(groupedKeys[lo:hi], k)
		}
		return groupedKeys, nil
	}

	groupedKeys := make([]*datastore.Key, len(keys))
	groupedErrs := make(appengine.MultiError, len(keys))
	for i, err := range errs {
		lo := i * putMultiLimit
		hi := (i + 1) * putMultiLimit
		if hi > len(keys) {
			hi = len(keys)
		}
		if me, ok := err.(appengine.MultiError); ok {
			for j, e := range me {
				if e == nil {
					groupedKeys[lo+j] = putKeys[i][j]
				} else {
					groupedErrs[lo+j] = e
				}
			}
		} else if err != nil {
			for j := lo; j < hi; j++ {
				groupedErrs[j] = err
			}
		}
	}

	return groupedKeys, groupedErrs
}

// Put saves the entity val into the datastore with key. val must be a struct
// pointer; if a struct pointer then any unexported fields of that struct will
// be skipped. If key is an incomplete key, the returned key will be a unique
// key generated by the datastore.
func Put(c context.Context,
	key *datastore.Key, val interface{}) (*datastore.Key, error) {

	keys := []*datastore.Key{key}
	vals := []interface{}{val}
	if err := checkKeysValues(keys, reflect.ValueOf(vals)); err != nil {
		return nil, err
	}

	keys, err := putMulti(c, keys, vals)
	switch e := err.(type) {
	case nil:
		return keys[0], nil
	case appengine.MultiError:
		return nil, e[0]
	default:
		return nil, err
	}
}

// putMulti puts the entities into the datastore and then its local cache.
func putMulti(c context.Context,
	keys []*datastore.Key, vals interface{}) ([]*datastore.Key, error) {

	lockMemcacheKeys := make([]string, 0, len(keys))
	lockMemcacheItems := make([]*memcache.Item, 0, len(keys))
	for _, key := range keys {
		if !key.Incomplete() {
			item := &memcache.Item{
				Key:        createMemcacheKey(key),
				Flags:      lockItem,
				Value:      itemLock(),
				Expiration: memcacheLockTime,
			}
			lockMemcacheItems = append(lockMemcacheItems, item)
			lockMemcacheKeys = append(lockMemcacheKeys, item.Key)
		}
	}

	memcacheCtx, err := memcacheContext(c)
	if err != nil {
		return nil, err
	}

	defer func() {
		if _, ok := transactionFromContext(c); !ok {
			// Remove the locks.
			if err := memcacheDeleteMulti(memcacheCtx,
				lockMemcacheKeys); err != nil {
				log.Warningf(c, "putMulti memcache.DeleteMulti %s", err)
			}
		}
	}()

	if tx, ok := transactionFromContext(c); ok {
		tx.Lock()
		tx.lockMemcacheItems = append(tx.lockMemcacheItems,
			lockMemcacheItems...)
		tx.Unlock()
	} else if err := memcacheSetMulti(memcacheCtx,
		lockMemcacheItems); err != nil {
		return nil, err
	}

	// Save to the datastore.
	return datastorePutMulti(c, keys, vals)
}