org.apache.bookkeeper.clients.impl.kv.PByteBufTableRangeImpl Maven / Gradle / Ivy

Go to download
/*
 * 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 org.apache.bookkeeper.clients.impl.kv;

import static org.apache.bookkeeper.clients.impl.kv.KvUtils.toProtoCompare;
import static org.apache.bookkeeper.clients.impl.kv.KvUtils.toProtoRequest;
import com.google.common.collect.Lists;
import com.google.protobuf.UnsafeByteOperations;
import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import org.apache.bookkeeper.api.kv.PTable;
import org.apache.bookkeeper.api.kv.Txn;
import org.apache.bookkeeper.api.kv.impl.result.KeyValueFactory;
import org.apache.bookkeeper.api.kv.impl.result.ResultFactory;
import org.apache.bookkeeper.api.kv.op.CompareOp;
import org.apache.bookkeeper.api.kv.op.Op;
import org.apache.bookkeeper.api.kv.op.OpFactory;
import org.apache.bookkeeper.api.kv.options.DeleteOption;
import org.apache.bookkeeper.api.kv.options.IncrementOption;
import org.apache.bookkeeper.api.kv.options.PutOption;
import org.apache.bookkeeper.api.kv.options.RangeOption;
import org.apache.bookkeeper.api.kv.result.DeleteResult;
import org.apache.bookkeeper.api.kv.result.IncrementResult;
import org.apache.bookkeeper.api.kv.result.PutResult;
import org.apache.bookkeeper.api.kv.result.RangeResult;
import org.apache.bookkeeper.api.kv.result.TxnResult;
import org.apache.bookkeeper.clients.impl.container.StorageContainerChannel;
import org.apache.bookkeeper.common.util.Backoff;
import org.apache.bookkeeper.stream.proto.RangeProperties;
import org.apache.bookkeeper.stream.proto.kv.rpc.RoutingHeader;
import org.apache.bookkeeper.stream.proto.kv.rpc.TxnRequest;

/**
 * A range of a table.
 */
class PByteBufTableRangeImpl implements PTable {
    private static final org.slf4j.Logger log = org.slf4j.LoggerFactory.getLogger(PByteBufTableRangeImpl.class);
    private final long streamId;
    private final RangeProperties rangeProps;
    private final StorageContainerChannel scChannel;
    private final ScheduledExecutorService executor;
    private final OpFactory opFactory;
    private final ResultFactory resultFactory;
    private final KeyValueFactory kvFactory;
    private final Backoff.Policy backoffPolicy;

    PByteBufTableRangeImpl(long streamId, RangeProperties rangeProps, StorageContainerChannel scChannel, ScheduledExecutorService executor, OpFactory opFactory, ResultFactory resultFactory, KeyValueFactory kvFactory, Backoff.Policy backoffPolicy) {
        this.streamId = streamId;
        this.rangeProps = rangeProps;
        this.scChannel = scChannel;
        this.executor = executor;
        this.opFactory = opFactory;
        this.resultFactory = resultFactory;
        this.kvFactory = kvFactory;
        this.backoffPolicy = backoffPolicy;
    }

    private RoutingHeader.Builder newRoutingHeader(ByteBuf pKey) {
        return RoutingHeader.newBuilder().setStreamId(streamId).setRangeId(rangeProps.getRangeId()).setRKey(UnsafeByteOperations.unsafeWrap(pKey.nioBuffer()));
    }

    @Override
    public CompletableFuture> get(ByteBuf pKey, ByteBuf lKey, RangeOption option) {
        pKey.retain();
        lKey.retain();
        if (null != option.endKey()) {
            option.endKey().retain();
        }
        return RangeRequestProcessor.of(KvUtils.newRangeRequest(lKey, option).setHeader(newRoutingHeader(pKey)).build(), response -> KvUtils.newRangeResult(response, resultFactory, kvFactory), scChannel, executor, backoffPolicy).process().whenComplete((value, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
            if (null != option.endKey()) {
                ReferenceCountUtil.release(option.endKey());
            }
        });
    }

