OmKeylocationinfogroup.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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.hadoop.ozone.om.helpers;

import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.KeyLocationList;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * A list of key locations. This class represents one single version of the
 * blocks of a key.
 */
public class OmKeyLocationInfoGroup {
  private final long version;
  // TODO: HDDS-5472 Store one version of locationInfo for each
  //   OmKeyLocationInfoGroup
  private final Map<Long, List<OmKeyLocationInfo>> locationVersionMap;
  private  boolean isMultipartKey;

  public OmKeyLocationInfoGroup(Builder objBuilder) {
    this.version = objBuilder.version;
    if (objBuilder.mapLocations != null) {
      this.locationVersionMap = objBuilder.mapLocations;
    } else {
      locationVersionMap = new HashMap<>();
      for (OmKeyLocationInfo info : objBuilder.listLocations) {
        locationVersionMap
            .computeIfAbsent(info.getCreateVersion(), v -> new ArrayList<>())
            .add(info);
      }
    }
    //prevent NPE
    this.locationVersionMap.putIfAbsent(objBuilder.version, new ArrayList<>());
    this.isMultipartKey = objBuilder.isMultipartKey;
  }

  /**
   *  Builder Class for OmKeyLocationInfoGroup.
   */
  public static class Builder {
    private Map<Long, List<OmKeyLocationInfo>> mapLocations;
    private boolean isMultipartKey = false;
    private List<OmKeyLocationInfo> listLocations;
    private long version;

    public Builder setMapLocations(
        Map<Long, List<OmKeyLocationInfo>> mapLocations) {
      this.mapLocations = mapLocations;
      return this;
    }

    public Builder setListLocations(List<OmKeyLocationInfo> listLocations) {
      this.listLocations = listLocations;
      return this;
    }

    public Builder setIsMultipartKey(boolean isMultipartKey) {
      if (isMultipartKey) {
        this.isMultipartKey = true;
      }
      return this;
    }

    public Builder setVersion(long version) {
      this.version = version;
      return this;
    }

    //Return the final constructed builder object
    public OmKeyLocationInfoGroup build() {
      return new OmKeyLocationInfoGroup(this);
    }
  }

  public void setMultipartKey(boolean isMpu) {
    this.isMultipartKey = isMpu;
  }

  public boolean isMultipartKey() {
    return isMultipartKey;
  }

  /**
   * Return only the blocks that are created in the most recent version.
   *
   * @return the list of blocks that are created in the latest version.
   */
  public List<OmKeyLocationInfo> getBlocksLatestVersionOnly() {
    return new ArrayList<>(locationVersionMap.get(version));
  }

  public long getVersion() {
    return version;
  }

  /**
   * Use this expensive method only when absolutely needed!
   * It creates a new list so it is not an O(1) operation.
   * Use getLocationLists() instead.
   * @return a list of OmKeyLocationInfo
   */
  public List<OmKeyLocationInfo> getLocationList() {
    return locationVersionMap.values().stream().flatMap(List::stream)
        .collect(Collectors.toList());
  }

  public Collection<List<OmKeyLocationInfo>> getLocationLists() {
    return locationVersionMap.values();
  }

  public long getLocationListCount() {
    return locationVersionMap.values().stream().mapToLong(List::size).sum();
  }

  @Deprecated
  public List<OmKeyLocationInfo> getLocationList(Long versionToFetch) {
    return new ArrayList<>(locationVersionMap.get(versionToFetch));
  }

  public KeyLocationList getProtobuf(boolean ignorePipeline,
      int clientVersion) {
    KeyLocationList.Builder builder = KeyLocationList.newBuilder()
        .setVersion(version).setIsMultipartKey(isMultipartKey);
    List<OzoneManagerProtocolProtos.KeyLocation> keyLocationList =
        new ArrayList<>();
    for (List<OmKeyLocationInfo> locationList : locationVersionMap.values()) {
      for (OmKeyLocationInfo keyInfo : locationList) {
        keyLocationList.add(keyInfo.getProtobuf(ignorePipeline, clientVersion));
      }
    }
    return  builder.addAllKeyLocations(keyLocationList).build();
  }

  public static OmKeyLocationInfoGroup getFromProtobuf(
      KeyLocationList keyLocationList) {
    return new OmKeyLocationInfoGroup.Builder()
        .setVersion(keyLocationList.getVersion())
        .setMapLocations(keyLocationList.getKeyLocationsList().stream()
            .map(OmKeyLocationInfo::getFromProtobuf)
            .collect(Collectors.groupingBy(
                OmKeyLocationInfo::getCreateVersion)))
        .setIsMultipartKey(keyLocationList.getIsMultipartKey())
        .build();
  }

  /**
   * Given a new block location, generate a new version list based upon this
   * one.
   *
   * @param newLocationList a list of new location to be added.
   * @return newly generated OmKeyLocationInfoGroup
   */
  OmKeyLocationInfoGroup generateNextVersion(
      List<OmKeyLocationInfo> newLocationList) {
    Map<Long, List<OmKeyLocationInfo>> newMap = new HashMap<>();
    newMap.put(version + 1, new ArrayList<>(newLocationList));
    return new OmKeyLocationInfoGroup.Builder()
        .setVersion(version+1)
        .setMapLocations(newMap).build();
  }

  void appendNewBlocks(List<OmKeyLocationInfo> newLocationList) {
    List<OmKeyLocationInfo> locationList = locationVersionMap.get(version);
    for (OmKeyLocationInfo info : newLocationList) {
      info.setCreateVersion(version);
      locationList.add(info);
    }
  }

  void removeBlocks(long versionToRemove) {
    locationVersionMap.remove(versionToRemove);
  }

  void addAll(long versionToAdd, List<OmKeyLocationInfo> locationInfoList) {
    locationVersionMap.putIfAbsent(versionToAdd, new ArrayList<>());
    List<OmKeyLocationInfo> list = locationVersionMap.get(versionToAdd);
    list.addAll(locationInfoList);
  }

  @Override
  public String toString() {
    StringBuilder sb = new StringBuilder();
    sb.append("version:").append(version).append(" ");
    sb.append("isMultipartKey:").append(isMultipartKey);
    for (List<OmKeyLocationInfo> kliList : locationVersionMap.values()) {
      for (OmKeyLocationInfo kli: kliList) {
        sb.append(kli.getLocalID()).append(" || ");
      }
    }
    return sb.toString();
  }
}