stream/storage/impl/src/main/java/org/apache/bookkeeper/stream/storage/impl/store/MVCCStoreFactoryImpl.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.stream.storage.impl.store;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import java.io.File;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.function.Supplier;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.common.coder.ByteArrayCoder;
import org.apache.bookkeeper.common.concurrent.FutureUtils;
import org.apache.bookkeeper.common.exceptions.ObjectClosedException;
import org.apache.bookkeeper.common.util.OrderedScheduler;
import org.apache.bookkeeper.common.util.SharedResourceManager;
import org.apache.bookkeeper.statelib.StateStores;
import org.apache.bookkeeper.statelib.api.StateStoreSpec;
import org.apache.bookkeeper.statelib.api.checkpoint.CheckpointStore;
import org.apache.bookkeeper.statelib.api.mvcc.MVCCAsyncStore;
import org.apache.bookkeeper.stream.protocol.RangeId;
import org.apache.bookkeeper.stream.storage.StorageResources;
import org.apache.bookkeeper.stream.storage.conf.StorageConfiguration;
import org.apache.distributedlog.api.namespace.Namespace;

/**
 * A default implementation of {@link MVCCStoreFactory}.
 */
@Accessors(fluent = true)
@Slf4j
public class MVCCStoreFactoryImpl implements MVCCStoreFactory {

    // store supplier
    private final Supplier<MVCCAsyncStore<byte[], byte[]>> storeSupplier;
    // storage resources
    private final StorageResources storageResources;
    // scheduler
    @Getter(value = AccessLevel.PACKAGE)
    private final OrderedScheduler writeIOScheduler;
    @Getter(value = AccessLevel.PACKAGE)
    private final OrderedScheduler readIOScheduler;
    @Getter(value = AccessLevel.PACKAGE)
    private final OrderedScheduler checkpointScheduler;
    // dirs
    private final File[] localStateDirs;
    // checkpoint manager
    private final Supplier<CheckpointStore> checkpointStoreSupplier;
    private CheckpointStore checkpointStore;
    // stores
    private final Map<Long, Map<RangeId, MVCCAsyncStore<byte[], byte[]>>> stores;
    private final boolean serveReadOnlyTable;
    private boolean closed = false;
    private final StorageConfiguration storageConf;

    public MVCCStoreFactoryImpl(Supplier<Namespace> namespaceSupplier,
                                Supplier<CheckpointStore> checkpointStoreSupplier,
                                File[] localStoreDirs,
                                StorageResources storageResources,
                                boolean serveReadOnlyTable, StorageConfiguration storageConf) {
        this.storeSupplier = StateStores.mvccKvBytesStoreSupplier(namespaceSupplier);
        this.storageResources = storageResources;
        this.writeIOScheduler =
            SharedResourceManager.shared().get(storageResources.ioWriteScheduler());
        this.readIOScheduler =
            SharedResourceManager.shared().get(storageResources.ioReadScheduler());
        this.checkpointScheduler =
            SharedResourceManager.shared().get(storageResources.checkpointScheduler());
        this.localStateDirs = localStoreDirs;
        this.checkpointStoreSupplier = checkpointStoreSupplier;
        this.storageConf = storageConf;
        this.stores = Maps.newHashMap();
        this.serveReadOnlyTable = serveReadOnlyTable;
    }

    private ScheduledExecutorService chooseWriteIOExecutor(long streamId) {
        return writeIOScheduler.chooseThread(streamId);
    }

    private ScheduledExecutorService chooseReadIOExecutor(long streamId) {
        return readIOScheduler.chooseThread(streamId);
    }

    private ScheduledExecutorService chooseCheckpointIOExecutor(long streamId) {
        return checkpointScheduler.chooseThread(streamId);
    }

    private File chooseLocalStoreDir(long streamId) {
        int idx = (int) (streamId % localStateDirs.length);
        return localStateDirs[idx];
    }

    static String normalizedName(long id) {
        return String.format("%018d", id);
    }

    static String streamName(long scId,
                             long streamId,
                             long rangeId) {
        // TODO: change to filesystem path
        return String.format(
            "%s_%018d_%018d_%018d",
            "streams",
            scId,
            streamId,
            rangeId);
    }

    private synchronized void addStore(long scId, long streamId, long rangeId,
                                       MVCCAsyncStore<byte[], byte[]> store) {
        Map<RangeId, MVCCAsyncStore<byte[], byte[]>> scStores = stores.get(scId);
        if (null == scStores) {
            scStores = Maps.newHashMap();
            stores.putIfAbsent(scId, scStores);
        }
        RangeId rid = RangeId.of(streamId, rangeId);
        MVCCAsyncStore<byte[], byte[]> oldStore = scStores.get(rid);
        if (null != oldStore) {
            store.closeAsync();
        } else {
            log.info("Add store (scId = {}, streamId = {}, rangeId = {}) at storage container ({})",
                scId, streamId, rangeId, scId);
            scStores.put(rid, store);
        }
    }

    private synchronized MVCCAsyncStore<byte[], byte[]> getStore(long scId, long streamId, long rangeId) {
        Map<RangeId, MVCCAsyncStore<byte[], byte[]>> scStores = stores.get(scId);
        if (null == scStores) {
            return null;
        } else {
            RangeId rid = RangeId.of(streamId, rangeId);
            return scStores.get(rid);
        }
    }

