JvmSizeUtils.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.doris.datasource.metacache;

import com.sun.management.HotSpotDiagnosticMXBean;

import java.lang.management.ManagementFactory;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/** HotSpot-aware formulas used by production cache size estimators. */
public final class JvmSizeUtils {
    // These fallbacks are deliberately conservative. OwnedObjectSizeEstimator rejects the graph
    // when the layout is unknown, but formula-only callers still must not silently choose the
    // smaller compressed-oops layout.
    private static final int DEFAULT_ALIGNMENT = 16;
    private static final Boolean CONFIGURED_COMPRESSED_OOPS = booleanVmOption("UseCompressedOops");
    private static final Boolean CONFIGURED_COMPRESSED_CLASS_POINTERS =
            booleanVmOption("UseCompressedClassPointers");
    private static final Boolean CONFIGURED_COMPACT_STRINGS =
            isJava9OrLater() ? booleanVmOption("CompactStrings") : Boolean.FALSE;
    private static final Integer CONFIGURED_OBJECT_ALIGNMENT = intVmOption("ObjectAlignmentInBytes");
    private static final boolean VM_LAYOUT_KNOWN = CONFIGURED_COMPRESSED_OOPS != null
            && CONFIGURED_COMPRESSED_CLASS_POINTERS != null
            && CONFIGURED_COMPACT_STRINGS != null
            && CONFIGURED_OBJECT_ALIGNMENT != null;
    private static final boolean COMPRESSED_OOPS = Boolean.TRUE.equals(CONFIGURED_COMPRESSED_OOPS);
    private static final boolean COMPRESSED_CLASS_POINTERS =
            Boolean.TRUE.equals(CONFIGURED_COMPRESSED_CLASS_POINTERS);
    private static final boolean COMPACT_STRINGS = Boolean.TRUE.equals(CONFIGURED_COMPACT_STRINGS);
    private static final int OBJECT_ALIGNMENT =
            CONFIGURED_OBJECT_ALIGNMENT == null ? DEFAULT_ALIGNMENT : CONFIGURED_OBJECT_ALIGNMENT;
    private static final int REFERENCE_BYTES = COMPRESSED_OOPS ? 4 : 8;
    private static final int CLASS_POINTER_BYTES = COMPRESSED_CLASS_POINTERS ? 4 : 8;
    private static final int OBJECT_HEADER_BYTES = 8 + CLASS_POINTER_BYTES;
    private static final int ARRAY_HEADER_BYTES = alignInt(OBJECT_HEADER_BYTES + Integer.BYTES);

    private static final ClassValue<Long> SHALLOW_SIZES = new ClassValue<Long>() {
        @Override
        protected Long computeValue(Class<?> type) {
            if (type.isArray()) {
                throw new IllegalArgumentException("array shallow size requires a length");
            }
            long fields = 0L;
            int hierarchyDepth = 0;
            for (Class<?> current = type; current != null; current = current.getSuperclass()) {
                hierarchyDepth++;
                for (Field field : current.getDeclaredFields()) {
                    if (!Modifier.isStatic(field.getModifiers())) {
                        fields = saturatedAdd(fields, fieldSize(field.getType()));
                    }
                }
            }
            // HotSpot can leave an alignment gap at an inheritance boundary. Adding at most one
            // alignment unit per superclass is conservative without accessing module-private offsets.
            long inheritancePadding = Math.max(0, hierarchyDepth - 1) * (long) (OBJECT_ALIGNMENT - 1);
            return align(saturatedAdd(OBJECT_HEADER_BYTES, saturatedAdd(fields, inheritancePadding)));
        }
    };

    private JvmSizeUtils() {
    }

    public static long shallowSizeOf(Class<?> type) {
        return SHALLOW_SIZES.get(type);
    }

    /** Whether every VM layout option needed by the retained-size formulas was read successfully. */
    public static boolean isVmLayoutKnown() {
        return VM_LAYOUT_KNOWN;
    }

