org.apache.bookkeeper.stream.storage.impl.kv.TableStoreUtils.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.bookkeeper.stream.storage.impl.kv.TableStoreUtils.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.stream.storage.impl.kv;

import com.google.common.collect.Lists;
import com.google.protobuf.ByteString;
import com.google.protobuf.UnsafeByteOperations;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.buffer.Unpooled;
import java.util.concurrent.ExecutionException;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.api.kv.op.CompareOp;
import org.apache.bookkeeper.api.kv.op.CompareResult;
import org.apache.bookkeeper.api.kv.op.OpFactory;
import org.apache.bookkeeper.api.kv.result.Code;
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.Result;
import org.apache.bookkeeper.api.kv.result.TxnResult;
import org.apache.bookkeeper.statelib.api.exceptions.MVCCStoreException;
import org.apache.bookkeeper.stream.proto.kv.KeyValue;
import org.apache.bookkeeper.stream.proto.kv.rpc.Compare;
import org.apache.bookkeeper.stream.proto.kv.rpc.DeleteRangeResponse;
import org.apache.bookkeeper.stream.proto.kv.rpc.IncrementResponse;
import org.apache.bookkeeper.stream.proto.kv.rpc.PutResponse;
import org.apache.bookkeeper.stream.proto.kv.rpc.RangeResponse;
import org.apache.bookkeeper.stream.proto.kv.rpc.ResponseHeader;
import org.apache.bookkeeper.stream.proto.kv.rpc.ResponseOp;
import org.apache.bookkeeper.stream.proto.kv.rpc.RoutingHeader;
import org.apache.bookkeeper.stream.proto.kv.rpc.TxnResponse;
import org.apache.bookkeeper.stream.proto.storage.StatusCode;

/**
 * Utils for accessing table stores.
 */
@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
final class TableStoreUtils {

    static final byte SEP = 0x0;
    static final byte NO_ROUTING_KEY = 0x0;
    static final byte HAS_ROUTING_KEY = 0x1;

    static boolean hasRKey(ByteString rKey) {
        return null != rKey && rKey.size() > 0;
    }

    static byte[] newStoreKey(ByteString rKey, ByteString lKey) {
        boolean hasRkey = hasRKey(rKey);

        int keyLen;
        if (hasRkey) {
            keyLen = rKey.size() + lKey.size() + 2;
        } else {
            keyLen = lKey.size() + 1;
        }
        ByteBuf keyBuf = Unpooled.buffer(keyLen);
        if (hasRkey) {
            keyBuf.writeByte(HAS_ROUTING_KEY);
            keyBuf.writeBytes(rKey.asReadOnlyByteBuffer());
            keyBuf.writeByte(SEP);
            keyBuf.writeBytes(lKey.asReadOnlyByteBuffer());
        } else {
            keyBuf.writeByte(NO_ROUTING_KEY);
            keyBuf.writeBytes(lKey.asReadOnlyByteBuffer());
        }

        return ByteBufUtil.getBytes(keyBuf);
    }

    static ByteString getLKey(byte[] storeKey, ByteString rKey) {
        boolean hasRKey = hasRKey(rKey);

        int lKeyOffset;
        if (hasRKey) {
            lKeyOffset = rKey.size() + 2;
        } else {
            lKeyOffset = 1;
        }

        return UnsafeByteOperations.unsafeWrap(storeKey, lKeyOffset, storeKey.length - lKeyOffset);
    }

    static StatusCode handleCause(Throwable cause) {
        if (cause instanceof ExecutionException) {
            return handleCause(cause.getCause());
        } else if (cause instanceof MVCCStoreException) {
            MVCCStoreException mse = (MVCCStoreException) cause;
            return mvccCodeToStatusCode(mse.getCode());
        } else {
            return StatusCode.INTERNAL_SERVER_ERROR;
        }
    }

    static StatusCode mvccCodeToStatusCode(Code code) {
        switch (code) {
        case OK:
            return StatusCode.SUCCESS;
        case INTERNAL_ERROR:
            return StatusCode.INTERNAL_SERVER_ERROR;
        case INVALID_ARGUMENT:
            return StatusCode.BAD_REQUEST;
        case ILLEGAL_OP:
            return StatusCode.BAD_REQUEST;
        case UNEXPECTED:
            return StatusCode.UNEXPECTED;
        case BAD_REVISION:
            return StatusCode.BAD_REVISION;
        case SMALLER_REVISION:
            return StatusCode.BAD_REVISION;
        case KEY_NOT_FOUND:
            return StatusCode.KEY_NOT_FOUND;
        case KEY_EXISTS:
            return StatusCode.KEY_EXISTS;
        default:
            return StatusCode.INTERNAL_SERVER_ERROR;
        }
    }

    static KeyValue newKeyValue(ByteString rKey, org.apache.bookkeeper.api.kv.result.KeyValue<byte[], byte[]> kv) {
        if (null == kv) {
            return null;
        }
        return KeyValue.newBuilder().setKey(getLKey(kv.key(), rKey))
                .setValue(UnsafeByteOperations.unsafeWrap(kv.value())).setCreateRevision(kv.createRevision())
                .setModRevision(kv.modifiedRevision()).setVersion(kv.version()).setIsNumber(kv.isNumber())
                .setNumberValue(kv.numberValue()).build();
    }

