From 620898c73e8278c666f817d0e69e7f4ee21eca45 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Mon, 28 Oct 2019 12:08:29 +0300
Subject: [PATCH] Use local single-thread cached value instead of thread locals
 (KT-28940)

It looks like thread locals are not needed here as their global state
is not utilized but creates troubles with memory management.
---
 .../storage/LockBasedStorageManager.java      | 34 ++++++------------
 .../kotlin/storage/SingleThreadValue.java     | 36 +++++++++++++++++++
 2 files changed, 47 insertions(+), 23 deletions(-)
 create mode 100644 core/util.runtime/src/org/jetbrains/kotlin/storage/SingleThreadValue.java

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 dce530a67a5..e92c234d8a5 100644
--- a/core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java
+++ b/core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java
@@ -403,9 +403,13 @@ public class LockBasedStorageManager implements StorageManager {
         }
     }
 
+    /**
+     * Computed value has an early publication and accessible from the same thread while executing a post-compute lambda.
+     * For other threads value will be accessible only after post-compute lambda is finished (when a real lock is used).
+     */
     private static abstract class LockBasedLazyValueWithPostCompute<T> extends LockBasedLazyValue<T> {
         @Nullable
-        private volatile ThreadLocal<Object> valuePostCompute = null;
+        private volatile SingleThreadValue<T> valuePostCompute = null;
 
         public LockBasedLazyValueWithPostCompute(
                 @NotNull LockBasedStorageManager storageManager,
@@ -416,17 +420,9 @@ public class LockBasedStorageManager implements StorageManager {
 
         @Override
         public T invoke() {
-            ThreadLocal<Object> postComputeCache = valuePostCompute;
-            if (postComputeCache != null) {
-                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();
-                }
+            SingleThreadValue<T> postComputeCache = valuePostCompute;
+            if (postComputeCache != null && postComputeCache.hasValue()) {
+                return postComputeCache.getValue();
             }
 
             return super.invoke();
@@ -435,21 +431,13 @@ public class LockBasedStorageManager implements StorageManager {
         // Doing something in post-compute helps prevent infinite recursion
         @Override
         protected final void postCompute(T value) {
-            ThreadLocal<Object> postComputeCache = new ThreadLocal<Object>() {
-                @Override
-                protected Object initialValue() {
-                    return NotValue.NOT_COMPUTED;
-                }
-            };
-
-            postComputeCache.set(value);
-
+            // Protected from rewrites in other threads because it is executed under lock in invoke().
+            // May be overwritten when NO_LOCK is used.
+            valuePostCompute = new SingleThreadValue<T>(value);
             try {
-                valuePostCompute = postComputeCache;
                 doPostCompute(value);
             } finally {
                 valuePostCompute = null;
-                postComputeCache.remove();
             }
         }
 
diff --git a/core/util.runtime/src/org/jetbrains/kotlin/storage/SingleThreadValue.java b/core/util.runtime/src/org/jetbrains/kotlin/storage/SingleThreadValue.java
new file mode 100644
index 00000000000..b688b4e1198
--- /dev/null
+++ b/core/util.runtime/src/org/jetbrains/kotlin/storage/SingleThreadValue.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.storage;
+
+/**
+ * A storage for the value that should exist and be accessible in the single thread.
+ *
+ * Unlike ThreadLocal, thread doesn't store a reference to the value that makes it inaccessible globally, but simplifies memory
+ * management.
+ *
+ * The other difference from ThreadLocal is inability to have different values per each thread, so SingleThreadValue instance
+ * should be protected with external lock from rewrites.
+ *
+ * @param <T>
+ */
+class SingleThreadValue<T> {
+    private final T value;
+    private final Thread thread;
+
+    SingleThreadValue(T value) {
+        this.value = value;
+        thread = Thread.currentThread();
+    }
+
+    public boolean hasValue() {
+        return thread == Thread.currentThread();
+    }
+
+    public T getValue() {
+        if (!hasValue()) throw new IllegalStateException("No value in this thread (hasValue should be checked before)");
+        return value;
+    }
+}