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.celeborn.common.protocol;
  
  import java.io.Serializable;
  import java.util.*;
  
  public class StorageInfo implements Serializable {
    public enum Type {
      MEMORY(0),
      HDD(1),
      SSD(2),
      HDFS(3),
      OSS(4),
      S3(5);
  
      private final int value;
  
      Type(int value) {
        this.value = value;
      }
  
      public int getValue() {
        return value;
      }
    }
  
    public static final Map<Integer, Type> typesMap = new HashMap<>();
    public static final Set<String> typeNames = new HashSet<>();
    public static final Map<String, Type> types = new HashMap<>();
  
    static {
      for (Type type : Type.values()) {
        typesMap.put(type.value, type);
        typeNames.add(type.name());
        types.put(type.name(), type);
      }
    }
  
    public static final int MEMORY_MASK = 0b1;
    public static final int LOCAL_DISK_MASK = 0b10;
    public static final int HDFS_MASK = 0b100;
    public static final int OSS_MASK = 0b1000;
    public static final int S3_MASK = 0b10000;
    public static final int ALL_TYPES_AVAILABLE_MASK = 0;
  
    // Default storage Type is MEMORY.
    private Type type = Type.MEMORY;
    private String mountPoint = "";
    // if a file is committed, field "finalResult" will be true
    private boolean finalResult = false;
    private String filePath;
  
    public int availableStorageTypes = 0;
  
    public long fileSize;
  
    public List<Long> chunkOffsets;
  
    public StorageInfo() {}
  
    public StorageInfo(Type type, boolean isFinal, String filePath) {
      this.type = type;
      this.finalResult = isFinal;
      this.filePath = filePath;
    }
  
    public StorageInfo(String mountPoint, StorageInfo.Type type, int availableStorageTypes) {
      this.mountPoint = mountPoint;
      this.type = type;
      this.availableStorageTypes = availableStorageTypes;
    }
  
    public StorageInfo(
        Type type,
        String mountPoint,
        boolean finalResult,
        String filePath,
        int availableStorageTypes) {
      this.type = type;
      this.mountPoint = mountPoint;
      this.finalResult = finalResult;
      this.filePath = filePath;
      this.availableStorageTypes = availableStorageTypes;
    }
  
    public StorageInfo(
        Type type,
        String mountPoint,
        boolean finalResult,
        String filePath,
        int availableStorageTypes,
        long fileSize,
        List<Long> chunkOffsets) {
      this.type = type;
      this.mountPoint = mountPoint;
      this.finalResult = finalResult;
      this.filePath = filePath;
      this.availableStorageTypes = availableStorageTypes;
      this.fileSize = fileSize;
      this.chunkOffsets = chunkOffsets;
    }
  
    public boolean isFinalResult() {
      return finalResult;
    }
  
    public void setFinalResult(boolean finalResult) {
      this.finalResult = finalResult;
    }
  
    public String getMountPoint() {
      return mountPoint;
    }
  
    public void setMountPoint(String mountPoint) {
      this.mountPoint = mountPoint;
    }
  
    public Type getType() {
      return type;
    }
  
    public String getFilePath() {
      return filePath;
    }
  
    public void setChunkOffsets(List<Long> chunkOffsets) {
      this.chunkOffsets = chunkOffsets;
    }
  
    public List<Long> getChunkOffsets() {
      return this.chunkOffsets;
    }
  
    public void setFileSize(long fileSize) {
      this.fileSize = fileSize;
    }
  
    public long getFileSize() {
      return fileSize;
    }
  
    @Override
    public String toString() {
      return "StorageInfo{"
          + "type="
          + type
          + ", mountPoint='"
          + mountPoint
          + '\''
          + ", finalResult="
          + finalResult
          + ", filePath="
          + filePath
          + ", fileSize="
          + fileSize
          + ", chunkOffsets="
          + chunkOffsets
          + '}';
    }
  
    public boolean memoryAvailable() {
      return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
          || (availableStorageTypes & MEMORY_MASK) > 0;
    }
  
    public static boolean localDiskAvailable(int availableStorageTypes) {
      return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
          || (availableStorageTypes & LOCAL_DISK_MASK) > 0;
    }
  
    public boolean localDiskAvailable() {
      return StorageInfo.localDiskAvailable(availableStorageTypes);
    }
  
    public static boolean HDFSAvailable(int availableStorageTypes) {
      return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
          || (availableStorageTypes & HDFS_MASK) > 0;
    }
  
    public boolean HDFSAvailable() {
      return StorageInfo.HDFSAvailable(availableStorageTypes);
    }
  
    public static boolean HDFSOnly(int availableStorageTypes) {
      return availableStorageTypes == HDFS_MASK;
    }
  
    public static boolean S3Only(int availableStorageTypes) {
      return availableStorageTypes == S3_MASK;
    }
  
    public static boolean OSSOnly(int availableStorageTypes) {
      return availableStorageTypes == OSS_MASK;
    }
  
    public static boolean OSSAvailable(int availableStorageTypes) {
      return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
          || (availableStorageTypes & OSS_MASK) > 0;
    }
  
    public static boolean S3Available(int availableStorageTypes) {
      return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
          || (availableStorageTypes & S3_MASK) > 0;
    }
  
    public boolean OSSAvailable() {
      return StorageInfo.OSSAvailable(availableStorageTypes);
    }
  
    public boolean S3Available() {
      return StorageInfo.S3Available(availableStorageTypes);
    }
  
    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;
      StorageInfo that = (StorageInfo) o;
      return finalResult == that.finalResult
          && availableStorageTypes == that.availableStorageTypes
          && type == that.type
          && Objects.equals(mountPoint, that.mountPoint)
          && Objects.equals(filePath, that.filePath);
    }
  
    @Override
    public int hashCode() {
      return Objects.hash(type, mountPoint, finalResult, filePath, availableStorageTypes);
    }
  
    public static boolean validate(String typeStr) {
      return typeNames.contains(typeStr);
    }
  
    public static PbStorageInfo toPb(StorageInfo storageInfo) {
      String filePath = storageInfo.getFilePath();
      PbStorageInfo.Builder builder = PbStorageInfo.newBuilder();
      builder
          .setType(storageInfo.type.value)
          .setFinalResult(storageInfo.finalResult)
          .setMountPoint(storageInfo.mountPoint)
          .setAvailableStorageTypes(storageInfo.availableStorageTypes)
          .setFileSize(storageInfo.getFileSize());
      if (storageInfo.getChunkOffsets() != null) {
        builder.addAllChunkOffsets(storageInfo.getChunkOffsets());
      }
      if (filePath != null) {
        builder.setFilePath(filePath);
      }
      return builder.build();
    }
  
    public static StorageInfo fromPb(PbStorageInfo pbStorageInfo) {
      return new StorageInfo(
          typesMap.get(pbStorageInfo.getType()),
          pbStorageInfo.getMountPoint(),
          pbStorageInfo.getFinalResult(),
          pbStorageInfo.getFilePath(),
          pbStorageInfo.getAvailableStorageTypes(),
          pbStorageInfo.getFileSize(),
          pbStorageInfo.getChunkOffsetsList());
    }
  
    public static int getAvailableTypes(List<Type> types) {
      int ava = 0;
      for (Type type : types) {
        switch (type) {
          case MEMORY:
            ava = ava | MEMORY_MASK;
            break;
          case HDD:
          case SSD:
            ava = ava | LOCAL_DISK_MASK;
            break;
          case HDFS:
            ava = ava | HDFS_MASK;
            break;
          case OSS:
            ava = ava | OSS_MASK;
            break;
          case S3:
            ava = ava | S3_MASK;
            break;
        }
      }
      return ava;
    }
  
    public static Type fromStrToType(String typeStr) {
      return types.get(typeStr);
    }
  }
  