    static PutResponse processPutResult(RoutingHeader routingHeader, PutResult<byte[], byte[]> result) {
        ByteString rKey = routingHeader.getRKey();
        PutResponse.Builder putRespBuilder = PutResponse.newBuilder().setHeader(ResponseHeader.newBuilder()
                .setCode(mvccCodeToStatusCode(result.code())).setRoutingHeader(routingHeader).build());
        if (null != result.prevKv()) {
            putRespBuilder = putRespBuilder.setPrevKv(newKeyValue(rKey, result.prevKv()));
        }
        return putRespBuilder.build();
    }

    static IncrementResponse processIncrementResult(RoutingHeader routingHeader,
            IncrementResult<byte[], byte[]> result) {
        IncrementResponse.Builder putRespBuilder = IncrementResponse.newBuilder().setHeader(ResponseHeader
                .newBuilder().setCode(mvccCodeToStatusCode(result.code())).setRoutingHeader(routingHeader).build())
                .setTotalAmount(result.totalAmount());
        return putRespBuilder.build();
    }

    static RangeResponse processRangeResult(RoutingHeader routingHeader, RangeResult<byte[], byte[]> result) {
        ByteString rKey = routingHeader.getRKey();
        return RangeResponse.newBuilder().setCount(result.count())
                .setHeader(ResponseHeader.newBuilder().setCode(mvccCodeToStatusCode(result.code()))
                        .setRoutingHeader(routingHeader).build())
                .addAllKvs(Lists.transform(result.kvs(), kv -> newKeyValue(rKey, kv))).setMore(result.more())
                .build();
    }

    static DeleteRangeResponse processDeleteResult(RoutingHeader routingHeader,
            DeleteResult<byte[], byte[]> result) {
        ByteString rKey = routingHeader.getRKey();
        return DeleteRangeResponse.newBuilder()
                .setHeader(ResponseHeader.newBuilder().setCode(mvccCodeToStatusCode(result.code()))
                        .setRoutingHeader(routingHeader).build())
                .setDeleted(result.numDeleted())
                .addAllPrevKvs(Lists.transform(result.prevKvs(), kv -> newKeyValue(rKey, kv))).build();
    }

    static TxnResponse processTxnResult(RoutingHeader routingHeader, TxnResult<byte[], byte[]> txnResult) {
        return TxnResponse.newBuilder()
                .setHeader(ResponseHeader.newBuilder().setCode(mvccCodeToStatusCode(txnResult.code()))
                        .setRoutingHeader(routingHeader).build())
                .setSucceeded(txnResult.isSuccess())
                .addAllResponses(
                        Lists.transform(txnResult.results(), result -> processTxnResult(routingHeader, result)))
                .build();
    }

    static ResponseOp processTxnResult(RoutingHeader routingHeader, Result<byte[], byte[]> result) {

        ResponseOp.Builder respBuilder = ResponseOp.newBuilder();
        switch (result.type()) {
        case PUT:
            PutResult<byte[], byte[]> putResult = (PutResult<byte[], byte[]>) result;
            respBuilder.setResponsePut(processPutResult(routingHeader, putResult));
            break;
        case DELETE:
            DeleteResult<byte[], byte[]> delResult = (DeleteResult<byte[], byte[]>) result;
            respBuilder.setResponseDeleteRange(processDeleteResult(routingHeader, delResult));
            break;
        case RANGE:
            RangeResult<byte[], byte[]> rangeResult = (RangeResult<byte[], byte[]>) result;
            respBuilder.setResponseRange(processRangeResult(routingHeader, rangeResult));
            break;
        default:
            break;
        }
        return respBuilder.build();
    }

    static CompareOp<byte[], byte[]> fromProtoCompare(OpFactory<byte[], byte[]> opFactory, RoutingHeader header,
            Compare compare) {
        ByteString rKey = header.getRKey();
        ByteString lKey = compare.getKey();
        byte[] storeKey = newStoreKey(rKey, lKey);
        CompareResult result = fromProtoCompareResult(compare.getResult());
        switch (compare.getTarget()) {
        case MOD:
            return opFactory.compareModRevision(result, storeKey, compare.getModRevision());
        case CREATE:
            return opFactory.compareCreateRevision(result, storeKey, compare.getCreateRevision());
        case VERSION:
            return opFactory.compareVersion(result, storeKey, compare.getVersion());
        case VALUE:
            return opFactory.compareValue(result, storeKey,
                    (null == compare.getValue() || compare.getValue().size() == 0) ? null
                            : compare.getValue().toByteArray());
        default:
            throw new IllegalArgumentException("Invalid compare target " + compare.getTarget());
        }

    }

    static CompareResult fromProtoCompareResult(Compare.CompareResult result) {
        switch (result) {
        case LESS:
            return CompareResult.LESS;
        case EQUAL:
            return CompareResult.EQUAL;
        case GREATER:
            return CompareResult.GREATER;
        case NOT_EQUAL:
            return CompareResult.NOT_EQUAL;
        default:
            throw new IllegalArgumentException("Invalid compare result " + result);
        }
    }

}