001/*
002 *  Copyright 2020 Anyware Services
003 *
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package org.ametys.core.util;
017
018import java.lang.annotation.ElementType;
019import java.lang.annotation.Retention;
020import java.lang.annotation.RetentionPolicy;
021import java.lang.annotation.Target;
022import java.lang.reflect.Field;
023import java.util.Arrays;
024import java.util.Collections;
025import java.util.IdentityHashMap;
026import java.util.List;
027import java.util.Map;
028import java.util.Set;
029import java.util.concurrent.ConcurrentHashMap;
030import java.util.stream.Collectors;
031
032import org.apache.commons.lang3.tuple.Pair;
033import org.openjdk.jol.info.ClassData;
034import org.openjdk.jol.info.FieldData;
035import org.openjdk.jol.layouters.CurrentLayouter;
036import org.openjdk.jol.util.ObjectUtils;
037import org.openjdk.jol.vm.VM;
038import org.slf4j.LoggerFactory;
039
040/**
041 * Helper for manipulating objects heap size.
042 */
043public final class SizeUtils
044{
045    private static Map<Class<?>, Pair<Long, List<Field>>> __cache = new ConcurrentHashMap<>();
046    
047    static
048    {
049        System.setProperty("jol.magicFieldOffset", "true");
050
051        // add opens to JOL for all java.base packages, to avoid illegal access warnings
052        Object.class.getModule().getPackages().forEach(p -> 
053        {
054            if (Object.class.getModule().isOpen(p, ObjectUtils.class.getModule()))
055            {
056                Object.class.getModule().addOpens(p, ObjectUtils.class.getModule());
057            }
058        });
059    }
060    
061    private SizeUtils()
062    {
063        // empty
064    }
065    
066    /**
067     * Returns true if size computation is enabled for this VM
068     * @return true if size computation is enabled
069     */
070    public static boolean enabled()
071    {
072        try
073        {
074            VM.current();
075            return true;
076        }
077        catch (IllegalStateException e)
078        {
079            // only Hotspot/OpenJDK VM are supported
080            return false;
081        }
082    }
083    
084    /**
085     * Calculates the heap size of the given object.
086     * @param object the object
087     * @return the size of the object.
088     */
089    public static long sizeOf(Object object)
090    {
091        if (object == null)
092        {
093            return 0;
094        }
095        
096        try
097        {
098            return _sizeOf(object, Collections.newSetFromMap(new IdentityHashMap<Object, Boolean>()));
099        }
100        catch (Error e)
101        {
102            // Special case for intercepting eg. StackOverflowError and log the concerned object
103            LoggerFactory.getLogger(SizeUtils.class.getName()).error("Error computing size of {}", object, e);
104            throw e;
105        }
106    }
107    
108    private static long _sizeOf(Object object, Set<Object> visited)
109    {
110        if (!visited.add(object))
111        {
112            return 0;
113        }
114        
115        Class<?> klass = object.getClass();
116        
117        if (klass.isArray())
118        {
119            long size = new CurrentLayouter().layout(ClassData.parseInstance(object)).instanceSize();
120
121            if (klass.getComponentType().isPrimitive()) 
122            {
123                return size;
124            }
125            
126            Object[] array = (Object[]) object;
127            
128            return Arrays.stream(array)
129                    .filter(SizeUtils::_isComputable)
130                    .map(value -> _sizeOf(value, visited))
131                    .reduce(size, Long::sum);
132        }
133        
134        // we cache the shallow class size and the list of class' fields
135        Pair<Long, List<Field>> data = __cache.computeIfAbsent(klass, c -> 
136        {
137            ClassData classData = ClassData.parseClass(c);
138
139            List<Field> fields = classData.fields().stream()
140                                                   .map(FieldData::refField)
141                                                   .filter(field -> !field.getType().isPrimitive())
142                                                   .filter(field -> !field.isAnnotationPresent(ExcludeFromSizeCalculation.class))
143                                                   .collect(Collectors.toList());
144
145            return Pair.of(new CurrentLayouter().layout(classData).instanceSize(), fields);
146        });
147        
148        // the actual size of a non-array object is the sum of its shallow size and the size of each of its fields
149        long size = data.getLeft();
150        
151        for (Field field : data.getRight())
152        {
153            Object value = ObjectUtils.value(object, field);
154            
155            if (_isComputable(value))
156            {
157                size += _sizeOf(value, visited);
158            }
159        }
160        
161        return size;
162    }
163
164    private static boolean _isComputable(Object object)
165    {
166        return object != null && !(object instanceof Class);
167    }
168    
169    /**
170     * Fields annotated with {@link ExcludeFromSizeCalculation} are excluded for size calculation
171     */
172    @Retention(RetentionPolicy.RUNTIME)
173    @Target(ElementType.FIELD)
174    public static @interface ExcludeFromSizeCalculation
175    {
176        // marker annotation
177    }
178}