test/unit/org/apache/cassandra/utils/concurrent/RefCountedTest.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.cassandra.utils.concurrent;

import java.lang.ref.WeakReference;
import java.time.Duration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;

import org.junit.Assert;
import org.junit.Test;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.dht.ByteOrderedPartitioner;
import org.apache.cassandra.io.sstable.Descriptor;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.util.File;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.schema.TableMetadataRef;
import org.apache.cassandra.utils.ObjectSizes;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.concurrent.Ref.Visitor;
import org.awaitility.Awaitility;

import static org.assertj.core.api.Assertions.assertThat;

@SuppressWarnings({"unused", "unchecked", "rawtypes"})
public class RefCountedTest
{
    static
    {
        if (Ref.STRONG_LEAK_DETECTOR != null)
            Ref.STRONG_LEAK_DETECTOR.submit(() -> { Thread.sleep(Integer.MAX_VALUE); return null; });
    }

    private static final class Tidier implements RefCounted.Tidy
    {
        boolean tidied;

        public void tidy()
        {
            tidied = true;
        }

        public String name()
        {
            return "test tidy";
        }
    }

    @Test
    public void testLeak() throws InterruptedException
    {
        Tidier tidier = new Tidier();
        Ref<?> obj = new Ref(null, tidier);
        obj.tryRef();
        obj.release();
        System.gc();
        System.gc();
        Thread.sleep(1000);
        Assert.assertTrue(tidier.tidied);
    }

    @Test
    public void testSeriousLeak() throws InterruptedException
    {
        Tidier tidier = new Tidier();
        new Ref(null, tidier);
        System.gc();
        System.gc();
        System.gc();
        System.gc();
        Thread.sleep(1000);
        Assert.assertTrue(tidier.tidied);
    }

    @Test
    public void testDoubleRelease() throws InterruptedException
    {
        Tidier tidier = null;
        try
        {
            tidier = new Tidier();
            Ref<?> obj = new Ref(null, tidier);
            obj.release();
            obj.release();
            Assert.assertTrue(false);
        }
        catch (Exception e)
        {
        }
    }

    @Test
    public void testMemoryLeak()
    {
        Tidier tidier = new Tidier();
        Ref<Object> ref = new Ref(null, tidier);
        long initialSize = ObjectSizes.measureDeep(ref);
        for (int i = 0 ; i < 1000 ; i++)
            ref.ref().release();
        long finalSize = ObjectSizes.measureDeep(ref);
        if (finalSize > initialSize * 2)
            throw new AssertionError();
        ref.release();
    }

    static final int entryCount = 1000000;
    static final int fudgeFactor = 20;

    @Test
    public void testLinkedList()
    {
        final List<Object> iterable = new LinkedList<Object>();
        Pair<Object, Object> p = Pair.create(iterable, iterable);
        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            Object ref = iterable;
            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };
        Ref<Object> ref = new Ref(new AtomicReference<List<Object>>(iterable), tidier);
        for (int i = 0; i < entryCount; i++)
        {
            iterable.add(p);
        }
        Visitor visitor = new Visitor();
        visitor.run();
        ref.close();

