Handle map representation of cache in SafeCaches

core/trino-main/src/main/java/io/trino/execution/SqlTaskManager.java

@@ -592,7 +592,7 @@ void removeOldTasks()
                         if (endTime != null && endTime.isBefore(oldestAllowedTask)) {
                             // The removal here is concurrency safe with respect to any concurrent loads: the cache has no expiration,
                             // the taskId is in the cache, so there mustn't be an ongoing load.
-                            tasks.asMap().remove(taskId);
+                            tasks.unsafeInvalidate(taskId);
                         }
                     }
                     catch (RuntimeException e) {

lib/trino-collect/src/main/java/io/trino/collect/cache/NonEvictableCacheImpl.java

@@ -14,6 +14,9 @@
 package io.trino.collect.cache;
 
 import com.google.common.cache.Cache;
+import com.google.common.collect.ForwardingConcurrentMap;
+
+import java.util.concurrent.ConcurrentMap;
 
 // package-private. The interface provides deprecation and javadoc to help at call sites
 final class NonEvictableCacheImpl<K, V>
@@ -32,4 +35,25 @@ public void invalidateAll()
                 "Use EvictableCache if you need invalidation, or use SafeCaches.buildNonEvictableCacheWithWeakInvalidateAll() " +
                 "if invalidateAll is not required for correctness");
     }