    public static long sizeOfString(String value) {
        if (value == null) {
            return 0L;
        }
        int elementBytes = COMPACT_STRINGS && isLatin1(value) ? Byte.BYTES : Character.BYTES;
        return saturatedAdd(shallowSizeOf(String.class), sizeOfPrimitiveArray(value.length(), elementBytes));
    }

    public static long sizeOfObjectArray(int length) {
        return sizeOfPrimitiveArray(length, REFERENCE_BYTES);
    }

    public static long byteArraySize(int length) {
        return sizeOfPrimitiveArray(length, Byte.BYTES);
    }

    public static long intArraySize(int length) {
        return sizeOfPrimitiveArray(length, Integer.BYTES);
    }

    public static long longArraySize(int length) {
        return sizeOfPrimitiveArray(length, Long.BYTES);
    }

    public static long sizeOfPrimitiveArray(int length, int elementBytes) {
        if (length < 0 || elementBytes < 0) {
            throw new IllegalArgumentException("array length and element size must be non-negative");
        }
        return align(saturatedAdd(ARRAY_HEADER_BYTES, saturatedMultiply(length, elementBytes)));
    }

    static int hashMapInitialCapacityForSize(int size) {
        if (size < 0) {
            throw new IllegalArgumentException("map size must be non-negative");
        }
        // Match the capacity requested by HashMap's copy constructor without relying on float
        // rounding: floor(size / 0.75) + 1. The extra one matters at 3 * 2^n boundaries.
        long requested = saturatedAdd(saturatedMultiply(size, 4L) / 3L, 1L);
        return (int) Math.min(requested, 1L << 30);
    }

    static int hashMapTableCapacityForSize(int size) {
        return size == 0 ? 0 : tableSizeFor(hashMapInitialCapacityForSize(size));
    }

    public static long sizeOfKnownCapacityHashMapStorage(int size, int tableCapacity) {
        if (size < 0 || tableCapacity < 0
                || (tableCapacity != 0 && (tableCapacity & (tableCapacity - 1)) != 0)) {
            throw new IllegalArgumentException("map size and power-of-two table capacity are invalid");
        }
        long table = sizeOfObjectArray(tableCapacity);
        // KnownCapacityHashMap.Node: object header + hash + key/value/next references.
        long node = align(OBJECT_HEADER_BYTES + Integer.BYTES + 3L * REFERENCE_BYTES);
        return saturatedAdd(table, saturatedMultiply(size, node));
    }

    /**
     * Conservative storage model for a map whose concrete backing capacity is not exposed.
     *
     * <p>The model uses the HashMap copy-constructor capacity and one node per mapping. Compact
     * immutable maps normally retain less; using this upper model prevents third-party SDK map
     * implementations from being under-counted merely because their private table is inaccessible.
     */
    public static long sizeOfConservativeMapStorage(int size) {
        if (size < 0) {
            throw new IllegalArgumentException("map size must be non-negative");
        }
        return sizeOfKnownCapacityHashMapStorage(size, hashMapTableCapacityForSize(size));
    }

    public static long sizeOfHashBiMapStorage(int size) {
        if (size <= 0) {
            return saturatedMultiply(2L, sizeOfObjectArray(2));
        }
        // HashBiMap keeps forward and inverse hash indexes plus one entry record per mapping.
        int modeledTableCapacity = hashMapTableCapacityForSize(size);
        long modeledIndex = sizeOfKnownCapacityHashMapStorage(size, modeledTableCapacity);
        return saturatedAdd(modeledIndex, modeledIndex);
    }

    /** Conservative retained storage for SDK lists exposed through public collection APIs. */
    public static long sizeOfConservativeListStorage(List<?> values) {
        if (values == null) {
            return 0L;
        }
        long bytes = shallowSizeOf(values.getClass());
        // SDKs commonly expose an unmodifiable wrapper around an ArrayList.
        bytes = saturatedAdd(bytes, shallowSizeOf(ArrayList.class));
        return saturatedAdd(bytes, sizeOfObjectArray(values.size()));
    }

