package org.nd4j.rng.deallocator;

import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.nd4j.nativeblas.NativeOpsHolder;

import java.lang.ref.ReferenceQueue;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.LockSupport;

/**
 * Since NativeRandom assumes some native resources, we have to track their use, and deallocate them as soon they are released by JVM GC
 *
 * @author raver119@gmail.com
 */
@Slf4j
public class NativeRandomDeallocator {
    private static final NativeRandomDeallocator INSTANCE = new NativeRandomDeallocator();

    // we don't really need concurrency here, so 1 queue will be just fine
    private final ReferenceQueue<NativePack> queue;
    private final Map<Long, GarbageStateReference> referenceMap;
    private List<DeallocatorThread> deallocatorThreads = new ArrayList<>();

    private NativeRandomDeallocator() {
        this.queue = new ReferenceQueue<>();
        this.referenceMap = new ConcurrentHashMap<>();

        DeallocatorThread thread = new DeallocatorThread(0, queue, referenceMap);
        thread.start();

        deallocatorThreads.add(thread);
    }

    public static NativeRandomDeallocator getInstance() {
        return INSTANCE;
    }


    /**
     * This method is used internally from NativeRandom deallocators
     * This method doesn't accept Random interface implementations intentionally.
     *
     * @param random
     */
    public void trackStatePointer(NativePack random) {
        if (random.getStatePointer() != null) {
            GarbageStateReference reference = new GarbageStateReference(random, queue);
            referenceMap.put(random.getStatePointer().address(), reference);
        }
    }


    /**
     * This class provides garbage collection for NativeRandom state memory. It's not too big amount of memory used, but we don't want any leaks.
     *
     */
    protected class DeallocatorThread extends Thread implements Runnable {
        private final ReferenceQueue<NativePack> queue;
        private final Map<Long, GarbageStateReference> referenceMap;

        protected DeallocatorThread(int threadId, @NonNull ReferenceQueue<NativePack> queue,
                        Map<Long, GarbageStateReference> referenceMap) {
            this.queue = queue;
            this.referenceMap = referenceMap;
            this.setName("NativeRandomDeallocator thread " + threadId);
            this.setDaemon(true);
        }

        @Override
        public void run() {
            while (true) {
                GarbageStateReference reference = (GarbageStateReference) queue.poll();
                if (reference != null) {
                    if (reference.getStatePointer() != null) {
                        referenceMap.remove(reference.getStatePointer().address());
                        NativeOpsHolder.getInstance().getDeviceNativeOps().destroyRandom(reference.getStatePointer());
                    }
                } else {
                    LockSupport.parkNanos(5000000L);
                }
            }
        }
    }
}

@@ -2587,24 +2587,50 @@
 le (true) %7B%0A
+                try %7B%0A    
             
@@ -2701,16 +2701,22 @@
 eue.
-poll();%0A
+remove();%0A    
     
@@ -2744,32 +2744,36 @@
 ence != null) %7B%0A
+    
                 
@@ -2835,32 +2835,36 @@
                 
+    
 referenceMap.rem
@@ -2903,24 +2903,28 @@
 address());%0A
+    
             
@@ -3042,34 +3042,42 @@
                 
-%7D%0A
+    %7D%0A    
                 
@@ -3081,24 +3081,28 @@
    %7D else %7B%0A
+    
             
@@ -3139,12 +3139,207 @@
 5000
-000L
+L);%0A                    %7D%0A                %7D catch (InterruptedException e) %7B%0A                    // do nothing%0A                %7D catch (Exception e) %7B%0A                    throw new RuntimeException(e
 );%0A 

fix for blocking thread