From 0664bf6334b78cd7f72d959ff9ea7982a7bce5d9 Mon Sep 17 00:00:00 2001
From: Karl Lessard <karl.lessard@gmail.com>
Date: Mon, 20 Jan 2020 15:01:56 -0500
Subject: [PATCH] Prevent cleanup thread to hang on exit

---
 .../main/java/org/tensorflow/EagerSession.java    | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerSession.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerSession.java
index 71619b516c5..8ca098c3327 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerSession.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerSession.java
@@ -274,13 +274,6 @@ public static EagerSession getDefault() {
       synchronized (EagerSession.class) {
         if (defaultSession == null) {
           defaultSession = options().build();
-
-          Runtime.getRuntime().addShutdownHook(new Thread() {
-            @Override
-            public void run() {
-              defaultSession.doClose();
-            }
-          });
         }
       }
     }
@@ -495,7 +488,13 @@ void stopCleanupThread() {
       cleanupService.shutdownNow(); // returns without waiting for the thread to stop
     }
 
-    private final ExecutorService cleanupService = Executors.newSingleThreadExecutor();
+    private final ExecutorService cleanupService = Executors.newSingleThreadExecutor(r -> {
+      Thread thread = Executors.defaultThreadFactory().newThread(r);
+      thread.setDaemon(true);
+      thread.setPriority(Thread.MAX_PRIORITY);
+      thread.setContextClassLoader(null);
+      return thread;
+    });
     private final Map<NativeReference, Void> nativeRefs = new IdentityHashMap<>();
     private final ReferenceQueue<Object> garbageQueue;
     private volatile boolean cleanupInBackground = false;