bookkeeper-server/src/main/java/org/apache/bookkeeper/discover/ZKRegistrationClient.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.discover;

import static org.apache.bookkeeper.util.BookKeeperConstants.AVAILABLE_NODE;
import static org.apache.bookkeeper.util.BookKeeperConstants.COOKIE_NODE;
import static org.apache.bookkeeper.util.BookKeeperConstants.READONLY;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Sets;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.client.BKException;
import org.apache.bookkeeper.client.BKException.ZKException;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.net.BookieId;
import org.apache.bookkeeper.proto.DataFormats.BookieServiceInfoFormat;
import org.apache.bookkeeper.versioning.LongVersion;
import org.apache.bookkeeper.versioning.Version;
import org.apache.bookkeeper.versioning.Version.Occurred;
import org.apache.bookkeeper.versioning.Versioned;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

/**
 * ZooKeeper based {@link RegistrationClient}.
 */

@Slf4j
public class ZKRegistrationClient implements RegistrationClient {

    static final int ZK_CONNECT_BACKOFF_MS = 200;

    class WatchTask
        implements Runnable,
                   Watcher,
                   BiConsumer<Versioned<Set<BookieId>>, Throwable>,
                   AutoCloseable {

        private final String regPath;
        private final Set<RegistrationListener> listeners;
        private volatile boolean closed = false;
        private Set<BookieId> bookies = null;
        private Version version = Version.NEW;
        private final CompletableFuture<Void> firstRunFuture;

        WatchTask(String regPath, CompletableFuture<Void> firstRunFuture) {
            this.regPath = regPath;
            this.listeners = new CopyOnWriteArraySet<>();
            this.firstRunFuture = firstRunFuture;
        }

        public int getNumListeners() {
            return listeners.size();
        }

        public boolean addListener(RegistrationListener listener) {
            if (listeners.add(listener)) {
                if (null != bookies) {
                    scheduler.execute(() -> {
                            listener.onBookiesChanged(
                                    new Versioned<>(bookies, version));
                        });
                }
            }
            return true;
        }

        public boolean removeListener(RegistrationListener listener) {
            return listeners.remove(listener);
        }

        void watch() {
            scheduleWatchTask(0L);
        }

        private void scheduleWatchTask(long delayMs) {
            try {
                scheduler.schedule(this, delayMs, TimeUnit.MILLISECONDS);
            } catch (RejectedExecutionException ree) {
                log.warn("Failed to schedule watch bookies task", ree);
            }
        }

        @Override
        public void run() {
            if (isClosed()) {
                return;
            }

            getChildren(regPath, this)
                .whenCompleteAsync(this, scheduler);
        }

        @Override
        public void accept(Versioned<Set<BookieId>> bookieSet, Throwable throwable) {
            if (throwable != null) {
                if (firstRunFuture.isDone()) {
                    scheduleWatchTask(ZK_CONNECT_BACKOFF_MS);
                } else {
                    firstRunFuture.completeExceptionally(throwable);
                }
                return;
            }

            if (this.version.compare(bookieSet.getVersion()) == Occurred.BEFORE) {
                this.version = bookieSet.getVersion();
                this.bookies = bookieSet.getValue();
                if (!listeners.isEmpty()) {
                    for (RegistrationListener listener : listeners) {
                        listener.onBookiesChanged(bookieSet);
                    }
                }
            }
            FutureUtils.complete(firstRunFuture, null);
        }

        @Override
        public void process(WatchedEvent event) {
            if (EventType.None == event.getType()) {
                if (KeeperState.Expired == event.getState()) {
                    scheduleWatchTask(ZK_CONNECT_BACKOFF_MS);
                }
                return;
            }

            // re-read the bookie list
            scheduleWatchTask(0L);
        }

        boolean isClosed() {
            return closed;
        }

        @Override
        public void close() {
            closed = true;
        }
    }

    private final ZooKeeper zk;
    private final ScheduledExecutorService scheduler;
    @Getter(AccessLevel.PACKAGE)
    private WatchTask watchWritableBookiesTask = null;
    @Getter(AccessLevel.PACKAGE)
    private WatchTask watchReadOnlyBookiesTask = null;
    private final ConcurrentHashMap<BookieId, Versioned<BookieServiceInfo>> bookieServiceInfoCache =
                                                                            new ConcurrentHashMap<>();
    private final Watcher bookieServiceInfoCacheInvalidation;
    private final boolean bookieAddressTracking;
    // registration paths
    private final String bookieRegistrationPath;
    private final String bookieAllRegistrationPath;
    private final String bookieReadonlyRegistrationPath;

