metadata-drivers/etcd/src/main/java/org/apache/bookkeeper/metadata/etcd/Etcd64bitIdGenerator.java

/*
 * 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.metadata.etcd;

import static com.google.common.base.Preconditions.checkArgument;
import static org.apache.bookkeeper.metadata.etcd.EtcdConstants.EMPTY_BS;

import io.etcd.jetcd.ByteSequence;
import io.etcd.jetcd.KV;
import io.etcd.jetcd.KeyValue;
import io.etcd.jetcd.Txn;
import io.etcd.jetcd.kv.GetResponse;
import io.etcd.jetcd.op.Cmp;
import io.etcd.jetcd.op.Cmp.Op;
import io.etcd.jetcd.op.CmpTarget;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.client.BKException.Code;
import org.apache.bookkeeper.meta.LedgerIdGenerator;
import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks.GenericCallback;

/**
 * Generate 64-bit ledger ids from a bucket.
 *
 * <p>The most significant 8 bits is used as bucket id. The remaining 56 bits are
 * used as the id generated per bucket.
 */
@Slf4j
class Etcd64bitIdGenerator implements LedgerIdGenerator {

    static final long MAX_ID_PER_BUCKET = 0x00ffffffffffffffL;
    static final long BUCKET_ID_MASK    = 0xff00000000000000L;
    static final int BUCKET_ID_SHIFT    = 56;
    static final int NUM_BUCKETS        = 0x80;

    static int getBucketId(long lid) {
        return (int) ((lid & BUCKET_ID_MASK) >>> BUCKET_ID_SHIFT);
    }

    static long getIdInBucket(long lid) {
        return lid & MAX_ID_PER_BUCKET;
    }

    private static final AtomicIntegerFieldUpdater<Etcd64bitIdGenerator> nextBucketIdUpdater =
        AtomicIntegerFieldUpdater.newUpdater(Etcd64bitIdGenerator.class, "nextBucketId");

    private final String scope;
    private final KV kvClient;
    private volatile int nextBucketId;

    Etcd64bitIdGenerator(KV kvClient, String scope) {
        this.kvClient = kvClient;
        this.scope = scope;
        this.nextBucketId = ThreadLocalRandom.current().nextInt(NUM_BUCKETS);
    }

    int nextBucketId() {
        while (true) {
            int bucketId = nextBucketIdUpdater.incrementAndGet(this);
            if (bucketId >= NUM_BUCKETS) {
                if (nextBucketIdUpdater.compareAndSet(this, bucketId, 0)) {
                    bucketId = 0;
                } else {
                    // someone has been updated bucketId, try it again.
                    continue;
                }
            }
            return bucketId;
        }
    }

    @Override
    public void generateLedgerId(GenericCallback<Long> cb) {
        int bucketId = nextBucketId();
        checkArgument(bucketId >= 0 && bucketId < NUM_BUCKETS,
            "Invalid bucket id : " + bucketId);

        ByteSequence bucketKey = ByteSequence.from(EtcdUtils.getBucketPath(scope, bucketId), StandardCharsets.UTF_8);
        Txn txn = kvClient.txn()
            .If(new Cmp(bucketKey, Op.GREATER, CmpTarget.createRevision(0)))
            .Then(
                io.etcd.jetcd.op.Op.put(bucketKey, EMPTY_BS, PutOption.DEFAULT),
                io.etcd.jetcd.op.Op.get(bucketKey, GetOption.DEFAULT)
            )
            .Else(
                io.etcd.jetcd.op.Op.put(bucketKey, EMPTY_BS, PutOption.DEFAULT),
                io.etcd.jetcd.op.Op.get(bucketKey, GetOption.DEFAULT)
            );
        txn.commit()
            .thenAccept(txnResponse -> {
                if (txnResponse.getGetResponses().size() <= 0) {
                    cb.operationComplete(Code.UnexpectedConditionException, null);
                } else {
                    GetResponse resp = txnResponse.getGetResponses().get(0);
                    if (resp.getCount() > 0) {
                        KeyValue kv = resp.getKvs().get(0);
                        if (kv.getVersion() > MAX_ID_PER_BUCKET) {
                            log.warn("Etcd bucket '{}' is overflowed", bucketKey.toString(StandardCharsets.UTF_8));
                            // the bucket is overflowed, moved to next bucket.
                            generateLedgerId(cb);
                        } else {
                            long version = kv.getVersion();
                            long lid = ((((long) bucketId) << BUCKET_ID_SHIFT) & BUCKET_ID_MASK)
                                | (version & MAX_ID_PER_BUCKET);
                            cb.operationComplete(Code.OK, lid);
                        }
                    } else {
                        cb.operationComplete(Code.UnexpectedConditionException, null);
                    }
                }
            })
            .exceptionally(cause -> {
                cb.operationComplete(Code.MetaStoreException, null);
                return null;
            });
    }

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