DefaultAllocationManagerOption.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
 *
 *    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.arrow.memory;

import java.lang.reflect.Field;
import org.apache.arrow.util.VisibleForTesting;
import org.checkerframework.checker.nullness.qual.Nullable;

/** A class for choosing the default allocation manager. */
public class DefaultAllocationManagerOption {

  /** The environmental variable to set the default allocation manager type. */
  public static final String ALLOCATION_MANAGER_TYPE_ENV_NAME = "ARROW_ALLOCATION_MANAGER_TYPE";

  /** The system property to set the default allocation manager type. */
  public static final String ALLOCATION_MANAGER_TYPE_PROPERTY_NAME =
      "arrow.allocation.manager.type";

  static final org.slf4j.Logger LOGGER =
      org.slf4j.LoggerFactory.getLogger(DefaultAllocationManagerOption.class);

  /** The default allocation manager factory. */
  private static AllocationManager.@Nullable Factory DEFAULT_ALLOCATION_MANAGER_FACTORY = null;

  /** The allocation manager type. */
  public enum AllocationManagerType {
    /** Netty based allocation manager. */
    Netty,

    /** Unsafe based allocation manager. */
    Unsafe,

    /** Unknown type. */
    Unknown,
  }

  /**
   * Returns the default allocation manager type.
   *
   * @return the default allocation manager type.
   */
  @SuppressWarnings("nullness:argument") // enum types valueOf are implicitly non-null
  @VisibleForTesting
  public static AllocationManagerType getDefaultAllocationManagerType() {
    AllocationManagerType ret = AllocationManagerType.Unknown;

    try {
      String envValue = System.getenv(ALLOCATION_MANAGER_TYPE_ENV_NAME);
      ret = AllocationManagerType.valueOf(envValue);
    } catch (IllegalArgumentException | NullPointerException e) {
      // ignore the exception, and make the allocation manager type remain unchanged
    }

    // system property takes precedence
    try {
      String propValue = System.getProperty(ALLOCATION_MANAGER_TYPE_PROPERTY_NAME);
      ret = AllocationManagerType.valueOf(propValue);
    } catch (IllegalArgumentException | NullPointerException e) {
      // ignore the exception, and make the allocation manager type remain unchanged
    }
    return ret;
  }

  static AllocationManager.Factory getDefaultAllocationManagerFactory() {
    if (DEFAULT_ALLOCATION_MANAGER_FACTORY != null) {
      return DEFAULT_ALLOCATION_MANAGER_FACTORY;
    }
    AllocationManagerType type = getDefaultAllocationManagerType();
    switch (type) {
      case Netty:
        DEFAULT_ALLOCATION_MANAGER_FACTORY = getNettyFactory();
        break;
      case Unsafe:
        DEFAULT_ALLOCATION_MANAGER_FACTORY = getUnsafeFactory();
        break;
      case Unknown:
        LOGGER.info("allocation manager type not specified, using netty as the default type");
        DEFAULT_ALLOCATION_MANAGER_FACTORY = getFactory(CheckAllocator.check());
        break;
      default:
        throw new IllegalStateException("Unknown allocation manager type: " + type);
    }
    return DEFAULT_ALLOCATION_MANAGER_FACTORY;
  }

  @SuppressWarnings({"nullness:argument", "nullness:return"})
  // incompatible argument for parameter obj of Field.get
  // Static member qualifying type may not be annotated
  private static AllocationManager.Factory getFactory(String clazzName) {
    try {
      Field field = Class.forName(clazzName).getDeclaredField("FACTORY");
      field.setAccessible(true);
      return (AllocationManager.Factory) field.get(null);
    } catch (Exception e) {
      throw new RuntimeException("Unable to instantiate Allocation Manager for " + clazzName, e);
    }
  }

  private static AllocationManager.Factory getUnsafeFactory() {
    try {
      return getFactory("org.apache.arrow.memory.unsafe.UnsafeAllocationManager");
    } catch (RuntimeException e) {
      throw new RuntimeException(
          "Please add arrow-memory-unsafe to your classpath,"
              + " No DefaultAllocationManager found to instantiate an UnsafeAllocationManager",
          e);
    }
  }

  private static AllocationManager.Factory getNettyFactory() {
    try {
      return getFactory("org.apache.arrow.memory.netty.NettyAllocationManager");
    } catch (RuntimeException e) {
      throw new RuntimeException(
          "Please add arrow-memory-netty to your classpath,"
              + " No DefaultAllocationManager found to instantiate an NettyAllocationManager",
          e);
    }
  }
}