Examples with ReorderingMultiFuture com.google.appengine.api.datastore.Batcher.ReorderingMultiFuture used on opensource projects

Example 1 with ReorderingMultiFuture

use of com.google.appengine.api.datastore.Batcher.ReorderingMultiFuture in project appengine-java-standard by GoogleCloudPlatform.

the class AsyncCloudDatastoreV1ServiceImpl method doBatchPut.

@Override
protected Future<List<Key>> doBatchPut(@Nullable final Transaction txn, final List<Entity> entities) {
    if (txn == null) {
        CommitRequest.Builder baseReq = CommitRequest.newBuilder();
        baseReq.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
        final Multimap<Integer, Integer> dedupedIndexMap = HashMultimap.create();
        final List<Entity> dedupedEntities = dedupeByKey(entities, dedupedIndexMap);
        final List<Integer> order = Lists.newArrayListWithCapacity(dedupedEntities.size());
        Iterator<CommitRequest.Builder> batches = putBatcher.getBatches(dedupedEntities, baseReq, baseReq.build().getSerializedSize(), /* group= */
        true, order);
        List<Future<CommitResponse>> futures = putBatcher.makeCalls(batches);
        return new ReorderingMultiFuture<CommitResponse, List<Key>>(futures, order) {

            @Override
            protected List<Key> aggregate(CommitResponse intermediateResult, Iterator<Integer> indexItr, List<Key> result) {
                for (MutationResult mutationResult : intermediateResult.getMutationResultsList()) {
                    int index = indexItr.next();
                    Key key = dedupedEntities.get(index).getKey();
                    if (mutationResult.hasKey()) {
                        // Update the Key object from the original Entity with the allocated id.
                        List<PathElement> pathElements = mutationResult.getKey().getPathList();
                        key.setId(pathElements.get(pathElements.size() - 1).getId());
                    }
                    for (Integer dedupedIndex : dedupedIndexMap.get(index)) {
                        result.set(dedupedIndex, key);
                    }
                }
                return result;
            }

            @Override
            protected List<Key> initResult() {
                // Pre-deduped size.
                int size = entities.size();
                // Elements are set into this list by the reordering batcher to recreate the
                // original key order. The list size must be remain mutable for compatibility
                // with previous SDK releases.
                List<Key> keyList = Lists.newArrayListWithCapacity(size);
                keyList.addAll(Collections.<Key>nCopies(size, null));
                return keyList;
            }
        };
    }
    // Handle transactional put.
    // v1 builds transactions locally: we defer the put until transaction commit. However,
    // the Java API requires us to the return keys corresponding to the entities. Insertions
    // might have incomplete keys, so we need to allocate these IDs.
    TransactionImpl.ensureTxnActive(txn);
    final InternalTransactionCloudDatastoreV1 txnV1 = InternalTransactionCloudDatastoreV1.get(txn);
    ImmutableList.Builder<Key> keyListBuilder = ImmutableList.builder();
    final List<Key> incompleteKeys = Lists.newArrayList();
    final List<com.google.datastore.v1.Entity.Builder> incompleteEntityBldrs = Lists.newArrayList();
    for (Entity entity : entities) {
        Key key = entity.getKey();
        keyListBuilder.add(key);
        if (key.isComplete()) {
            // Mutations on complete keys and incomplete keys can't conflict with each other,
            // so it's safe to add mutations with completed keys immediately (out of order)
            // even though we need to wait for the allocations for incomplete keys.
            txnV1.deferPut(entity);
        } else {
            // Freeze entity.  We will add it to the transaction in a callback from the
            // AllocateId call below.
            com.google.datastore.v1.Entity.Builder entityV1 = com.google.datastore.v1.Entity.newBuilder();
            DataTypeTranslator.addPropertiesToPb(entity.getPropertyMap(), entityV1);
            incompleteEntityBldrs.add(entityV1);
            incompleteKeys.add(key);
        }
    }
    final List<Key> allKeys = keyListBuilder.build();
    if (incompleteKeys.isEmpty()) {
        return new FutureHelper.FakeFuture<List<Key>>(allKeys);
    }
    return registerInTransaction(txn, new FutureWrapper<List<com.google.datastore.v1.Key>, List<Key>>(allocateIds(incompleteKeys)) {

        @Override
        protected List<Key> wrap(List<com.google.datastore.v1.Key> completedKeyPbs) {
            Iterator<com.google.datastore.v1.Entity.Builder> entityPbBldrIt = incompleteEntityBldrs.iterator();
            Iterator<Key> incompleteKeysIt = incompleteKeys.iterator();
            for (com.google.datastore.v1.Key keyV1 : completedKeyPbs) {
                // Set the id field of the original key object.
                updateKey(keyV1, incompleteKeysIt.next());
                // Build the entity proto and add it to the transaction.
                txnV1.deferPut(entityPbBldrIt.next().setKey(keyV1));
            }
            return allKeys;
        }

        @Override
        protected Throwable convertException(Throwable cause) {
            return cause;
        }
    });
}

Example 2 with ReorderingMultiFuture

use of com.google.appengine.api.datastore.Batcher.ReorderingMultiFuture in project appengine-java-standard by GoogleCloudPlatform.

the class AsyncDatastoreServiceImpl method doBatchPut.

@Override
protected Future<List<Key>> doBatchPut(@Nullable Transaction txn, final List<Entity> entities) {
    PutRequest baseReq = new PutRequest();
    if (txn != null) {
        TransactionImpl.ensureTxnActive(txn);
        baseReq.setTransaction(InternalTransactionV3.toProto(txn));
    }
    // Do not group when inside a transaction.
    boolean group = !baseReq.hasTransaction();
    final List<Integer> order = Lists.newArrayListWithCapacity(entities.size());
    Iterator<PutRequest> batches = putBatcher.getBatches(entities, baseReq, baseReq.getSerializedSize(), group, order);
    List<Future<PutResponse>> futures = putBatcher.makeCalls(batches);
    return registerInTransaction(txn, new ReorderingMultiFuture<PutResponse, List<Key>>(futures, order) {

        @Override
        protected List<Key> aggregate(PutResponse intermediateResult, Iterator<Integer> indexItr, List<Key> result) {
            for (Reference reference : intermediateResult.keys()) {
                int index = indexItr.next();
                Key key = entities.get(index).getKey();
                KeyTranslator.updateKey(reference, key);
                result.set(index, key);
            }
            return result;
        }

        @Override
        protected List<Key> initResult() {
            // Create an array pre-populated with null values (twice :-))
            List<Key> result = new ArrayList<Key>(Collections.<Key>nCopies(order.size(), null));
            return result;
        }
    });
}