From 02f9b255e6f370ac0cb0b1c5a8751eeebd26725b Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Mon, 28 Oct 2019 11:48:44 +0300
Subject: [PATCH] Clean using thread local in all threads (KT-28940)

---
 .../kotlin/storage/LockBasedStorageManager.java    | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java b/core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java
index 8d939ee41e4..dce530a67a5 100644
--- a/core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java
+++ b/core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java
@@ -418,10 +418,14 @@ public class LockBasedStorageManager implements StorageManager {
         public T invoke() {
             ThreadLocal<Object> postComputeCache = valuePostCompute;
             if (postComputeCache != null) {
-                Object _value = postComputeCache.get();
-                if (!(_value instanceof NotValue)) {
-                    // This thread is counting the value so allow an early publication
-                    return WrappedValues.unescapeThrowable(_value);
+                try {
+                    Object _value = postComputeCache.get();
+                    if (!(_value instanceof NotValue)) {
+                        // This thread is counting the value so allow an early publication
+                        return WrappedValues.unescapeThrowable(_value);
+                    }
+                } finally {
+                    postComputeCache.remove();
                 }
             }
 
@@ -444,8 +448,8 @@ public class LockBasedStorageManager implements StorageManager {
                 valuePostCompute = postComputeCache;
                 doPostCompute(value);
             } finally {
-                postComputeCache.remove();
                 valuePostCompute = null;
+                postComputeCache.remove();
             }
         }