Use TruffleSafepoint.setBlockedThreadInterruptible for blocking operations in LRUCache

Author: steve-s

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/StructModuleBuiltins.java

@@ -78,12 +78,12 @@ public void postInitialize(Python3Core core) {
         structModule.setModuleState(cache);
     }
 
-    protected static PStruct getStruct(PythonModule structModule, Object format, StructBuiltins.ConstructStructNode constructStructNode) {
+    protected static PStruct getStruct(Node location, PythonModule structModule, Object format, StructBuiltins.ConstructStructNode constructStructNode) {
         LRUStructCache cache = structModule.getModuleState(LRUStructCache.class);
-        PStruct pStruct = cache.get(format);
+        PStruct pStruct = cache.get(location, format);
         if (pStruct == null) {
             pStruct = constructStructNode.execute(format);
-            cache.put(format, pStruct);
+            cache.put(location, format, pStruct);
         }
         return pStruct;
     }
@@ -94,8 +94,8 @@ protected static PStruct getStruct(PythonModule structModule, Object format, Str
     abstract static class GetStructNode extends PNodeWithContext {
         abstract PStruct execute(Node inliningTarget, PythonModule module, Object format, StructBuiltins.ConstructStructNode constructStructNode);
 
-        protected PStruct getStructInternal(PythonModule module, Object format, StructBuiltins.ConstructStructNode constructStructNode) {
-            return getStruct(module, format, constructStructNode);
+        protected PStruct getStructInternal(Node location, PythonModule module, Object format, StructBuiltins.ConstructStructNode constructStructNode) {
+            return getStruct(location, module, format, constructStructNode);
         }
 
         protected boolean eq(TruffleString s1, TruffleString s2, TruffleString.EqualNode eqNode) {
@@ -107,7 +107,7 @@ protected boolean eq(TruffleString s1, TruffleString s2, TruffleString.EqualNode
         static PStruct doCachedString(PythonModule module, TruffleString format, StructBuiltins.ConstructStructNode constructStructNode,
                         @Cached("format") TruffleString cachedFormat,
                         @Cached TruffleString.EqualNode eqNode,
-                        @Cached(value = "getStructInternal(module, format, constructStructNode)", weak = true) PStruct cachedStruct) {
+                        @Cached(value = "getStructInternal($node, module, format, constructStructNode)", weak = true) PStruct cachedStruct) {
             return cachedStruct;
         }
 
@@ -116,13 +116,14 @@ static PStruct doCachedString(PythonModule module, TruffleString format, StructB
         static PStruct doCachedBytes(PythonModule module, PBytes format, StructBuiltins.ConstructStructNode constructStructNode,
                         @CachedLibrary("format") PythonBufferAccessLibrary bufferLib,
                         @Cached(value = "bufferLib.getCopiedByteArray(format)", dimensions = 1) byte[] cachedFormat,
-                        @Cached(value = "getStructInternal(module, format, constructStructNode)", weak = true) PStruct cachedStruct) {
+                        @Cached(value = "getStructInternal($node, module, format, constructStructNode)", weak = true) PStruct cachedStruct) {
             return cachedStruct;
         }
 
         @Specialization(replaces = {"doCachedString", "doCachedBytes"})
-        static PStruct doGeneric(PythonModule module, Object format, StructBuiltins.ConstructStructNode constructStructNode) {
-            return getStruct(module, format, constructStructNode);
+        static PStruct doGeneric(PythonModule module, Object format, StructBuiltins.ConstructStructNode constructStructNode,
+                        @Bind Node location) {
+            return getStruct(location, module, format, constructStructNode);
         }
     }
 
@@ -241,7 +242,7 @@ abstract static class ClearCacheNode extends PythonUnaryBuiltinNode {
         @Specialization
         Object clearCache(PythonModule self) {
             LRUStructCache cache = self.getModuleState(LRUStructCache.class);
-            cache.clear();
+            cache.clear(this);
             return PNone.NONE;
         }
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/util/Cache.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * The Universal Permissive License (UPL), Version 1.0
@@ -46,9 +46,11 @@
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
-import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import com.oracle.truffle.api.TruffleSafepoint;
+import com.oracle.truffle.api.nodes.Node;
 
-public class LRUCache<K, V> implements Cache<K, V> {
+public class LRUCache<K, V> {
     private final LRUHashMap<K, V> lruHashMap;
     private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
     private final Lock read = readWriteLock.readLock();
@@ -73,47 +75,33 @@ protected boolean removeEldestEntry(Map.Entry<K, V> eldest) {
         }
     }
 
-    @Override
-    @CompilerDirectives.TruffleBoundary
-    public V get(K key) {
+    @TruffleBoundary
+    public V get(Node location, K key) {
         try {
-            read.lock();
+            TruffleSafepoint.setBlockedThreadInterruptible(location, Lock::lockInterruptibly, read);
             return lruHashMap.get(key);
         } finally {
             read.unlock();
         }
     }
 
-    @Override
-    @CompilerDirectives.TruffleBoundary
-    public V put(K key, V value) {
+    @TruffleBoundary
+    public V put(Node location, K key, V value) {
         try {
-            write.lock();
+            TruffleSafepoint.setBlockedThreadInterruptible(location, Lock::lockInterruptibly, write);
             return lruHashMap.putIfAbsent(key, value);
         } finally {
             write.unlock();
         }
     }
 
-    @Override
-    @CompilerDirectives.TruffleBoundary
-    public void clear() {
+    @TruffleBoundary
+    public void clear(Node location) {
         try {
-            write.lock();
+            TruffleSafepoint.setBlockedThreadInterruptible(location, Lock::lockInterruptibly, write);
             lruHashMap.clear();
         } finally {
             write.unlock();
         }
     }
-
-    @Override
-    @CompilerDirectives.TruffleBoundary
-    public int size() {
-        try {
-            read.lock();
-            return lruHashMap.size();
-        } finally {
-            read.unlock();
-        }
-    }
 }