        System.out.println("LinkedList visited " + visitor.lastVisitedCount + " iterations " + visitor.iterations);
        //Should visit a lot of list nodes, but no more since there is only one object stored in the list
        Assert.assertTrue(visitor.lastVisitedCount > entryCount && visitor.lastVisitedCount < entryCount + fudgeFactor);
        //Should have a lot of iterations to walk the list, but linear to the number of entries
        Assert.assertTrue(visitor.iterations > (entryCount * 3) && visitor.iterations < (entryCount * 3) + fudgeFactor);
    }

    /*
     * There was a traversal error terminating traversal for an object upon encountering a null
     * field. Test for the bug here using CLQ.
     */
    @Test
    public void testCLQBug()
    {
        Ref.concurrentIterables.remove(ConcurrentLinkedQueue.class);
        try
        {
            testConcurrentLinkedQueueImpl(true);
        }
        finally
        {
            Ref.concurrentIterables.add(ConcurrentLinkedQueue.class);
        }
    }

    private void testConcurrentLinkedQueueImpl(boolean bugTest)
    {
        final Queue<Object> iterable = new ConcurrentLinkedQueue<Object>();
        Pair<Object, Object> p = Pair.create(iterable, iterable);
        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            Object ref = iterable;
            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };
        Ref<Object> ref = new Ref(new AtomicReference<Queue<Object>>(iterable), tidier);
        for (int i = 0; i < entryCount; i++)
        {
            iterable.add(p);
        }
        Visitor visitor = new Visitor();
        visitor.run();
        ref.close();

        System.out.println("ConcurrentLinkedQueue visited " + visitor.lastVisitedCount + " iterations " + visitor.iterations + " bug test " + bugTest);

        if (bugTest)
        {
            //Should have to visit a lot of queue nodes
            Assert.assertTrue(visitor.lastVisitedCount > entryCount && visitor.lastVisitedCount < entryCount + fudgeFactor);
            //Should have a lot of iterations to walk the queue, but linear to the number of entries
            Assert.assertTrue(visitor.iterations > (entryCount * 2) && visitor.iterations < (entryCount * 2) + fudgeFactor);
        }
        else
        {
            //There are almost no objects in this linked list once it's iterated as a collection so visited count
            //should be small
            Assert.assertTrue(visitor.lastVisitedCount < 10);
            //Should have a lot of iterations to walk the collection, but linear to the number of entries
            Assert.assertTrue(visitor.iterations > entryCount && visitor.iterations < entryCount + fudgeFactor);
        }
    }

    @Test
    public void testConcurrentLinkedQueue()
    {
        testConcurrentLinkedQueueImpl(false);
    }

    @Test
    public void testBlockingQueue()
    {
        final BlockingQueue<Object> iterable = new LinkedBlockingQueue<Object>();
        Pair<Object, Object> p = Pair.create(iterable, iterable);
        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            Object ref = iterable;
            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };
        Ref<Object> ref = new Ref(new AtomicReference<BlockingQueue<Object>>(iterable), tidier);
        for (int i = 0; i < entryCount; i++)
        {
            iterable.add(p);
        }
        Visitor visitor = new Visitor();
        visitor.run();
        ref.close();

        System.out.println("BlockingQueue visited " + visitor.lastVisitedCount + " iterations " + visitor.iterations);
        //There are almost no objects in this queue once it's iterated as a collection so visited count
        //should be small
        Assert.assertTrue(visitor.lastVisitedCount < 10);
        //Should have a lot of iterations to walk the collection, but linear to the number of entries
        Assert.assertTrue(visitor.iterations > entryCount && visitor.iterations < entryCount + fudgeFactor);
    }

    @Test
    public void testConcurrentMap()
    {
        final Map<Object, Object> map = new ConcurrentHashMap<Object, Object>();
        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            Object ref = map;
            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };
        Ref<Object> ref = new Ref(new AtomicReference<Map<Object, Object>>(map), tidier);

        Object o = new Object();
        for (int i = 0; i < entryCount; i++)
        {
            map.put(new Object(), o);
        }
        Visitor visitor = new Visitor();
        visitor.run();
        ref.close();

        System.out.println("ConcurrentHashMap visited " + visitor.lastVisitedCount + " iterations " + visitor.iterations);

        //Should visit roughly the same number of objects as entries because the value object is constant
        //Map.Entry objects shouldn't be counted since it is iterated as a collection
        Assert.assertTrue(visitor.lastVisitedCount > entryCount && visitor.lastVisitedCount < entryCount + fudgeFactor);
        //Should visit 2x the number of entries since we have to traverse the key and value separately
        Assert.assertTrue(visitor.iterations > entryCount * 2 && visitor.iterations < entryCount * 2 + fudgeFactor);
    }

    @Test
    public void testHashMap()
    {
        final Map<Object, Object> map = new HashMap<Object, Object>();
        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            Object ref = map;
            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };
        Ref<Object> ref = new Ref(new AtomicReference<Map<Object, Object>>(map), tidier);

        Object o = new Object();
        for (int i = 0; i < entryCount; i++)
        {
            map.put(new Object(), o);
        }
        Visitor visitor = new Visitor();
        visitor.run();
        ref.close();

        System.out.println("HashMap visited " + visitor.lastVisitedCount + " iterations " + visitor.iterations);

        //Should visit 2x  the number of entries because of the wrapper Map.Entry objects
        Assert.assertTrue(visitor.lastVisitedCount > (entryCount * 2) && visitor.lastVisitedCount < (entryCount * 2) + fudgeFactor);
        //Should iterate 3x the number of entries since we have to traverse the key and value separately
        Assert.assertTrue(visitor.iterations > (entryCount * 3) && visitor.iterations < (entryCount * 3) + fudgeFactor);
    }

    @Test
    public void testArray() throws Exception
    {
        final Object objects[] = new Object[entryCount];
        for (int i = 0; i < entryCount; i += 2)
            objects[i] = new Object();

        File f = FileUtils.createTempFile("foo", "bar");
        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            Object ref = objects;
            //Checking we don't get an infinite loop out of traversing file refs
            File fileRef = f;

            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };
        Ref<Object> ref = new Ref(new AtomicReference<Object[]>(objects), tidier);

        Visitor visitor = new Visitor();
        visitor.run();
        ref.close();

        System.out.println("Array visited " + visitor.lastVisitedCount + " iterations " + visitor.iterations);
        //Should iterate the elements in the array and get a unique object from every other one
        Assert.assertTrue(visitor.lastVisitedCount > (entryCount / 2) && visitor.lastVisitedCount < (entryCount / 2) + fudgeFactor);
        //Should iterate over the array touching roughly the same number of objects as entries
        Assert.assertTrue(visitor.iterations > (entryCount / 2) && visitor.iterations < (entryCount / 2) + fudgeFactor);
    }

    //Make sure a weak ref is ignored by the visitor looking for strong ref leaks
    @Test
    public void testWeakRef() throws Exception
    {
        AtomicReference dontRefMe = new AtomicReference();

        WeakReference<Object> weakRef = new WeakReference(dontRefMe);

        RefCounted.Tidy tidier = new RefCounted.Tidy() {
            WeakReference<Object> ref = weakRef;

            @Override
            public void tidy() throws Exception
            {
            }

            @Override
            public String name()
            {
                return "42";
            }
        };

        Ref<Object> ref = new Ref(dontRefMe, tidier);
        dontRefMe.set(ref);

        Visitor visitor = new Visitor();
        visitor.haveLoops = new HashSet<>();
        visitor.run();
        ref.close();

        Assert.assertTrue(visitor.haveLoops.isEmpty());
    }

    static class LambdaTestClassTidier implements RefCounted.Tidy
    {
        Runnable runOnClose;

        @Override
        public void tidy() throws Exception
        {
            runOnClose.run();
        }

        @Override
        public String name()
        {
            return "42";
        }
    }

    static class LambdaTestClass
    {
        String a = "x";
        Ref<Object> ref;

        Runnable getRunOnCloseLambda()
        {
            return () -> System.out.println("aaa");
        }

        Runnable getRunOnCloseInner()
        {
            return new Runnable()
            {
                @Override
                public void run()
                {
                    System.out.println("aaa");
                }
            };
        }

        Runnable getRunOnCloseLambdaWithThis()
        {
            return () -> System.out.println(a);
        }
    }

    private Set<Ref.GlobalState> testCycles(Function<LambdaTestClass, Runnable> runOnCloseSupplier)
    {
        LambdaTestClass test = new LambdaTestClass();
        Runnable weakRef = runOnCloseSupplier.apply(test);
        RefCounted.Tidy tidier = new RefCounted.Tidy()
        {
            Runnable ref = weakRef;

            public void tidy()
            {
            }

            public String name()
            {
                return "42";
            }
        };

        Ref<Object> ref = new Ref(test, tidier);
        test.ref = ref;

        Visitor visitor = new Visitor();
        visitor.haveLoops = new HashSet<>();
        visitor.run();
        ref.close();

        return visitor.haveLoops;
    }

    /**
     * The intention of this test is to confirm that lambda without a reference to `this`, does not implicitly
     * include a reference to the enclosing class. If it did, we would detect cycles, as we do when we deal with
     * anonymous inner classes or lamba which references `this` explicitly (see the sanity checks).
     * <p>
     * This test aims to confirm that JVM works as we assumed because we couldn't find that in the specification.
     */
    @Test
    public void testCycles()
    {
        assertThat(testCycles(LambdaTestClass::getRunOnCloseLambdaWithThis)).isNotEmpty(); // sanity test
        assertThat(testCycles(LambdaTestClass::getRunOnCloseInner)).isNotEmpty(); // sanity test

        assertThat(testCycles(LambdaTestClass::getRunOnCloseLambda)).isEmpty();
    }

    @Test
    public void testSSTableReaderLeakIsDetected()
    {
        DatabaseDescriptor.clientInitialization();
        DatabaseDescriptor.setPartitionerUnsafe(ByteOrderedPartitioner.instance);
        Descriptor descriptor = Descriptor.fromFileWithComponent(new File("test/data/legacy-sstables/nb/legacy_tables/legacy_nb_simple/nb-1-big-Data.db"), false).left;
        TableMetadata tm = TableMetadata.builder("legacy_tables", "legacy_nb_simple").addPartitionKeyColumn("pk", UTF8Type.instance).addRegularColumn("val", UTF8Type.instance).build();
        AtomicBoolean leakDetected = new AtomicBoolean();
        AtomicBoolean runOnCloseExecuted1 = new AtomicBoolean();
        AtomicBoolean runOnCloseExecuted2 = new AtomicBoolean();
        Ref.OnLeak prevOnLeak = Ref.ON_LEAK;

        try
        {
            Ref.ON_LEAK = state -> {
                leakDetected.set(true);
            };
            {
                SSTableReader reader = SSTableReader.openNoValidation(null, descriptor, TableMetadataRef.forOfflineTools(tm));
                reader.runOnClose(() -> runOnCloseExecuted1.set(true));
                reader.runOnClose(() -> runOnCloseExecuted2.set(true)); // second time to actually create lambda referencing to lambda, see runOnClose impl
                //noinspection UnusedAssignment
                reader = null; // this is required, otherwise GC will not attempt to collect the created reader
            }
            Awaitility.await().atMost(Duration.ofSeconds(30)).pollDelay(Duration.ofSeconds(1)).untilAsserted(() -> {
                System.gc();
                System.gc();
                System.gc();
                System.gc();
                assertThat(leakDetected.get()).isTrue();
            });
            assertThat(runOnCloseExecuted1.get()).isTrue();
            assertThat(runOnCloseExecuted2.get()).isTrue();
        }
        finally
        {
            Ref.ON_LEAK = prevOnLeak;
        }
    }
}