    @Override
    public CompletableFuture> put(ByteBuf pKey, ByteBuf lKey, ByteBuf value, PutOption option) {
        pKey.retain();
        lKey.retain();
        value.retain();
        return PutRequestProcessor.of(KvUtils.newPutRequest(lKey, value, option).setHeader(newRoutingHeader(pKey)).build(), response -> KvUtils.newPutResult(response, resultFactory, kvFactory), scChannel, executor, backoffPolicy).process().whenComplete((ignored, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
            ReferenceCountUtil.release(value);
        });
    }

    @Override
    public CompletableFuture> delete(ByteBuf pKey, ByteBuf lKey, DeleteOption option) {
        pKey.retain();
        lKey.retain();
        if (null != option.endKey()) {
            option.endKey().retain();
        }
        return DeleteRequestProcessor.of(KvUtils.newDeleteRequest(lKey, option).setHeader(newRoutingHeader(pKey)).build(), response -> KvUtils.newDeleteResult(response, resultFactory, kvFactory), scChannel, executor, backoffPolicy).process().whenComplete((ignored, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
            if (null != option.endKey()) {
                ReferenceCountUtil.release(option.endKey());
            }
        });
    }

    @Override
    public CompletableFuture> increment(ByteBuf pKey, ByteBuf lKey, long amount, IncrementOption option) {
        pKey.retain();
        lKey.retain();
        return IncrementRequestProcessor.of(KvUtils.newIncrementRequest(lKey, amount, option).setHeader(newRoutingHeader(pKey)).build(), response -> KvUtils.newIncrementResult(response, resultFactory, kvFactory), scChannel, executor, backoffPolicy).process().whenComplete((ignored, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
        });
    }

    @Override
    public Txn txn(ByteBuf pKey) {
        return new TxnImpl(pKey);
    }

    @Override
    public void close() {
        // no-op
    }

    @Override
    public OpFactory opFactory() {
        return opFactory;
    }

    //
    // Txn Implementation
    //
    class TxnImpl implements Txn {
        private final ByteBuf pKey;
        private final TxnRequest.Builder txnBuilder;
        private final List resourcesToRelease;

        TxnImpl(ByteBuf pKey) {
            this.pKey = pKey.retain();
            this.txnBuilder = TxnRequest.newBuilder();
            this.resourcesToRelease = Lists.newArrayList();
        }

        @SuppressWarnings("unchecked")
        @Override
        public Txn If(CompareOp... cmps) {
            for (CompareOp cmp : cmps) {
                txnBuilder.addCompare(toProtoCompare(cmp));
                resourcesToRelease.add(cmp);
            }
            return this;
        }

        @SuppressWarnings("unchecked")
        @Override
        public Txn Then(Op... ops) {
            for (Op op : ops) {
                txnBuilder.addSuccess(toProtoRequest(op));
                resourcesToRelease.add(op);
            }
            return this;
        }

        @SuppressWarnings("unchecked")
        @Override
        public Txn Else(Op... ops) {
            for (Op op : ops) {
                txnBuilder.addFailure(toProtoRequest(op));
                resourcesToRelease.add(op);
            }
            return this;
        }

        @Override
        public CompletableFuture> commit() {
            return TxnRequestProcessor.of(txnBuilder.setHeader(newRoutingHeader(pKey)).build(), response -> KvUtils.newKvTxnResult(response, resultFactory, kvFactory), scChannel, executor, backoffPolicy).process().whenComplete((ignored, cause) -> {
                ReferenceCountUtil.release(pKey);
                for (AutoCloseable resource : resourcesToRelease) {
                    closeResource(resource);
                }
            });
        }

        private void closeResource(AutoCloseable resource) {
            try {
                resource.close();
            } catch (Exception e) {
                log.warn("Fail to close resource {}", resource, e);
            }
        }
    }
}