    /** Conservative retained storage for SDK maps whose backing capacity is not public. */
    public static long sizeOfConservativeMap(Map<?, ?> values) {
        if (values == null) {
            return 0L;
        }
        long bytes = saturatedAdd(shallowSizeOf(values.getClass()),
                sizeOfConservativeMapStorage(values.size()));
        // Covers linked-map nodes and unmodifiable/delegating wrappers.
        return saturatedAdd(bytes, saturatedMultiply(values.size(), 16L));
    }

    public static long sizeOfConservativeIndexMap(int size) {
        return saturatedAdd(shallowSizeOf(java.util.HashMap.class),
                sizeOfConservativeMapStorage(size));
    }

    public static long sizeOfStringList(List<String> values) {
        if (values == null) {
            return 0L;
        }
        long bytes = sizeOfConservativeListStorage(values);
        for (String value : values) {
            bytes = saturatedAdd(bytes, sizeOfString(value));
        }
        return bytes;
    }

    public static long sizeOfStringMap(Map<String, String> values) {
        if (values == null) {
            return 0L;
        }
        long bytes = sizeOfConservativeMap(values);
        for (Map.Entry<String, String> entry : values.entrySet()) {
            bytes = saturatedAdd(bytes, sizeOfString(entry.getKey()));
            bytes = saturatedAdd(bytes, sizeOfString(entry.getValue()));
        }
        return bytes;
    }

    public static long align(long bytes) {
        if (bytes >= Long.MAX_VALUE - OBJECT_ALIGNMENT) {
            return Long.MAX_VALUE;
        }
        long mask = OBJECT_ALIGNMENT - 1L;
        return (bytes + mask) & ~mask;
    }

    public static long saturatedAdd(long left, long right) {
        if (left < 0 || right < 0 || Long.MAX_VALUE - left < right) {
            return Long.MAX_VALUE;
        }
        return left + right;
    }

    public static long saturatedMultiply(long left, long right) {
        if (left < 0 || right < 0 || (left != 0L && right > Long.MAX_VALUE / left)) {
            return Long.MAX_VALUE;
        }
        return left * right;
    }

    private static long fieldSize(Class<?> type) {
        if (!type.isPrimitive()) {
            return REFERENCE_BYTES;
        }
        if (type == boolean.class || type == byte.class) {
            return 1L;
        }
        if (type == char.class || type == short.class) {
            return 2L;
        }
        if (type == int.class || type == float.class) {
            return 4L;
        }
        return 8L;
    }

    private static boolean isLatin1(String value) {
        if (!COMPACT_STRINGS) {
            return false;
        }
        for (int i = 0; i < value.length(); i++) {
            if (value.charAt(i) > 0xFF) {
                return false;
            }
        }
        return true;
    }

    private static int tableSizeFor(long requested) {
        long value = 1L;
        while (value < requested && value < (1L << 30)) {
            value <<= 1;
        }
        return (int) Math.min(value, 1L << 30);
    }

    private static int alignInt(int bytes) {
        return (int) (((long) bytes + OBJECT_ALIGNMENT - 1L) & ~(OBJECT_ALIGNMENT - 1L));
    }

    private static boolean isJava9OrLater() {
        String version = System.getProperty("java.specification.version", "8");
        try {
            return version.indexOf('.') >= 0
                    ? Double.parseDouble(version) >= 1.9
                    : Integer.parseInt(version) >= 9;
        } catch (NumberFormatException e) {
            return false;
        }
    }

    private static Boolean booleanVmOption(String name) {
        try {
            return Boolean.parseBoolean(hotSpotBean().getVMOption(name).getValue());
        } catch (RuntimeException e) {
            return null;
        }
    }

    private static Integer intVmOption(String name) {
        try {
            int value = Integer.parseInt(hotSpotBean().getVMOption(name).getValue());
            return value > 0 && (value & (value - 1)) == 0 ? value : null;
        } catch (RuntimeException e) {
            return null;
        }
    }

    private static HotSpotDiagnosticMXBean hotSpotBean() {
        return ManagementFactory.getPlatformMXBean(HotSpotDiagnosticMXBean.class);
    }
}