bookkeeper-server/src/test/java/org/apache/bookkeeper/meta/MockLedgerManager.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.meta;

import com.google.common.collect.Lists;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import org.apache.bookkeeper.client.BKException;
import org.apache.bookkeeper.client.api.LedgerMetadata;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks.LedgerMetadataListener;
import org.apache.bookkeeper.proto.BookkeeperInternalCallbacks.Processor;
import org.apache.bookkeeper.versioning.LongVersion;
import org.apache.bookkeeper.versioning.Version;
import org.apache.bookkeeper.versioning.Versioned;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.zookeeper.AsyncCallback;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Mock implementation of Ledger Manager.
 */
public class MockLedgerManager implements LedgerManager {
    static final Logger LOG = LoggerFactory.getLogger(MockLedgerManager.class);

    /**
     * Hook for injecting errors or delays.
     */
    public interface Hook {
        CompletableFuture<Void> runHook(long ledgerId, LedgerMetadata metadata);
    }

    final Map<Long, Pair<LongVersion, byte[]>> metadataMap;
    final ExecutorService executor;
    final boolean ownsExecutor;
    final LedgerMetadataSerDe serDe;
    private Hook preWriteHook = (ledgerId, metadata) -> FutureUtils.value(null);

    public MockLedgerManager() {
        this(new ConcurrentHashMap<>(),
             Executors.newSingleThreadExecutor((r) -> new Thread(r, "MockLedgerManager")), true);
    }

    private MockLedgerManager(Map<Long, Pair<LongVersion, byte[]>> metadataMap,
                              ExecutorService executor, boolean ownsExecutor) {
        this.metadataMap = metadataMap;
        this.executor = executor;
        this.ownsExecutor = ownsExecutor;
        this.serDe = new LedgerMetadataSerDe();
    }

    public MockLedgerManager newClient() {
        return new MockLedgerManager(metadataMap, executor, false);
    }

    private Versioned<LedgerMetadata> readMetadata(long ledgerId) throws Exception {
        Pair<LongVersion, byte[]> pair = metadataMap.get(ledgerId);
        if (pair == null) {
            return null;
        } else {
            return new Versioned<>(serDe.parseConfig(pair.getRight(), ledgerId, Optional.empty()), pair.getLeft());
        }
    }

    public void setPreWriteHook(Hook hook) {
        this.preWriteHook = hook;
    }

    public void executeCallback(Runnable r) {
        r.run();
    }

    @Override
    public CompletableFuture<Versioned<LedgerMetadata>> createLedgerMetadata(long ledgerId, LedgerMetadata metadata) {
        CompletableFuture<Versioned<LedgerMetadata>> promise = new CompletableFuture<>();
        executor.submit(() -> {
                if (metadataMap.containsKey(ledgerId)) {
                    executeCallback(() -> promise.completeExceptionally(new BKException.BKLedgerExistException()));
                } else {
                    try {
                        metadataMap.put(ledgerId, Pair.of(new LongVersion(0L), serDe.serialize(metadata)));
                        Versioned<LedgerMetadata> readBack = readMetadata(ledgerId);
                        executeCallback(() -> promise.complete(readBack));
                    } catch (Exception e) {
                        LOG.error("Error reading back written metadata", e);
                        executeCallback(() -> promise.completeExceptionally(new BKException.MetaStoreException()));
                    }
                }
            });
        return promise;
    }

    @Override
    public CompletableFuture<Void> removeLedgerMetadata(long ledgerId, Version version) {
        return CompletableFuture.completedFuture(null);
    }

    @Override
    public CompletableFuture<Versioned<LedgerMetadata>> readLedgerMetadata(long ledgerId) {
        CompletableFuture<Versioned<LedgerMetadata>> promise = new CompletableFuture<>();
        executor.submit(() -> {
                try {
                    Versioned<LedgerMetadata> metadata = readMetadata(ledgerId);
                    if (metadata == null) {
                        executeCallback(() -> promise.completeExceptionally(
                                                new BKException.BKNoSuchLedgerExistsOnMetadataServerException()));
                    } else {
                        executeCallback(() -> promise.complete(metadata));
                    }
                } catch (Exception e) {
                    LOG.error("Error reading metadata", e);
                    executeCallback(() -> promise.completeExceptionally(new BKException.MetaStoreException()));
                }
            });
        return promise;
    }

    @Override
    public CompletableFuture<Versioned<LedgerMetadata>> writeLedgerMetadata(long ledgerId, LedgerMetadata metadata,
                                                                            Version currentVersion) {
        CompletableFuture<Versioned<LedgerMetadata>> promise = new CompletableFuture<>();
        preWriteHook.runHook(ledgerId, metadata)
            .thenComposeAsync((ignore) -> {
                    try {
                        Versioned<LedgerMetadata> oldMetadata = readMetadata(ledgerId);
                        if (oldMetadata == null) {
                            return FutureUtils.exception(
                                    new BKException.BKNoSuchLedgerExistsOnMetadataServerException());
                        } else if (!oldMetadata.getVersion().equals(currentVersion)) {
                            return FutureUtils.exception(new BKException.BKMetadataVersionException());
                        } else {
                            LongVersion oldVersion = (LongVersion) oldMetadata.getVersion();
                            metadataMap.put(ledgerId, Pair.of(new LongVersion(oldVersion.getLongVersion() + 1),
                                                              serDe.serialize(metadata)));
                            Versioned<LedgerMetadata> readBack = readMetadata(ledgerId);
                            return FutureUtils.value(readBack);
                        }
                    } catch (Exception e) {
                        LOG.error("Error writing metadata", e);
                        return FutureUtils.exception(e);
                    }
                }, executor)
            .whenComplete((res, ex) -> {
                    if (ex != null) {
                        Throwable cause = (ex instanceof CompletionException) ? ex.getCause() : ex;
                        executeCallback(() -> promise.completeExceptionally(cause));
                    } else {
                        executeCallback(() -> promise.complete(res));
                    }
                });
        return promise;
    }

    @Override
    public void registerLedgerMetadataListener(long ledgerId, LedgerMetadataListener listener) {}

    @Override
    public void unregisterLedgerMetadataListener(long ledgerId, LedgerMetadataListener listener) {}

    @Override
    public void asyncProcessLedgers(Processor<Long> processor, AsyncCallback.VoidCallback finalCb,
                                    Object context, int successRc, int failureRc) {
    }

    @Override
    public LedgerRangeIterator getLedgerRanges(long zkOpTimeoutMs) {
        List<Long> ledgerIds = new ArrayList<>(metadataMap.keySet());
        ledgerIds.sort(Comparator.naturalOrder());
        List<List<Long>> partitions = Lists.partition(ledgerIds, 100);
        return new LedgerRangeIterator() {
            int i = 0;
            @Override
            public boolean hasNext() {
                if (i >= partitions.size()) {
                    return false;
                } else {
                    return true;
                }
            }

            @Override
            public LedgerRange next() {
                return new LedgerRange(new HashSet<>(partitions.get(i++)));
            }
        };
    }

    @Override
    public void close() {
        if (ownsExecutor) {
            executor.shutdownNow();
        }
    }

}