    @Override
    public CompletableFuture<MVCCAsyncStore<byte[], byte[]>> openStore(long scId, long streamId, long rangeId,
                                                                       int ttlSeconds) {
        MVCCAsyncStore<byte[], byte[]> store = getStore(scId, streamId, rangeId);
        if (null == store) {
            return newStore(scId, streamId, rangeId, ttlSeconds);
        } else {
            return FutureUtils.value(store);
        }
    }

    CompletableFuture<MVCCAsyncStore<byte[], byte[]>> newStore(long scId, long streamId, long rangeId, int ttlSeconds) {
        synchronized (this) {
            if (closed) {
                return FutureUtils.exception(new ObjectClosedException("MVCCStoreFactory"));
            }
        }

        log.info("Initializing stream({})/range({}) at storage container ({})",
            streamId, rangeId, scId);

        MVCCAsyncStore<byte[], byte[]> store = storeSupplier.get();

        File targetDir = chooseLocalStoreDir(streamId);
        // used for store ranges
        Path rangeStorePath = Paths.get(
            targetDir.getAbsolutePath(),
            "ranges",
            normalizedName(scId),
            normalizedName(streamId),
            normalizedName(rangeId));
        String storeName = String.format(
            "%s/%s/%s",
            normalizedName(scId),
            normalizedName(streamId),
            normalizedName(rangeId));

        if (null == checkpointStore) {
            checkpointStore = checkpointStoreSupplier.get();
        }

        // build a spec
        StateStoreSpec spec = StateStoreSpec.builder()
            .name(storeName)
            .keyCoder(ByteArrayCoder.of())
            .valCoder(ByteArrayCoder.of())
            .localStateStoreDir(rangeStorePath.toFile())
            .stream(streamName(scId, streamId, rangeId))
            .writeIOScheduler(chooseWriteIOExecutor(streamId))
            .readIOScheduler(chooseReadIOExecutor(streamId))
            .checkpointStore(checkpointStore)
            .checkpointDuration(Duration.ofMinutes(15))
            .checkpointIOScheduler(chooseCheckpointIOExecutor(streamId))
            .isReadonly(serveReadOnlyTable)
            .checkpointChecksumEnable(storageConf.getCheckpointChecksumEnable())
            .checkpointChecksumCompatible(storageConf.getCheckpointChecksumCompatible())
            .localStorageCleanupEnable(storageConf.getLocalStorageCleanupEnable())
            .checkpointRestoreIdleLimit(
                Duration.ofMillis(storageConf.getCheckpointRestoreIdleLimitMs()))
            .ttlSeconds(ttlSeconds)
            .build();


        return store.init(spec).whenComplete((ignored, throwable) -> {
            // since the store has not been added, so can't release its resources during close sc
            if (null != throwable) {
                log.info("Clearing resources hold by stream({})/range({}) at storage container ({}) ",
                    streamId, rangeId, scId);
                store.closeAsync().whenComplete((i, t) -> {
                    if (null != t) {
                        log.error("Clear resources hold by {} fail", store.name());
                    }
                });
            }
        }).thenApply(ignored -> {
            log.info("Successfully initialize stream({})/range({}) at storage container ({})",
                streamId, rangeId, scId);
            addStore(scId, streamId, rangeId, store);
            return store;
        });
    }

    @Override
    public CompletableFuture<Void> closeStores(long scId) {
        Map<RangeId, MVCCAsyncStore<byte[], byte[]>> scStores;
        synchronized (this) {
            scStores = stores.remove(scId);
        }
        if (null == scStores) {
            log.info("scStores for {} on store factory is null, return directly", scId);
            return FutureUtils.Void();
        }

        List<CompletableFuture<Void>> closeFutures = Lists.newArrayList();
        for (MVCCAsyncStore<byte[], byte[]> store : scStores.values()) {
            log.info("Closing {} of sc {}", store.name(), scId);
            closeFutures.add(store.closeAsync());
        }

        return FutureUtils.collect(closeFutures).thenApply(ignored -> null);
    }

    @Override
    public void close() {
        Map<Long, Map<RangeId, MVCCAsyncStore<byte[], byte[]>>> storesToClose;
        synchronized (this) {
            if (closed) {
                return;
            }
            storesToClose = Maps.newHashMap(stores);
            closed = true;
        }

        List<CompletableFuture<Void>> closeFutures = Lists.newArrayList();
        for (Map<RangeId, MVCCAsyncStore<byte[], byte[]>> scStores : storesToClose.values()) {
            for (MVCCAsyncStore<byte[], byte[]> store : scStores.values()) {
                closeFutures.add(store.closeAsync());
            }
        }
        try {
            FutureUtils.result(FutureUtils.collect(closeFutures));
            log.info("Successfully closed all the range stores opened by this range factory");
        } catch (Exception e) {
            log.info("Encountered issue on closing all the range stores opened by this range factory");
        }
        if (null != checkpointStore) {
            checkpointStore.close();
            checkpointStore = null;
        }

        SharedResourceManager.shared().release(
            storageResources.ioWriteScheduler(), writeIOScheduler);
        SharedResourceManager.shared().release(
            storageResources.ioReadScheduler(), readIOScheduler);
        SharedResourceManager.shared().release(
            storageResources.checkpointScheduler(), checkpointScheduler);
    }
}