content: /*
   * 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.metadata.stream;
  
  import static com.google.common.base.Preconditions.checkArgument;
  import static com.google.common.base.Preconditions.checkState;
  import static org.apache.bookkeeper.stream.protocol.ProtocolConstants.MIN_DATA_RANGE_ID;
  import static org.apache.bookkeeper.stream.protocol.util.ProtoUtils.isStreamCreated;
  import static org.apache.bookkeeper.stream.protocol.util.ProtoUtils.split;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Maps;
  import com.google.protobuf.InvalidProtocolBufferException;
  import java.util.List;
  import java.util.NavigableMap;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutorService;
  import java.util.function.Consumer;
  import java.util.function.Supplier;
  import javax.annotation.concurrent.GuardedBy;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.api.kv.op.CompareResult;
  import org.apache.bookkeeper.api.kv.op.Op;
  import org.apache.bookkeeper.api.kv.op.TxnOp;
  import org.apache.bookkeeper.api.kv.result.KeyValue;
  import org.apache.bookkeeper.api.kv.result.Result;
  import org.apache.bookkeeper.common.concurrent.FutureUtils;
  import org.apache.bookkeeper.common.util.Bytes;
  import org.apache.bookkeeper.statelib.api.mvcc.MVCCAsyncStore;
  import org.apache.bookkeeper.stream.proto.RangeMetadata;
  import org.apache.bookkeeper.stream.proto.RangeProperties;
  import org.apache.bookkeeper.stream.proto.RangeState;
  import org.apache.bookkeeper.stream.proto.StreamConfiguration;
  import org.apache.bookkeeper.stream.proto.StreamMetadata;
  import org.apache.bookkeeper.stream.proto.StreamMetadata.LifecycleState;
  import org.apache.bookkeeper.stream.proto.StreamMetadata.ServingState;
  import org.apache.bookkeeper.stream.proto.StreamProperties;
  import org.apache.bookkeeper.stream.protocol.util.StorageContainerPlacementPolicy;
  import org.apache.bookkeeper.stream.storage.api.metadata.stream.MetaRange;
  import org.apache.bookkeeper.stream.storage.exceptions.DataRangeNotFoundException;
  
  /**
   * The default implementation of {@link MetaRange}.
   */
  @Slf4j
  public class MetaRangeImpl implements MetaRange {
  
      private static final byte METADATA_SEP = (byte) 0x1;
      private static final byte RANGE_SEP = (byte) 0x2;
      private static final byte END_SEP = (byte) 0xff;
  
      public static final byte[] getStreamMetadataKey(long streamId) {
          byte[] metadataKey = new byte[Long.BYTES + 1];
          Bytes.toBytes(streamId, metadataKey, 0);
          metadataKey[Long.BYTES] = METADATA_SEP;
          return metadataKey;
      }
  
      public static final byte[] getStreamRangeKey(long streamId, long rangeId) {
          byte[] rangeKey = new byte[2 * Long.BYTES + 1];
          Bytes.toBytes(streamId, rangeKey, 0);
          rangeKey[Long.BYTES] = RANGE_SEP;
          Bytes.toBytes(rangeId, rangeKey, Long.BYTES + 1);
          return rangeKey;
      }
  
      public static final byte[] getStreamMetadataEndKey(long streamId) {
          byte[] metadataKey = new byte[Long.BYTES + 1];
          Bytes.toBytes(streamId, metadataKey, 0);
          metadataKey[Long.BYTES] = END_SEP;
          return metadataKey;
      }
  
      public static final boolean isMetadataKey(byte[] key) {
          return key.length == (Long.BYTES + 1) && key[Long.BYTES] == METADATA_SEP;
      }
  
      static final boolean isStreamRangeKey(byte[] key) {
          return key.length == (2 * Long.BYTES + 1) && key[Long.BYTES] == RANGE_SEP;
      }
  
      private final MVCCAsyncStore<byte[], byte[]> store;
      private final ExecutorService executor;
      private final StorageContainerPlacementPolicy placementPolicy;
  
      @GuardedBy("this")
      private long streamId;
      @GuardedBy("this")
      private StreamProperties streamProps;
      private LifecycleState lifecycleState = LifecycleState.UNINIT;
      private ServingState servingState = ServingState.WRITABLE;
      private long cTime = 0L;
      private long mTime = 0L;
  
      private long nextRangeId = MIN_DATA_RANGE_ID;
      /**
       * All the ranges in the stream.
       *
       * <p>The ranges are sorted by range id. It implies that these
       * ranges are also sorted in the increasing order of their creation time.
       */
      private final NavigableMap<Long, RangeMetadata> ranges;
      /**
       * List of ranges that are currently active in the stream.
       */
      private final List<Long> currentRanges;
  
      // the revision of this meta range in local store.
      private long revision;
  
      public MetaRangeImpl(MVCCAsyncStore<byte[], byte[]> store,
                           ExecutorService executor,
                           StorageContainerPlacementPolicy placementPolicy) {
          this(
              store,
              executor,
              placementPolicy,
              Maps.newTreeMap(),
              Lists.newArrayList(),
              StreamMetadata
                  .newBuilder()
                  .setLifecycleState(LifecycleState.UNINIT)
                  .setServingState(ServingState.WRITABLE)
                  .setNextRangeId(MIN_DATA_RANGE_ID)
                  .build(),
              0L,
              0L);
      }
  
      private MetaRangeImpl(MVCCAsyncStore<byte[], byte[]> store,
                            ExecutorService executor,
                            StorageContainerPlacementPolicy placementPolicy,
                            NavigableMap<Long, RangeMetadata> ranges,
                            List<Long> currentRanges,
                            StreamMetadata meta,
                            long cTime,
                            long mTime) {
          this.store = store;
          this.executor = executor;
          this.placementPolicy = placementPolicy;
          // construct the state
          this.ranges = ranges;
          this.currentRanges = currentRanges;
          this.cTime = cTime;
          this.mTime = mTime;
          this.streamProps = meta.getProps();
          this.streamId = streamProps.getStreamId();
          this.lifecycleState = meta.getLifecycleState();
          this.servingState = meta.getServingState();
          this.nextRangeId = meta.getNextRangeId();
      }
  
      @VisibleForTesting
      public long unsafeGetCreationTime() {
          return cTime;
      }
  
      @VisibleForTesting
      public long unsafeGetModificationTime() {
          return mTime;
      }
  
      @VisibleForTesting
      public LifecycleState unsafeGetLifecycleState() {
          return lifecycleState;
      }
  
      @VisibleForTesting
      public synchronized StreamProperties unsafeGetStreamProperties() {
          return streamProps;
      }
  
      @VisibleForTesting
      public synchronized long unsafeGetStreamId() {
          return streamId;
      }
  
      @VisibleForTesting
      public NavigableMap<Long, RangeMetadata> unsafeGetRanges() {
          return ranges;
      }
  
      @VisibleForTesting
      private synchronized StreamMetadata toStreamMetadata(LifecycleState state) {
          StreamMetadata.Builder builder = StreamMetadata
              .newBuilder()
              .setProps(streamProps)
              .setLifecycleState(state)
              .setServingState(servingState)
              .setNextRangeId(nextRangeId)
              .setCTime(cTime)
              .setMTime(mTime)
              .addAllCurrentRanges(currentRanges);
          return builder.build();
      }
  
      private synchronized StreamMetadata toStreamMetadata(ServingState state, long mTime) {
          StreamMetadata.Builder builder = StreamMetadata
              .newBuilder()
              .setProps(streamProps)
              .setLifecycleState(lifecycleState)
              .setServingState(state)
              .setNextRangeId(nextRangeId)
              .setCTime(cTime)
              .setMTime(mTime)
              .addAllCurrentRanges(currentRanges);
          return builder.build();
      }
  
      @VisibleForTesting
      public List<Long> unsafeGetCurrentRanges() {
          return currentRanges;
      }
  
      private <T> CompletableFuture<T> checkStreamCreated(Supplier<CompletableFuture<T>> supplier) {
          CompletableFuture<T> future = FutureUtils.createFuture();
          executor.submit(() -> {
              try {
                  if (isStreamCreated(lifecycleState)) {
                      supplier.get()
                          .thenApplyAsync(value -> future.complete(value), executor)
                          .exceptionally(cause -> future.completeExceptionally(cause));
                  } else {
                      throw new IllegalStateException("Stream isn't created yet.");
                  }
              } catch (Throwable cause) {
                  future.completeExceptionally(cause);
              }
          });
          return future;
      }
  
      private void checkLifecycleState(LifecycleState expected) {
          checkState(expected == lifecycleState, "Unexpected state " + lifecycleState + ", expected to be " + expected);
      }
  
      @Override
      public synchronized String getName() {
          checkState(null != streamProps);
          return streamProps.getStreamName();
      }
  
      private <T> CompletableFuture<T> executeTask(Consumer<CompletableFuture<T>> consumer) {
          CompletableFuture<T> executeFuture = FutureUtils.createFuture();
          executor.submit(() -> {
              try {
                  consumer.accept(executeFuture);
              } catch (Throwable cause) {
                  executeFuture.completeExceptionally(cause);
              }
          });
          return executeFuture;
      }
  
      @Override
      public CompletableFuture<Boolean> create(StreamProperties streamProps) {
          return executeTask((future) -> unsafeCreate(future, streamProps));
      }
  
      private void unsafeCreate(CompletableFuture<Boolean> createFuture,
                                StreamProperties streamProps) {
          // 1. verify the state
          checkLifecycleState(LifecycleState.UNINIT);
          this.lifecycleState = LifecycleState.CREATING;
  
          // 2. store the props/configuration
          synchronized (this) {
              this.streamProps = streamProps;
          }
          this.streamId = streamProps.getStreamId();
          this.cTime = this.mTime = System.currentTimeMillis();
  
          // 3. create the ranges
          List<RangeProperties> propertiesList =
              split(streamId, streamProps.getStreamConf().getInitialNumRanges(), nextRangeId, placementPolicy);
          List<Op<byte[], byte[]>> successOps = Lists.newArrayListWithExpectedSize(
              propertiesList.size() + 1);
          for (RangeProperties props : propertiesList) {
              RangeMetadata meta = RangeMetadata.newBuilder()
                  .setProps(props)
                  .setCreateTime(cTime)
                  .setFenceTime(Long.MAX_VALUE)
                  .setState(RangeState.RANGE_ACTIVE)
                  .addAllParents(Lists.newArrayList())
                  .build();
              ranges.put(props.getRangeId(), meta);
              currentRanges.add(props.getRangeId());
  
              successOps.add(
                  store.newPut(
                      getStreamRangeKey(streamId, props.getRangeId()),
                      meta.toByteArray()));
          }
          nextRangeId += propertiesList.size();
  
          // serialize the stream metadata
          byte[] streamMetadataKey = getStreamMetadataKey(streamId);
          successOps.add(
              store.newPut(
                  streamMetadataKey,
                  toStreamMetadata(LifecycleState.CREATED).toByteArray()));
  
          TxnOp<byte[], byte[]> txn = store.newTxn()
              // create stream only when stream doesn't exists
              .If(store.newCompareValue(CompareResult.EQUAL, streamMetadataKey, null))
              .Then(successOps.toArray(new Op[successOps.size()]))
              .build();
  
          if (log.isTraceEnabled()) {
              log.trace("Execute create stream metadata range txn {}", streamProps);
          }
          store.txn(txn)
              .thenApplyAsync(txnResult -> {
                  try {
                      if (log.isTraceEnabled()) {
                          log.trace("Create stream metadata range txn result = {}", txnResult.isSuccess());
                      }
                      if (txnResult.isSuccess()) {
                          List<Result<byte[], byte[]>> results = txnResult.results();
                          MetaRangeImpl.this.revision = results.get(results.size() - 1).revision();
  
                          // mark the state to CREATED
                          this.lifecycleState = LifecycleState.CREATED;
                          createFuture.complete(true);
                      } else {
                          createFuture.complete(false);
                      }
                      return null;
                  } finally {
                      txnResult.close();
                  }
              }, executor)
              .exceptionally(cause -> {
                  createFuture.completeExceptionally(cause);
                  return null;
              });
      }
  
      @Override
      public CompletableFuture<MetaRange> load(long streamId) {
          byte[] streamMetadataKey = getStreamMetadataKey(streamId);
          byte[] streamMetadataEndKey = getStreamMetadataEndKey(streamId);
  
          return store.range(streamMetadataKey, streamMetadataEndKey)
              .thenApplyAsync(kvs -> {
                  if (kvs.isEmpty()) {
                      return null;
                  } else {
                      loadMetadata(kvs);
                      return MetaRangeImpl.this;
                  }
              }, executor);
      }
  
      private void loadMetadata(List<KeyValue<byte[], byte[]>> kvs) {
          for (KeyValue<byte[], byte[]> kv : kvs) {
              if (isMetadataKey(kv.key())) {
                  this.revision = kv.modifiedRevision();
                  long streamId = Bytes.toLong(kv.key(), 0);
                  loadStreamMetadata(streamId, kv.value());
              } else if (isStreamRangeKey(kv.key())) {
                  long streamId = Bytes.toLong(kv.key(), 0);
                  long rangeId = Bytes.toLong(kv.key(), Long.BYTES + 1);
                  loadRangeMetadata(streamId, rangeId, kv.value());
              }
          }
      }
  
      private synchronized void loadStreamMetadata(long streamId, byte[] streamMetadataBytes) {
          this.streamId = streamId;
          StreamMetadata metadata;
          try {
              metadata = StreamMetadata.parseFrom(streamMetadataBytes);
          } catch (InvalidProtocolBufferException e) {
              throw new RuntimeException("Invalid stream metadata of stream " + streamId, e);
          }
  
          this.streamProps = metadata.getProps();
          this.lifecycleState = metadata.getLifecycleState();
          this.servingState = metadata.getServingState();
          this.currentRanges.clear();
          this.currentRanges.addAll(metadata.getCurrentRangesList());
          this.nextRangeId = metadata.getNextRangeId();
          this.cTime = metadata.getCTime();
          this.mTime = metadata.getMTime();
      }
  
      private void loadRangeMetadata(long streamId, long rangeId, byte[] rangeMetadataBytes) {
          checkArgument(this.streamId == streamId);
          checkArgument(rangeId >= 0L);
  
          RangeMetadata metadata;
          try {
              metadata = RangeMetadata.parseFrom(rangeMetadataBytes);
          } catch (InvalidProtocolBufferException e) {
              throw new RuntimeException("Invalid range metadata of range (" + streamId + ", " + rangeId + ")", e);
          }
  
          this.ranges.put(rangeId, metadata);
      }
  
      @Override
      public CompletableFuture<Boolean> delete(long streamId) {
          byte[] streamMetadataKey = getStreamMetadataKey(streamId);
          byte[] streamEndKey = getStreamMetadataEndKey(streamId);
  
          return store.deleteRange(streamMetadataKey, streamEndKey)
              .thenApplyAsync(kvs -> !kvs.isEmpty(), executor);
      }
  
      @Override
      public CompletableFuture<ServingState> getServingState() {
          return checkStreamCreated(() -> FutureUtils.value(servingState));
      }
  
      @Override
      public CompletableFuture<ServingState> updateServingState(ServingState state) {
          return checkStreamCreated(() -> {
              long mTime = System.currentTimeMillis();
              byte[] streamMetadataKey = getStreamMetadataKey(streamId);
              byte[] streamMetadata = toStreamMetadata(state, mTime).toByteArray();
              return store.rPut(streamMetadataKey, streamMetadata, revision)
                  .thenApplyAsync(newRev -> {
                      this.servingState = state;
                      this.mTime = mTime;
                      this.revision = newRev;
                      return state;
                  }, executor);
          });
      }
  
      @Override
      public CompletableFuture<StreamConfiguration> getConfiguration() {
          return checkStreamCreated(() -> FutureUtils.value(unsafeGetStreamProperties().getStreamConf()));
      }
  
      private RangeMetadata unsafeGetDataRange(long rangeId) throws DataRangeNotFoundException {
          RangeMetadata rangeMeta = ranges.get(rangeId);
          if (null == rangeMeta) {
              throw new DataRangeNotFoundException(streamId, rangeId);
          }
          return rangeMeta;
      }
  
      @Override
      public CompletableFuture<List<RangeMetadata>> getActiveRanges() {
          return checkStreamCreated(() -> {
              List<Long> rangesIds = ImmutableList.copyOf(currentRanges);
              List<RangeMetadata> properties = Lists.newArrayListWithExpectedSize(rangesIds.size());
              for (long rangeId : rangesIds) {
                  properties.add(unsafeGetDataRange(rangeId));
              }
              return FutureUtils.value(properties);
          });
      }
  
  }
  