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.arrow.memory.util;
  
  import java.lang.reflect.Constructor;
  import java.lang.reflect.Field;
  import java.lang.reflect.InvocationTargetException;
  import java.nio.ByteBuffer;
  import java.nio.ByteOrder;
  import java.security.AccessController;
  import java.security.PrivilegedAction;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import sun.misc.Unsafe;
  
  /** Utilities for memory related operations. */
  public class MemoryUtil {
    private static final org.slf4j.Logger logger =
        org.slf4j.LoggerFactory.getLogger(MemoryUtil.class);
  
    private static final @Nullable Constructor<?> DIRECT_BUFFER_CONSTRUCTOR;
  
    /** The unsafe object from which to access the off-heap memory. */
    private static final Unsafe UNSAFE;
  
    /** The start offset of array data relative to the start address of the array object. */
    private static final long BYTE_ARRAY_BASE_OFFSET;
  
    /** The offset of the address field with the {@link java.nio.ByteBuffer} object. */
    private static final long BYTE_BUFFER_ADDRESS_OFFSET;
  
    /** If the native byte order is little-endian. */
    public static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN;
  
    // Java 1.8, 9, 11, 17, 21 becomes 1, 9, 11, 17, and 21.
    @SuppressWarnings("StringSplitter")
    private static final int majorVersion =
        Integer.parseInt(System.getProperty("java.specification.version").split("\\D+")[0]);
  
    static {
      try {
        // try to get the unsafe object
        final Object maybeUnsafe =
            AccessController.doPrivileged(
                new PrivilegedAction<Object>() {
                  @Override
                  @SuppressWarnings({"nullness:argument", "nullness:return"})
                  // incompatible argument for parameter obj of Field.get
                  // incompatible types in return
                  public Object run() {
                    try {
                      final Field unsafeField = Unsafe.class.getDeclaredField("theUnsafe");
                      unsafeField.setAccessible(true);
                      return unsafeField.get(null);
                    } catch (Throwable e) {
                      return e;
                    }
                  }
                });
  
        if (maybeUnsafe instanceof Throwable) {
          throw (Throwable) maybeUnsafe;
        }
  
        UNSAFE = (Unsafe) maybeUnsafe;
  
        // get the offset of the data inside a byte array object
        BYTE_ARRAY_BASE_OFFSET = UNSAFE.arrayBaseOffset(byte[].class);
  
        // get the offset of the address field in a java.nio.Buffer object
        Field addressField = java.nio.Buffer.class.getDeclaredField("address");
        addressField.setAccessible(true);
        BYTE_BUFFER_ADDRESS_OFFSET = UNSAFE.objectFieldOffset(addressField);
  
        Constructor<?> directBufferConstructor;
        long address = -1;
        final ByteBuffer direct = ByteBuffer.allocateDirect(1);
        try {
  
          final Object maybeDirectBufferConstructor =
              AccessController.doPrivileged(
                  new PrivilegedAction<Object>() {
                    @Override
                    public Object run() {
                      try {
                        final Constructor<?> constructor =
                            (majorVersion >= 21)
                                ? direct.getClass().getDeclaredConstructor(long.class, long.class)
                                : direct.getClass().getDeclaredConstructor(long.class, int.class);
                        constructor.setAccessible(true);
                        logger.debug("Constructor for direct buffer found and made accessible");
                        return constructor;
                      } catch (NoSuchMethodException e) {
                        logger.debug("Cannot get constructor for direct buffer allocation", e);
                        return e;
                      } catch (SecurityException e) {
                        logger.debug("Cannot get constructor for direct buffer allocation", e);
                        return e;
                      }
                    }
                  });
  
          if (maybeDirectBufferConstructor instanceof Constructor<?>) {
            address = UNSAFE.allocateMemory(1);
            // try to use the constructor now
            try {
              ((Constructor<?>) maybeDirectBufferConstructor).newInstance(address, 1);
              directBufferConstructor = (Constructor<?>) maybeDirectBufferConstructor;
              logger.debug("direct buffer constructor: available");
            } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
              logger.warn("unable to instantiate a direct buffer via constructor", e);
              directBufferConstructor = null;
            }
          } else {
            logger.debug(
                "direct buffer constructor: unavailable", (Throwable) maybeDirectBufferConstructor);
            directBufferConstructor = null;
          }
        } finally {
          if (address != -1) {
            UNSAFE.freeMemory(address);
          }
        }
        DIRECT_BUFFER_CONSTRUCTOR = directBufferConstructor;
      } catch (Throwable e) {
        // This exception will get swallowed, but it's necessary for the static analysis that ensures
        // the static fields above get initialized
        final RuntimeException failure =
            new RuntimeException(
                "Failed to initialize MemoryUtil. You must start Java with "
                    + "`--add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED` "
                    + "(See https://arrow.apache.org/docs/java/install.html)",
                e);
        failure.printStackTrace();
        throw failure;
      }
    }
  
    /**
     * Given a {@link ByteBuffer}, gets the address the underlying memory space.
     *
     * @param buf the byte buffer.
     * @return address of the underlying memory.
     */
    public static long getByteBufferAddress(ByteBuffer buf) {
      return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET);
    }
  
    private MemoryUtil() {}
  
    /** Create nio byte buffer. */
    public static ByteBuffer directBuffer(long address, int capacity) {
      if (DIRECT_BUFFER_CONSTRUCTOR != null) {
        if (capacity < 0) {
          throw new IllegalArgumentException("Capacity is negative, has to be positive or 0");
        }
        try {
          return (ByteBuffer) DIRECT_BUFFER_CONSTRUCTOR.newInstance(address, capacity);
        } catch (Throwable cause) {
          throw new Error(cause);
        }
      }
      throw new UnsupportedOperationException(
          "sun.misc.Unsafe or java.nio.DirectByteBuffer.<init>(long, int) not available");
    }
  
    @SuppressWarnings(
        "nullness:argument") // to handle null assignment on third party dependency: Unsafe
    private static void copyMemory(
        @Nullable Object srcBase,
        long srcOffset,
        @Nullable Object destBase,
        long destOffset,
        long bytes) {
      UNSAFE.copyMemory(srcBase, srcOffset, destBase, destOffset, bytes);
    }
  
    public static void copyMemory(long srcAddress, long destAddress, long bytes) {
      UNSAFE.copyMemory(srcAddress, destAddress, bytes);
    }
  
    public static void copyToMemory(byte[] src, long srcIndex, long destAddress, long bytes) {
      copyMemory(src, BYTE_ARRAY_BASE_OFFSET + srcIndex, null, destAddress, bytes);
    }
  
    public static void copyFromMemory(long srcAddress, byte[] dest, long destIndex, long bytes) {
      copyMemory(null, srcAddress, dest, BYTE_ARRAY_BASE_OFFSET + destIndex, bytes);
    }
  
    public static byte getByte(long address) {
      return UNSAFE.getByte(address);
    }
  
    public static void putByte(long address, byte value) {
      UNSAFE.putByte(address, value);
    }
  
    public static short getShort(long address) {
      return UNSAFE.getShort(address);
    }
  
    public static void putShort(long address, short value) {
      UNSAFE.putShort(address, value);
    }
  
    public static int getInt(long address) {
      return UNSAFE.getInt(address);
    }
  
    public static void putInt(long address, int value) {
      UNSAFE.putInt(address, value);
    }
  
    public static long getLong(long address) {
      return UNSAFE.getLong(address);
    }
  
    public static void putLong(long address, long value) {
      UNSAFE.putLong(address, value);
    }
  
    public static void setMemory(long address, long bytes, byte value) {
      UNSAFE.setMemory(address, bytes, value);
    }
  
    public static int getInt(byte[] bytes, int index) {
      return UNSAFE.getInt(bytes, BYTE_ARRAY_BASE_OFFSET + index);
    }
  
    public static long getLong(byte[] bytes, int index) {
      return UNSAFE.getLong(bytes, BYTE_ARRAY_BASE_OFFSET + index);
    }
  
    public static long allocateMemory(long bytes) {
      return UNSAFE.allocateMemory(bytes);
    }
  
    public static void freeMemory(long address) {
      UNSAFE.freeMemory(address);
    }
  }
  