alluxio-master-2.x/microbench/src/main/java/alluxio/inode/RocksBenchBase.java
/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.inode;

import alluxio.AlluxioTestDirectory;
import alluxio.collections.Pair;
import alluxio.conf.Configuration;
import alluxio.master.file.contexts.CreateDirectoryContext;
import alluxio.master.file.contexts.CreateFileContext;
import alluxio.master.file.meta.Inode;
import alluxio.master.file.meta.MutableInode;
import alluxio.master.file.meta.MutableInodeDirectory;
import alluxio.master.file.meta.MutableInodeFile;
import alluxio.master.metastore.rocks.RocksInodeStore;
import alluxio.proto.meta.InodeMeta;

import com.google.common.primitives.Longs;
import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;

public class RocksBenchBase {

  static final String SER_READ = "serRead";
  static final String NO_SER_READ = "noSerRead";
  static final String SER_NO_ALLOC_READ = "serNoAllocRead";
  static final String NO_SER_NO_ALLOC_READ = "noSerNoAllocRead";

  private final RocksInodeStore mRocksInodeStore;
  // RocksDB resources managed by the RocksInodeStore, no need to close manually
  private final RocksDB mDB;
  private final AtomicReference<ColumnFamilyHandle> mInodesColumn;
  // Created and managed in this class
  private final WriteOptions mDisableWAL;

  RocksBenchBase(String confType) throws IOException {
    Logger.getRootLogger().setLevel(Level.ERROR);
    String dir =
        AlluxioTestDirectory.createTemporaryDirectory("inode-store-bench").getAbsolutePath();
    RocksBenchConfig.setRocksConfig(confType, dir, Configuration.modifiableGlobal());
    mRocksInodeStore = new RocksInodeStore(dir);
    Pair<RocksDB, AtomicReference<ColumnFamilyHandle>> dbInfo = mRocksInodeStore.getDBInodeColumn();
    mDB = dbInfo.getFirst();
    mInodesColumn = dbInfo.getSecond();
    mDisableWAL = new WriteOptions().setDisableWAL(true);
  }

  static MutableInode<?> genInode(boolean mIsDirectory) {
    if (!mIsDirectory) {
      return MutableInodeFile.create(0, 1, "testFile", System.currentTimeMillis(),
          CreateFileContext.defaults());
    } else {
      return MutableInodeDirectory.create(0, -1, "someParent",
          CreateDirectoryContext.defaults());
    }
  }

  void after() {
    mRocksInodeStore.clear();
    mRocksInodeStore.close();
    mDisableWAL.close();
  }

  long getInodeReadId(long total, long nxt, long min, int threadCount, int threadId) {
    // since the larger ids were written last (and will be in the memtable), but the zipfan will
    // generate small numbers the most often, we subtract from the total count to read the
    // large ids most often
    // this assumes a workload where latest writes are read more often
    long id = total - (nxt - min);
    return (id * threadCount) + threadId;
  }

  Inode readInode(long totalCount, long nxtId, long min, int threadCount, int threadId) {
    long readId = getInodeReadId(totalCount, nxtId, min, threadCount, threadId);
    try {
      return mRocksInodeStore.get(readId).get();
    } catch (Exception e) {
      System.err.printf("Error with element %d, %s", readId, e);
      throw e;
    }
  }

  Inode readInode(long totalCount, long nxtId, long min, int threadCount, int threadId,
                  byte[] inodeBytes) {
    return get(getInodeReadId(totalCount, nxtId, min, threadCount, threadId),
        inodeBytes).map(Inode::wrap).get();
  }

  byte[] readInodeBytes(long totalCount, long nxtId, long min, int threadCount, int threadId) {
    return getBytes(getInodeReadId(totalCount, nxtId, min, threadCount, threadId)).get();
  }

  int readInodeBytes(long totalCount, long nxtId, long min, int threadCount, int threadId,
                     byte[] inodeBytes) {
    return getBytes(getInodeReadId(totalCount, nxtId, min, threadCount, threadId), inodeBytes);
  }

  void writeInode(long nxtId, int threadCount, int threadId, MutableInode<?> inode) {
    // we write the bytes directly instead of creating a new inode with the new id
    // since we just want to test the cost of storing the inode and not allocating new ones
    writeBytes(nxtId, threadCount, threadId, inode.toProto().toByteArray());
  }

  long getWriteId(long nxtId, int threadCount, int threadId) {
    return (nxtId * threadCount) + threadId;
  }

  void remove(long nxtId, int threadCount, int threadId) {
    long nodeId = getWriteId(nxtId, threadCount, threadId);
    try {
      mDB.delete(mInodesColumn.get(), mDisableWAL, Longs.toByteArray(nodeId));
    } catch (RocksDBException e) {
      throw new RuntimeException(e);
    }
  }

  void writeBytes(long nxtId, int threadCount, int threadId, byte[] inodeBytes) {
    long nodeId = getWriteId(nxtId, threadCount, threadId);
    try {
      mDB.put(mInodesColumn.get(), mDisableWAL, Longs.toByteArray(nodeId),
          inodeBytes);
    } catch (RocksDBException e) {
      throw new RuntimeException(e);
    }
  }

  public int getBytes(long id, byte[] inodeBytes) {
    try {
      return mDB.get(mInodesColumn.get(), Longs.toByteArray(id), inodeBytes);
    } catch (RocksDBException e) {
      throw new RuntimeException(e);
    }
  }

  public Optional<byte[]> getBytes(long id) {
    byte[] inodeBytes;
    try {
      inodeBytes = mDB.get(mInodesColumn.get(), Longs.toByteArray(id));
    } catch (RocksDBException e) {
      throw new RuntimeException(e);
    }
    if (inodeBytes == null) {
      return Optional.empty();
    }
    return Optional.of(inodeBytes);
  }

  public Optional<MutableInode<?>> get(long id, byte[] inode) {
    int count;
    try {
      count = mDB.get(mInodesColumn.get(), Longs.toByteArray(id), inode);
    } catch (RocksDBException e) {
      throw new RuntimeException(e);
    }
    if (count == RocksDB.NOT_FOUND) {
      return Optional.empty();
    }
    if (count > inode.length) {
      throw new RuntimeException("Byte array too small to read inode");
    }
    try {
      return Optional.of(MutableInode.fromProto(
          InodeMeta.Inode.parseFrom(ByteBuffer.wrap(inode, 0, count))));
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }
}