+
+    @Override
+    public ConcurrentMap<K, V> asMap()
+    {
+        ConcurrentMap<K, V> map = super.asMap();
+        return new ForwardingConcurrentMap<K, V>()
+        {
+            @Override
+            protected ConcurrentMap<K, V> delegate()
+            {
+                return map;
+            }
+
+            @Override
+            public void clear()
+            {
+                throw new UnsupportedOperationException("clear() does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+        };
+    }
 }

lib/trino-collect/src/main/java/io/trino/collect/cache/NonEvictableLoadingCacheImpl.java

@@ -14,6 +14,9 @@
 package io.trino.collect.cache;
 
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ForwardingConcurrentMap;
+
+import java.util.concurrent.ConcurrentMap;
 
 // package-private. The interface provides deprecation and javadoc to help at call sites
 final class NonEvictableLoadingCacheImpl<K, V>
@@ -29,7 +32,28 @@ final class NonEvictableLoadingCacheImpl<K, V>
     public void invalidateAll()
     {
         throw new UnsupportedOperationException("invalidateAll does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
-                "Use EvictableCacheBuilder if you need invalidation, or use SafeCaches.buildNonEvictableCacheWithWeakInvalidateAll() " +
+                "Use EvictableCache if you need invalidation, or use SafeCaches.buildNonEvictableCacheWithWeakInvalidateAll() " +
                 "if invalidateAll is not required for correctness");
     }
+
+    @Override
+    public ConcurrentMap<K, V> asMap()
+    {
+        ConcurrentMap<K, V> map = super.asMap();
+        return new ForwardingConcurrentMap<K, V>()
+        {
+            @Override
+            protected ConcurrentMap<K, V> delegate()
+            {
+                return map;
+            }
+
+            @Override
+            public void clear()
+            {
+                throw new UnsupportedOperationException("clear() does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+        };
+    }
 }

lib/trino-collect/src/main/java/io/trino/collect/cache/NonKeyEvictableCacheImpl.java

@@ -15,6 +15,9 @@
 
 import com.google.common.cache.Cache;
 import com.google.common.cache.ForwardingCache;
+import com.google.common.collect.ForwardingConcurrentMap;
+
+import java.util.concurrent.ConcurrentMap;
 
 import static java.util.Objects.requireNonNull;
 
@@ -49,4 +52,49 @@ public void invalidateAll(Iterable<?> keys)
         throw new UnsupportedOperationException("invalidateAll(keys) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
                 "Use EvictableCache if you need invalidation");
     }
+
+    @Override
+    public ConcurrentMap<K, V> asMap()
+    {
+        ConcurrentMap<K, V> map = delegate.asMap();
+        return new ForwardingConcurrentMap<K, V>()
+        {
+            @Override
+            protected ConcurrentMap<K, V> delegate()
+            {
+                return map;
+            }
+
+            @Override
+            public V remove(Object key)
+            {
+                throw new UnsupportedOperationException("remove(key) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+
+            @Override
+            public boolean remove(Object key, Object value)
+            {
+                // defensive forbid, it is not sure if this method is safe to call or not
+                throw new UnsupportedOperationException("remove(key, value) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+
+            @Override
+            public boolean replace(K key, V oldValue, V newValue)
+            {
+                // defensive forbid, it is not sure if this method is safe to call or not
+                throw new UnsupportedOperationException("replace(key, oldValue, newValue) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+
+            @Override
+            public V replace(K key, V value)
+            {
+                // defensive forbid, it is not sure if this method is safe to call or not
+                throw new UnsupportedOperationException("replace(key, value) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+        };
+    }
 }

lib/trino-collect/src/main/java/io/trino/collect/cache/NonKeyEvictableLoadingCache.java

@@ -28,6 +28,15 @@ public interface NonKeyEvictableLoadingCache<K, V>
     @Override
     void invalidate(Object key);
 
+    /**
+     * Allows invalidation for a key, but does not invalidate ongoing loads.
+     * It is caller responsibility to ensure correct use.
+     * The method is deprecated to discourage the usage of it.
+     * Use {@link EvictableCacheBuilder} to build a cache instead.
+     */
+    @Deprecated
+    void unsafeInvalidate(Object key);
+
     /**
      * @deprecated Not supported. Use {@link EvictableCacheBuilder} to build a cache instead.
      */

lib/trino-collect/src/main/java/io/trino/collect/cache/NonKeyEvictableLoadingCacheImpl.java

@@ -15,6 +15,9 @@
 
 import com.google.common.cache.ForwardingLoadingCache;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ForwardingConcurrentMap;
+
+import java.util.concurrent.ConcurrentMap;
 
 import static java.util.Objects.requireNonNull;
 
@@ -40,13 +43,61 @@ protected LoadingCache<K, V> delegate()
     public void invalidate(Object key)
     {
         throw new UnsupportedOperationException("invalidate(key) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
-                "Use EvictableCacheBuilder if you need invalidation");
+                "Use EvictableCache if you need invalidation");
+    }
+
+    @Override
+    public void unsafeInvalidate(Object key)
+    {
+        super.invalidate(key);
     }
 
     @Override
     public void invalidateAll(Iterable<?> keys)
     {
         throw new UnsupportedOperationException("invalidateAll(keys) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
-                "Use EvictableCacheBuilder if you need invalidation");
+                "Use EvictableCache if you need invalidation");
+    }
+
+    @Override
+    public ConcurrentMap<K, V> asMap()
+    {
+        ConcurrentMap<K, V> map = delegate.asMap();
+        return new ForwardingConcurrentMap<K, V>()
+        {
+            @Override
+            protected ConcurrentMap<K, V> delegate()
+            {
+                return map;
+            }
+
+            @Override
+            public V remove(Object key)
+            {
+                throw new UnsupportedOperationException("remove(key) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+
+            @Override
+            public boolean remove(Object key, Object value)
+            {
+                throw new UnsupportedOperationException("remove(key, value) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+
+            @Override
+            public boolean replace(K key, V oldValue, V newValue)
+            {
+                throw new UnsupportedOperationException("replace(key, oldValue, newValue) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+
+            @Override
+            public V replace(K key, V value)
+            {
+                throw new UnsupportedOperationException("replace(key, value) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. " +
+                        "Use EvictableCacheBuilder if you need invalidation");
+            }
+        };
     }
 }

lib/trino-collect/src/test/java/io/trino/collect/cache/TestSafeCaches.java

@@ -0,0 +1,132 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.collect.cache;
+
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.collect.ImmutableList;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static io.trino.collect.cache.SafeCaches.buildNonEvictableCache;
+import static io.trino.collect.cache.SafeCaches.buildNonEvictableCacheWithWeakInvalidateAll;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.testng.Assert.assertSame;
+
+public class TestSafeCaches
+{
+    @Test
+    public void testNonEvictableCache()
+            throws Exception
+    {
+        NonEvictableCache<Object, Object> cache = buildNonEvictableCache(CacheBuilder.newBuilder());
+        verifyKeyInvalidationIsImpossible(cache);
+        verifyClearIsImpossible(cache);
+        verifyLoadingIsPossible(cache);
+    }
+
+    @Test
+    public void testNonEvictableCacheWithWeakInvalidateAll()
+            throws Exception
+    {
+        NonKeyEvictableCache<Object, Object> cache = buildNonEvictableCacheWithWeakInvalidateAll(CacheBuilder.newBuilder());
+        verifyKeyInvalidationIsImpossible(cache);
+        verifyClearIsPossible(cache);
+        verifyLoadingIsPossible(cache);
+    }
+
+    @Test
+    public void testNonEvictableLoadingCache()
+            throws Exception
+    {
+        NonEvictableLoadingCache<Object, Object> cache = buildNonEvictableCache(CacheBuilder.newBuilder(), CacheLoader.from(key -> key));
+        verifyKeyInvalidationIsImpossible(cache);
+        verifyClearIsImpossible(cache);
+        verifyLoadingIsPossible(cache);
+    }
+
+    @Test
+    public void testNonEvictableLoadingCacheWithWeakInvalidateAll()
+            throws Exception
+    {
+        NonKeyEvictableLoadingCache<Object, Object> cache = buildNonEvictableCacheWithWeakInvalidateAll(CacheBuilder.newBuilder(), CacheLoader.from(key -> key));
+        verifyKeyInvalidationIsImpossible(cache);
+        verifyClearIsPossible(cache);
+        verifyLoadingIsPossible(cache);
+    }
+
+    private static void verifyLoadingIsPossible(Cache<Object, Object> cache)
+            throws Exception
+    {
+        Object key = new Object();
+        Object value = new Object();
+        // Verify the previous load was inserted into the cache
+        assertSame(cache.get(key, () -> value), value);
+    }
+
+    private static void verifyKeyInvalidationIsImpossible(Cache<Object, Object> cache)
+    {
+        assertThatThrownBy(() -> cache.invalidate(new Object()))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("invalidate(key) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCache if you need invalidation");
+        assertThatThrownBy(() -> cache.invalidateAll(ImmutableList.of()))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("invalidateAll(keys) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCache if you need invalidation");
+        Object object = new Object();
+        assertThatThrownBy(() -> cache.asMap().remove(object))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("remove(key) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCacheBuilder if you need invalidation");
+        assertThatThrownBy(() -> cache.asMap().remove(object, object))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("remove(key, value) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCacheBuilder if you need invalidation");
+        assertThatThrownBy(() -> cache.asMap().replace(object, object))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("replace(key, value) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCacheBuilder if you need invalidation");
+        assertThatThrownBy(() -> cache.asMap().replace(object, object, object))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("replace(key, oldValue, newValue) does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCacheBuilder if you need invalidation");
+    }
+
+    private static void verifyClearIsImpossible(Cache<Object, Object> cache)
+    {
+        assertThatThrownBy(cache::invalidateAll)
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("invalidateAll does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCache if you need invalidation, or use SafeCaches.buildNonEvictableCacheWithWeakInvalidateAll() if invalidateAll is not required for correctness");
+        Map<Object, Object> map = cache.asMap();
+        assertThatThrownBy(map::clear)
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("clear() does not invalidate ongoing loads, so a stale value may remain in the cache for ever. Use EvictableCacheBuilder if you need invalidation");
+    }
+
+    private static void verifyClearIsPossible(Cache<Object, Object> cache)
+            throws Exception
+    {
+        Object key = new Object();
+        Object firstValue = new Object();
+        cache.get(key, () -> firstValue);
+        assertSame(cache.getIfPresent(key), firstValue);
+
+        cache.invalidateAll();
+        assertThat(cache.getIfPresent(key)).isNull();
+
+        Object secondValue = new Object();
+        cache.get(key, () -> secondValue);
+        assertSame(cache.getIfPresent(key), secondValue);
+        cache.asMap().clear();
+        assertThat(cache.getIfPresent(key)).isNull();
+    }
+}