    public ZKRegistrationClient(ZooKeeper zk,
                                String ledgersRootPath,
                                ScheduledExecutorService scheduler,
                                boolean bookieAddressTracking) {
        this.zk = zk;
        this.scheduler = scheduler;
        // Following Bookie Network Address Changes is an expensive operation
        // as it requires additional ZooKeeper watches
        // we can disable this feature, in case the BK cluster has only
        // static addresses
        this.bookieAddressTracking = bookieAddressTracking;
        this.bookieServiceInfoCacheInvalidation = bookieAddressTracking
                                                    ? new BookieServiceInfoCacheInvalidationWatcher() : null;
        this.bookieRegistrationPath = ledgersRootPath + "/" + AVAILABLE_NODE;
        this.bookieAllRegistrationPath = ledgersRootPath + "/" + COOKIE_NODE;
        this.bookieReadonlyRegistrationPath = this.bookieRegistrationPath + "/" + READONLY;
    }

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

    public boolean isBookieAddressTracking() {
        return bookieAddressTracking;
    }

    public ZooKeeper getZk() {
        return zk;
    }

    @Override
    public CompletableFuture<Versioned<Set<BookieId>>> getWritableBookies() {
        return getChildren(bookieRegistrationPath, null);
    }

    @Override
    public CompletableFuture<Versioned<Set<BookieId>>> getAllBookies() {
        return getChildren(bookieAllRegistrationPath, null);
    }

    @Override
    public CompletableFuture<Versioned<Set<BookieId>>> getReadOnlyBookies() {
        return getChildren(bookieReadonlyRegistrationPath, null);
    }

    @Override
    public CompletableFuture<Versioned<BookieServiceInfo>> getBookieServiceInfo(BookieId bookieId) {
        // we can only serve data from cache here,
        // because it can happen than this method is called inside the main
        // zookeeper client event loop thread
        Versioned<BookieServiceInfo> resultFromCache = bookieServiceInfoCache.get(bookieId);
        if (log.isDebugEnabled()) {
            log.debug("getBookieServiceInfo {} -> {}", bookieId, resultFromCache);
        }
        if (resultFromCache != null) {
            return CompletableFuture.completedFuture(resultFromCache);
        } else {
            return FutureUtils.exception(new BKException.BKBookieHandleNotAvailableException());
        }
    }

