* 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,
 * See the License for the specific language governing permissions and
 * limitations under the License.

package org.apache.bookkeeper.client;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.bookkeeper.bookie.BookieShell.UpdateLedgerNotifier;
import org.apache.bookkeeper.net.BookieSocketAddress;
import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks.GenericCallback;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.RateLimiter;
import com.google.common.util.concurrent.SettableFuture;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

 * Encapsulates updating the ledger metadata operation
public class UpdateLedgerOp {

    private final static Logger LOG = LoggerFactory.getLogger(UpdateLedgerOp.class);
    private final BookKeeper bkc;
    private final BookKeeperAdmin admin;

    public UpdateLedgerOp(final BookKeeper bkc, final BookKeeperAdmin admin) {
        this.bkc = bkc;
        this.admin = admin;

     * Update the bookie id present in the ledger metadata.
     * @param oldBookieId
     *            current bookie id
     * @param newBookieId
     *            new bookie id
     * @param rate
     *            number of ledgers updating per second (default 5 per sec)
     * @param limit
     *            maximum number of ledgers to update (default: no limit). Stop
     *            update if reaching limit
     * @param progressable
     *            report progress of the ledger updates
     * @throws IOException
     *             if there is an error when updating bookie id in ledger
     *             metadata
     * @throws InterruptedException
     *             interrupted exception when update ledger meta
    public void updateBookieIdInLedgers(final BookieSocketAddress oldBookieId, final BookieSocketAddress newBookieId,
            final int rate, final int limit, final UpdateLedgerNotifier progressable) throws BKException, IOException {

        final ThreadFactoryBuilder tfb = new ThreadFactoryBuilder().setNameFormat("UpdateLedgerThread").setDaemon(true);
        final ExecutorService executor = Executors.newSingleThreadExecutor(tfb.build());
        final AtomicInteger issuedLedgerCnt = new AtomicInteger();
        final AtomicInteger updatedLedgerCnt = new AtomicInteger();
        final Future<?> updateBookieCb = executor.submit(new Runnable() {

            public void run() {
                updateLedgers(oldBookieId, newBookieId, rate, limit, progressable);

            private void updateLedgers(final BookieSocketAddress oldBookieId, final BookieSocketAddress newBookieId,
                    final int rate, final int limit, final UpdateLedgerNotifier progressable) {
                try {
                    final AtomicBoolean stop = new AtomicBoolean(false);
                    final Set<Long> outstandings = Collections.newSetFromMap(new ConcurrentHashMap<Long, Boolean>());
                    final RateLimiter throttler = RateLimiter.create(rate);
                    final Iterator<Long> ledgerItr = admin.listLedgers().iterator();
                    final CountDownLatch syncObj = new CountDownLatch(1);

                    // iterate through all the ledgers
                    while (ledgerItr.hasNext() && !stop.get()) {
                        // throttler to control updates per second

                        final Long lId = ledgerItr.next();
                        final ReadLedgerMetadataCb readCb = new ReadLedgerMetadataCb(bkc, lId, oldBookieId, newBookieId);

                        FutureCallback<Void> updateLedgerCb = new UpdateLedgerCb(lId, stop, issuedLedgerCnt,
                                updatedLedgerCnt, outstandings, syncObj, progressable);
                        Futures.addCallback(readCb.getFutureListener(), updateLedgerCb);

                        if (limit != Integer.MIN_VALUE && issuedLedgerCnt.get() >= limit || !ledgerItr.hasNext()) {
                        bkc.getLedgerManager().readLedgerMetadata(lId, readCb);
                    // waiting till all the issued ledgers are finished
                } catch (IOException ioe) {
                    LOG.error("Exception while updating ledger", ioe);
                    throw new RuntimeException("Exception while updating ledger", ioe.getCause());
                } catch (InterruptedException ie) {
                    LOG.error("Exception while updating ledger metadata", ie);
                    throw new RuntimeException("Exception while updating ledger", ie.getCause());
        try {
            // Wait to finish the issued ledgers.
        } catch (ExecutionException ee) {
            throw new IOException("Exception while updating ledger", ee);
        } catch (InterruptedException ie) {
            throw new IOException("Exception while updating ledger", ie);
        } finally {

    private final static class UpdateLedgerCb implements FutureCallback<Void> {
        final long ledgerId;
        final AtomicBoolean stop;
        final AtomicInteger issuedLedgerCnt;
        final AtomicInteger updatedLedgerCnt;
        final Set<Long> outstandings;
        final CountDownLatch syncObj;
        final UpdateLedgerNotifier progressable;

        public UpdateLedgerCb(long ledgerId, AtomicBoolean stop, AtomicInteger issuedLedgerCnt,
                AtomicInteger updatedLedgerCnt, Set<Long> outstandings, CountDownLatch syncObj,
                UpdateLedgerNotifier progressable) {
            this.ledgerId = ledgerId;
            this.stop = stop;
            this.issuedLedgerCnt = issuedLedgerCnt;
            this.updatedLedgerCnt = updatedLedgerCnt;
            this.outstandings = outstandings;
            this.syncObj = syncObj;
            this.progressable = progressable;

        public void onFailure(Throwable th) {
            LOG.error("Error updating ledger {}", ledgerId, th);

        public void onSuccess(Void obj) {
            // may print progress
            progressable.progress(updatedLedgerCnt.get(), issuedLedgerCnt.get());

        private void finishUpdateLedger() {
            if (outstandings.isEmpty() && stop.get()) {
                LOG.info("Total number of ledgers issued={} updated={}", issuedLedgerCnt.get(), updatedLedgerCnt.get());

    private final static class ReadLedgerMetadataCb implements GenericCallback<LedgerMetadata> {
        final BookKeeper bkc;
        final Long ledgerId;
        final BookieSocketAddress curBookieAddr;
        final BookieSocketAddress toBookieAddr;
        SettableFuture<Void> future = SettableFuture.create();
        public ReadLedgerMetadataCb(BookKeeper bkc, Long ledgerId, BookieSocketAddress curBookieAddr,
                BookieSocketAddress toBookieAddr) {
            this.bkc = bkc;
            this.ledgerId = ledgerId;
            this.curBookieAddr = curBookieAddr;
            this.toBookieAddr = toBookieAddr;

        ListenableFuture<Void> getFutureListener() {
            return future;

        public void operationComplete(int rc, LedgerMetadata metadata) {
            if (BKException.Code.NoSuchLedgerExistsException == rc) {
                return; // this is OK
            } else if (BKException.Code.OK != rc) {
                // open ledger failed.
                LOG.error("Get ledger metadata {} failed. Error code {}", ledgerId, rc);
            boolean updateEnsemble = false;
            for (ArrayList<BookieSocketAddress> ensembles : metadata.getEnsembles().values()) {
                int index = ensembles.indexOf(curBookieAddr);
                if (-1 != index) {
                    ensembles.set(index, toBookieAddr);
                    updateEnsemble = true;
            if (!updateEnsemble) {
                return; // ledger doesn't contains the given curBookieId
            final GenericCallback<Void> writeCb = new GenericCallback<Void>() {
                public void operationComplete(int rc, Void result) {
                    if (rc != BKException.Code.OK) {
                        // metadata update failed
                        LOG.error("Ledger {} metadata update failed. Error code {}", ledgerId, rc);
            bkc.getLedgerManager().writeLedgerMetadata(ledgerId, metadata, writeCb);