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.beam.sdk.transforms.resourcehints;
  
  import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
  
  import java.nio.charset.StandardCharsets;
  import java.util.List;
  import java.util.Map;
  import java.util.function.Function;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  import org.apache.beam.model.pipeline.v1.RunnerApi;
  import org.apache.beam.model.pipeline.v1.RunnerApi.StandardResourceHints;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ProtocolMessageEnum;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Pipeline authors can use resource hints to provide additional information to runners about the
   * desired aspects of the execution environment. Resource hints can be specified via {@link
   * org.apache.beam.sdk.transforms.PTransform PTransform#setResourceHints} for parts of the pipeline,
   * or globally via {@link ResourceHintsOptions resourceHints} pipeline option.
   *
   * <p>Interpretation of hints is provided by Beam runners.
   */
  public class ResourceHints {
    private static final Logger LOG = LoggerFactory.getLogger(ResourceHints.class);
    private static final String MIN_RAM_URN = "beam:resources:min_ram_bytes:v1";
    private static final String ACCELERATOR_URN = "beam:resources:accelerator:v1";
    private static final String CPU_COUNT_URN = "beam:resources:cpu_count:v1";
    private static final String MAX_ACTIVE_BUNDLES_PER_WORKER =
        "beam:resources:max_active_bundles_per_worker:v1";
  
    // TODO: reference this from a common location in all packages that use this.
    private static String getUrn(ProtocolMessageEnum value) {
      return value.getValueDescriptor().getOptions().getExtension(RunnerApi.beamUrn);
    }
  
    static {
      checkState(MIN_RAM_URN.equals(getUrn(StandardResourceHints.Enum.MIN_RAM_BYTES)));
      checkState(ACCELERATOR_URN.equals(getUrn(StandardResourceHints.Enum.ACCELERATOR)));
      checkState(CPU_COUNT_URN.equals(getUrn(StandardResourceHints.Enum.CPU_COUNT)));
      checkState(
          MAX_ACTIVE_BUNDLES_PER_WORKER.equals(
              (getUrn(StandardResourceHints.Enum.MAX_ACTIVE_BUNDLES_PER_WORKER))));
    }
  
    private static ImmutableMap<String, String> hintNameToUrn =
        ImmutableMap.<String, String>builder()
            .put("minRam", MIN_RAM_URN)
            .put("min_ram", MIN_RAM_URN) // Courtesy alias.
            .put("accelerator", ACCELERATOR_URN)
            .put("cpuCount", CPU_COUNT_URN)
            .put("cpu_count", CPU_COUNT_URN) // Courtesy alias.
            .put("max_active_bundles_per_worker", MAX_ACTIVE_BUNDLES_PER_WORKER)
            .put("maxActiveBundlesPerWorker", MAX_ACTIVE_BUNDLES_PER_WORKER) // Courtesy alias.
            .put("max_active_bundle_per_worker", MAX_ACTIVE_BUNDLES_PER_WORKER) // Courtesy alias.
            .build();
  
    private static ImmutableMap<String, Function<String, ResourceHint>> parsers =
        ImmutableMap.<String, Function<String, ResourceHint>>builder()
            .put(MIN_RAM_URN, s -> new BytesHint(BytesHint.parse(s)))
            .put(ACCELERATOR_URN, s -> new StringHint(s))
            .put(CPU_COUNT_URN, s -> new IntHint(IntHint.parse(s)))
            .put(MAX_ACTIVE_BUNDLES_PER_WORKER, s -> new IntHint(IntHint.parse(s)))
            .build();
  
    private static final ResourceHints EMPTY = new ResourceHints(ImmutableMap.of());
  
    private final ImmutableMap<String, ResourceHint> hints;
  
    private ResourceHints(ImmutableMap<String, ResourceHint> hints) {
      this.hints = hints;
    }
  
    /** Creates a {@link ResourceHints} instance with no hints. */
    public static ResourceHints create() {
      return EMPTY;
    }
  
    /** Creates a {@link ResourceHints} instance with hints supplied in options. */
    public static ResourceHints fromOptions(PipelineOptions options) {
      ResourceHintsOptions resourceHintsOptions = options.as(ResourceHintsOptions.class);
      ResourceHints result = create();
      List<String> hints = resourceHintsOptions.getResourceHints();
      Splitter splitter = Splitter.on('=').limit(2);
      for (String hint : hints) {
        List<String> parts = splitter.splitToList(hint);
        if (parts.size() != 2) {
          throw new IllegalArgumentException("Unparsable resource hint: " + hint);
        }
        String nameOrUrn = parts.get(0);
        String stringValue = parts.get(1);
        String urn;
        if (hintNameToUrn.containsKey(nameOrUrn)) {
          urn = hintNameToUrn.get(nameOrUrn);
        } else if (!nameOrUrn.startsWith("beam:resources:")) {
          // Allow unknown hints to be passed, but validate a little bit to prevent typos.
          throw new IllegalArgumentException("Unknown resource hint: " + hint);
        } else {
          urn = nameOrUrn;
        }
        ResourceHint value =
            Preconditions.checkNotNull(parsers.getOrDefault(urn, StringHint::new)).apply(stringValue);
        result = result.withHint(urn, value);
      }
      return result;
    }
  
    /*package*/ static class BytesHint extends ResourceHint {
      private static Map<String, Long> suffixes =
          ImmutableMap.<String, Long>builder()
              .put("B", 1L)
              .put("KB", 1000L)
              .put("MB", 1000_000L)
              .put("GB", 1000_000_000L)
              .put("TB", 1000_000_000_000L)
              .put("PB", 1000_000_000_000_000L)
              .put("KiB", 1L << 10)
              .put("MiB", 1L << 20)
              .put("GiB", 1L << 30)
              .put("TiB", 1L << 40)
              .put("PiB", 1L << 50)
              .build();
  
      private final long value;
  
      @Override
      public boolean equals(@Nullable Object other) {
        if (other == null) {
          return false;
        } else if (this == other) {
          return true;
        } else if (other instanceof BytesHint) {
          return ((BytesHint) other).value == value;
        } else {
          return false;
        }
      }
  
      @Override
      public int hashCode() {
        return Long.hashCode(value);
      }
  
      public BytesHint(long value) {
        this.value = value;
      }
  
      public static long parse(String s) {
        Matcher m = Pattern.compile("([\\d.]+)[\\s]?([\\D]+$)").matcher(s);
        if (m.find()) {
          String number = m.group(1);
          String suffix = m.group(2);
          if (number != null && suffix != null && suffixes.containsKey(suffix)) {
            return (long) (Double.parseDouble(number) * suffixes.get(suffix));
          }
        }
        throw new IllegalArgumentException("Unable to parse '" + s + "' as a byte value.");
      }
  
      @Override
      public ResourceHint mergeWithOuter(ResourceHint outer, boolean isSum) {
        return isSum
            ? new BytesHint(value + ((BytesHint) outer).value)
            : new BytesHint(Math.max(value, ((BytesHint) outer).value));
      }
  
      @Override
      public byte[] toBytes() {
        return String.valueOf(value).getBytes(StandardCharsets.US_ASCII);
      }
    }
  
    /*package*/ static class StringHint extends ResourceHint {
      private final String value;
  
      public StringHint(String value) {
        this.value = value;
      }
  
      public static String parse(String s) {
        return s;
      }
  
      @Override
      public byte[] toBytes() {
        return value.getBytes(StandardCharsets.US_ASCII);
      }
  
      @Override
      public boolean equals(@Nullable Object other) {
        if (other == null) {
          return false;
        } else if (this == other) {
          return true;
        } else if (other instanceof StringHint) {
          return ((StringHint) other).value.equals(value);
        } else {
          return false;
        }
      }
  
      @Override
      public int hashCode() {
        return value.hashCode();
      }
    }
  
    /*package*/ static class IntHint extends ResourceHint {
      private final int value;
  
      @Override
      public boolean equals(@Nullable Object other) {
        if (other == null) {
          return false;
        } else if (this == other) {
          return true;
        } else if (other instanceof IntHint) {
          return ((IntHint) other).value == value;
        } else {
          return false;
        }
      }
  
      @Override
      public int hashCode() {
        return Integer.hashCode(value);
      }
  
      public IntHint(int value) {
        this.value = value;
      }
  
      public static int parse(String s) {
        return Integer.parseInt(s, 10);
      }
  
      @Override
      public ResourceHint mergeWithOuter(ResourceHint outer, boolean isSum) {
        return isSum
            ? new IntHint(value + ((IntHint) outer).value)
            : new IntHint(Math.max(value, ((IntHint) outer).value));
      }
  
      @Override
      public byte[] toBytes() {
        return String.valueOf(value).getBytes(StandardCharsets.US_ASCII);
      }
    }
  
    /**
     * Sets desired minimal available RAM size to have in transform's execution environment.
     *
     * @param ramBytes specifies a positive RAM size in bytes. A number greater than 2G
     *     (Integer.MAX_VALUE) is typical.
     */
    public ResourceHints withMinRam(long ramBytes) {
      if (ramBytes <= 0L) {
        // TODO(yathu) ignore invalid value as of Beam v2.47.0. throw error in future version.
        LOG.error(
            "Encountered invalid non-positive minimum ram hint value {}.\n"
                + "Likely cause is an (overflowing) int expression is passed in. "
                + "The value is ignored. In the future, The method will require an object Long type "
                + "and throw an IllegalArgumentException for invalid values.",
            ramBytes);
        return this;
      } else if (ramBytes <= Integer.MAX_VALUE) {
        LOG.warn(
            "Minimum available RAM size ({}) is set too small.\n"
                + "Likely cause is an (overflowing) int expression is passed in.",
            ramBytes);
      }
      return withHint(MIN_RAM_URN, new BytesHint(ramBytes));
    }
  
    /**
     * Sets desired minimal available RAM size to have in transform's execution environment.
     *
     * @param ramBytes specifies a human-friendly size string, for example: '10.5 GiB', '4096 MiB',
     *     etc.
     */
    public ResourceHints withMinRam(String ramBytes) {
      return withMinRam(BytesHint.parse(ramBytes));
    }
  
    /** Declares hardware accelerators that are desired to have in the execution environment. */
    public ResourceHints withAccelerator(String accelerator) {
      return withHint(ACCELERATOR_URN, new StringHint(accelerator));
    }
  
    /** Declares a custom resource hint that has a specified URN. */
    public ResourceHints withHint(String urn, ResourceHint hint) {
      ImmutableMap.Builder<String, ResourceHint> newHints = ImmutableMap.builder();
      newHints.put(urn, hint);
      for (Map.Entry<String, ResourceHint> oldHint : hints.entrySet()) {
        if (!oldHint.getKey().equals(urn)) {
          newHints.put(oldHint.getKey(), oldHint.getValue());
        }
      }
      return new ResourceHints(newHints.build());
    }
  
    /**
     * Sets desired minimal CPU or vCPU count to have in transform's execution environment.
     *
     * @param cpuCount specifies a positive CPU count.
     */
    public ResourceHints withCPUCount(int cpuCount) {
      if (cpuCount <= 0) {
        LOG.error(
            "Encountered invalid non-positive cpu count hint value {}.\n"
                + "The value is ignored. In the future, The method will require an object Long type "
                + "and throw an IllegalArgumentException for invalid values.",
            cpuCount);
        return this;
      }
      return withHint(CPU_COUNT_URN, new IntHint(cpuCount));
    }
  
    public ResourceHints withMaxActiveBundlesPerWorker(int maxActiveBundlesPerWorker) {
      if (maxActiveBundlesPerWorker <= 0) {
        LOG.error(
            "Encountered invalid non-positive max active bundles per worker hint value {}.\n"
                + "The value is ignored.",
            maxActiveBundlesPerWorker);
        return this;
      }
      return withHint(MAX_ACTIVE_BUNDLES_PER_WORKER, new IntHint(maxActiveBundlesPerWorker));
    }
  
    public Map<String, ResourceHint> hints() {
      return hints;
    }
  
    public ResourceHints mergeWithOuter(ResourceHints outer) {
      if (outer.hints.isEmpty()) {
        return this;
      } else if (hints.isEmpty()) {
        return outer;
      } else {
        ImmutableMap.Builder<String, ResourceHint> newHints = ImmutableMap.builder();
        for (Map.Entry<String, ResourceHint> outerHint : outer.hints().entrySet()) {
          String key = outerHint.getKey();
          if (hints.containsKey(key)) {
            newHints.put(
                key,
                hints
                    .get(key)
                    .mergeWithOuter(
                        outerHint.getValue(), /*isSum*/ key.equals(MAX_ACTIVE_BUNDLES_PER_WORKER)));
          } else {
            newHints.put(outerHint);
          }
        }
        for (Map.Entry<String, ResourceHint> hint : hints.entrySet()) {
          if (!outer.hints.containsKey(hint.getKey())) {
            newHints.put(hint);
          }
        }
        return new ResourceHints(newHints.build());
      }
    }
  
    @Override
    public boolean equals(@Nullable Object other) {
      if (other == null) {
        return false;
      } else if (this == other) {
        return true;
      } else if (other instanceof ResourceHints) {
        return ((ResourceHints) other).hints.equals(hints);
      } else {
        return false;
      }
    }
  
    @Override
    public int hashCode() {
      return hints.hashCode();
    }
  }
  