    /**
     * Read BookieServiceInfo from ZooKeeper and updates the local cache.
     *
     * @param bookieId
     * @return an handle to the result of the operation.
     */
    private CompletableFuture<Versioned<BookieServiceInfo>> readBookieServiceInfoAsync(BookieId bookieId) {
        String pathAsWritable = bookieRegistrationPath + "/" + bookieId;
        String pathAsReadonly = bookieReadonlyRegistrationPath + "/" + bookieId;

        CompletableFuture<Versioned<BookieServiceInfo>> promise = new CompletableFuture<>();
        zk.getData(pathAsWritable, bookieServiceInfoCacheInvalidation,
                (int rc, String path, Object o, byte[] bytes, Stat stat) -> {
            if (KeeperException.Code.OK.intValue() == rc) {
                try {
                    BookieServiceInfo bookieServiceInfo = deserializeBookieServiceInfo(bookieId, bytes);
                    Versioned<BookieServiceInfo> result = new Versioned<>(bookieServiceInfo,
                            new LongVersion(stat.getCversion()));
                    log.info("Update BookieInfoCache (writable bookie) {} -> {}", bookieId, result.getValue());
                    bookieServiceInfoCache.put(bookieId, result);
                    promise.complete(result);
                } catch (IOException ex) {
                    log.error("Cannot update BookieInfo for ", ex);
                    promise.completeExceptionally(KeeperException.create(KeeperException.Code.get(rc), path)
                            .initCause(ex));
                    return;
                }
            } else if (KeeperException.Code.NONODE.intValue() == rc) {
                // not found, looking for a readonly bookie
                zk.getData(pathAsReadonly, bookieServiceInfoCacheInvalidation,
                        (int rc2, String path2, Object o2, byte[] bytes2, Stat stat2) -> {
                    if (KeeperException.Code.OK.intValue() == rc2) {
                        try {
                            BookieServiceInfo bookieServiceInfo = deserializeBookieServiceInfo(bookieId, bytes2);
                            Versioned<BookieServiceInfo> result =
                                    new Versioned<>(bookieServiceInfo, new LongVersion(stat2.getCversion()));
                            log.info("Update BookieInfoCache (readonly bookie) {} -> {}", bookieId, result.getValue());
                            bookieServiceInfoCache.put(bookieId, result);
                            promise.complete(result);
                        } catch (IOException ex) {
                            log.error("Cannot update BookieInfo for ", ex);
                            promise.completeExceptionally(KeeperException.create(KeeperException.Code.get(rc2), path2)
                                    .initCause(ex));
                            return;
                        }
                    } else {
                        // not found as writable and readonly, the bookie is offline
                        promise.completeExceptionally(BKException.create(BKException.Code.NoBookieAvailableException));
                    }
                }, null);
            } else {
                promise.completeExceptionally(KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }, null);
        return promise;
    }

    @SuppressWarnings("unchecked")
    @VisibleForTesting
    static BookieServiceInfo deserializeBookieServiceInfo(BookieId bookieId, byte[] bookieServiceInfo)
            throws IOException {
        if (bookieServiceInfo == null || bookieServiceInfo.length == 0) {
            return BookieServiceInfoUtils.buildLegacyBookieServiceInfo(bookieId.toString());
        }

        BookieServiceInfoFormat builder = BookieServiceInfoFormat.parseFrom(bookieServiceInfo);
        BookieServiceInfo bsi = new BookieServiceInfo();
        List<BookieServiceInfo.Endpoint> endpoints = builder.getEndpointsList().stream()
                .map(e -> {
                    BookieServiceInfo.Endpoint endpoint = new BookieServiceInfo.Endpoint();
                    endpoint.setId(e.getId());
                    endpoint.setPort(e.getPort());
                    endpoint.setHost(e.getHost());
                    endpoint.setProtocol(e.getProtocol());
                    endpoint.setAuth(e.getAuthList());
                    endpoint.setExtensions(e.getExtensionsList());
                    return endpoint;
                })
                .collect(Collectors.toList());

        bsi.setEndpoints(endpoints);
        bsi.setProperties(builder.getPropertiesMap());

        return bsi;
    }

    /**
     * Reads the list of bookies at the given path and eagerly caches the BookieServiceInfo
     * structure.
     *
     * @param regPath the path on ZooKeeper
     * @param watcher an optional watcher
     * @return an handle to the operation
     */
    private CompletableFuture<Versioned<Set<BookieId>>> getChildren(String regPath, Watcher watcher) {
        CompletableFuture<Versioned<Set<BookieId>>> future = FutureUtils.createFuture();
        zk.getChildren(regPath, watcher, (rc, path, ctx, children, stat) -> {
            if (KeeperException.Code.OK.intValue() != rc) {
                ZKException zke = new ZKException(KeeperException.create(KeeperException.Code.get(rc), path));
                future.completeExceptionally(zke.fillInStackTrace());
                return;
            }

            Version version = new LongVersion(stat.getCversion());
            Set<BookieId> bookies = convertToBookieAddresses(children);
            List<CompletableFuture<Versioned<BookieServiceInfo>>> bookieInfoUpdated = new ArrayList<>(bookies.size());
            for (BookieId id : bookies) {
                // update the cache for new bookies
                if (!bookieServiceInfoCache.containsKey(id)) {
                    bookieInfoUpdated.add(readBookieServiceInfoAsync(id));
                }
            }
            if (bookieInfoUpdated.isEmpty()) {
                future.complete(new Versioned<>(bookies, version));
            } else {
                FutureUtils
                        .collect(bookieInfoUpdated)
                        .whenComplete((List<Versioned<BookieServiceInfo>> info, Throwable error) -> {
                            // we are ignoring errors intentionally
                            // there could be bookies that publish unparseable information
                            // or other temporary/permanent errors
                            future.complete(new Versioned<>(bookies, version));
                        });
            }
        }, null);
        return future;
    }


    @Override
    public synchronized CompletableFuture<Void> watchWritableBookies(RegistrationListener listener) {
        CompletableFuture<Void> f;
        if (null == watchWritableBookiesTask) {
            f = new CompletableFuture<>();
            watchWritableBookiesTask = new WatchTask(bookieRegistrationPath, f);
            f = f.whenComplete((value, cause) -> {
                if (null != cause) {
                    unwatchWritableBookies(listener);
                }
            });
        } else {
            f = watchWritableBookiesTask.firstRunFuture;
        }

        watchWritableBookiesTask.addListener(listener);
        if (watchWritableBookiesTask.getNumListeners() == 1) {
            watchWritableBookiesTask.watch();
        }
        return f;
    }

    @Override
    public synchronized void unwatchWritableBookies(RegistrationListener listener) {
        if (null == watchWritableBookiesTask) {
            return;
        }

        watchWritableBookiesTask.removeListener(listener);
        if (watchWritableBookiesTask.getNumListeners() == 0) {
            watchWritableBookiesTask.close();
            watchWritableBookiesTask = null;
        }
    }

    @Override
    public synchronized CompletableFuture<Void> watchReadOnlyBookies(RegistrationListener listener) {
        CompletableFuture<Void> f;
        if (null == watchReadOnlyBookiesTask) {
            f = new CompletableFuture<>();
            watchReadOnlyBookiesTask = new WatchTask(bookieReadonlyRegistrationPath, f);
            f = f.whenComplete((value, cause) -> {
                if (null != cause) {
                    unwatchReadOnlyBookies(listener);
                }
            });
        } else {
            f = watchReadOnlyBookiesTask.firstRunFuture;
        }

        watchReadOnlyBookiesTask.addListener(listener);
        if (watchReadOnlyBookiesTask.getNumListeners() == 1) {
            watchReadOnlyBookiesTask.watch();
        }
        return f;
    }

    @Override
    public synchronized void unwatchReadOnlyBookies(RegistrationListener listener) {
        if (null == watchReadOnlyBookiesTask) {
            return;
        }

        watchReadOnlyBookiesTask.removeListener(listener);
        if (watchReadOnlyBookiesTask.getNumListeners() == 0) {
            watchReadOnlyBookiesTask.close();
            watchReadOnlyBookiesTask = null;
        }
    }

    private static HashSet<BookieId> convertToBookieAddresses(List<String> children) {
        // Read the bookie addresses into a set for efficient lookup
        HashSet<BookieId> newBookieAddrs = Sets.newHashSet();
        for (String bookieAddrString : children) {
            if (READONLY.equals(bookieAddrString)) {
                continue;
            }
            BookieId bookieAddr = BookieId.parse(bookieAddrString);
            newBookieAddrs.add(bookieAddr);
        }
        return newBookieAddrs;
    }

    private static BookieId stripBookieIdFromPath(String path) {
        if (path == null) {
            return null;
        }
        final int slash = path.lastIndexOf('/');
        if (slash >= 0) {
            try {
                return BookieId.parse(path.substring(slash + 1));
            } catch (IllegalArgumentException e) {
                log.warn("Cannot decode bookieId from {}", path, e);
            }
        }
        return null;
    }

    private class BookieServiceInfoCacheInvalidationWatcher implements Watcher {

        @Override
        public void process(WatchedEvent we) {
            if (log.isDebugEnabled()) {
                log.debug("zk event {} for {} state {}", we.getType(), we.getPath(), we.getState());
            }
            if (we.getState() == KeeperState.Expired) {
                log.info("zk session expired, invalidating cache");
                bookieServiceInfoCache.clear();
                return;
            }
            BookieId bookieId = stripBookieIdFromPath(we.getPath());
            if (bookieId == null) {
                return;
            }
            switch (we.getType()) {
                case NodeDeleted:
                    log.info("Invalidate cache for {}", bookieId);
                    bookieServiceInfoCache.remove(bookieId);
                    break;
                case NodeDataChanged:
                    log.info("refresh cache for {}", bookieId);
                    readBookieServiceInfoAsync(bookieId);
                    break;
                default:
                    if (log.isDebugEnabled()) {
                        log.debug("ignore cache event {} for {}", we.getType(), bookieId);
                    }
                    break;
            }
        }
    }

}