com.palantir.common.base.BatchingVisitableFromIterable Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of atlasdb-commons Show documentation
Palantir open source project
There is a newer version: 0.1152.0
/*
 * (c) Copyright 2018 Palantir Technologies Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.palantir.common.base;

import com.google.common.collect.Iterables;
import com.palantir.common.collect.IterableUtils;
import com.palantir.common.collect.IteratorUtils;
import com.palantir.logsafe.Preconditions;
import com.palantir.logsafe.exceptions.SafeIllegalStateException;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;

public class BatchingVisitableFromIterable extends AbstractBatchingVisitable {

    private static final int MAX_BATCH_SIZE = 10000;

    private final Iterable iterable;

    public static  BatchingVisitable create(Iterable iterable) {
        return new BatchingVisitableFromIterable(iterable);
    }

    public BatchingVisitableFromIterable(Iterable iterable) {
        this.iterable = IterableUtils.wrap(iterable);
    }

    /**
     * This creates a one time use visitable.  The only proper use of the returned class is to call
     * accept on it directly and never reference it again.
     */
    public static  BatchingVisitable create(final Iterator iterator) {
        return new SingleCallBatchingVisitable<>(create(() -> IteratorUtils.wrap(iterator)));
    }

    @Override
    public  void batchAcceptSizeHint(int batchSize, ConsistentVisitor v) throws K {
        /*
         * Iterables.partition allocates an array of size batchSize, so avoid an OOM by making sure
         * it's not too big.
         */
        batchSize = Math.min(batchSize, MAX_BATCH_SIZE);
        for (List list : Iterables.partition(iterable, batchSize)) {
            if (!v.visit(list)) {
                return;
            }
        }
    }

    private static final class SingleCallBatchingVisitable implements BatchingVisitable {

        @SuppressWarnings("rawtypes")
        private static final AtomicReferenceFieldUpdater
                delegateUpdater = AtomicReferenceFieldUpdater.newUpdater(
                        SingleCallBatchingVisitable.class, BatchingVisitable.class, "delegate");

        private volatile BatchingVisitable delegate;

        SingleCallBatchingVisitable(BatchingVisitable delegate) {
            this.delegate = Preconditions.checkNotNull(delegate, "Delegate BatchingVisitable is required");
        }

        @Override
        @SuppressWarnings("unchecked")
        public  boolean batchAccept(int batchSize, AbortingVisitor, K> visitor)
                throws K {
            BatchingVisitable current = (BatchingVisitable) delegateUpdater.getAndSet(this, null);
            if (current == null) {
                throw new SafeIllegalStateException("This class has already been called once before");
            }
            return current.batchAccept(batchSize, visitor);
        }

        @Override
        public String toString() {
            return "SingleCallBatchingVisitable{" + delegate + '}';
        }
    }
}