Add a hook for cache invalidation

This is an additional Guice object that the module may or may not inject.
trinodb · Jul 27, 2023 · 35c091b · 35c091b
1 parent 8798669
commit 35c091b
Show file tree

Hide file tree

Showing 5 changed files with 176 additions and 22 deletions.
diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/group/CachingGroupProvider.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/group/CachingGroupProvider.java
@@ -26,7 +26,7 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 public class CachingGroupProvider
-        implements GroupProvider
+        implements GroupProvider, GroupCacheInvalidationListener
 {
     private final LoadingCache<String, Set<String>> cache;
 
@@ -46,4 +46,16 @@ public Set<String> getGroups(String user)
     {
         return cache.getUnchecked(user);
     }
+
+    @Override
+    public void invalidate(String user)
+    {
+        cache.invalidate(user);
+    }
+
+    @Override
+    public void invalidateAll()
+    {
+        cache.invalidateAll();
+    }
 }
diff --git a/...o-plugin-toolkit/src/main/java/io/trino/plugin/base/group/CachingGroupProviderModule.java b/...o-plugin-toolkit/src/main/java/io/trino/plugin/base/group/CachingGroupProviderModule.java
@@ -51,6 +51,9 @@
  * </ul>
  * These properties can optionally have an arbitrary prefix ({@link Builder#withPrefix(String)})
  * and/or a binding annotation for the resulting binding of {@link GroupProvider} ({@link Builder#withBindingAnnotation(Class)}).
+ * <p>
+ * An additional object of type {@link GroupCacheInvalidationListener} will also be bound, with which one can invalidate
+ * all or part of the cache.
  */
 public class CachingGroupProviderModule
         extends AbstractConfigurationAwareModule
@@ -104,6 +107,9 @@ public void configure(Binder binder)
                             .orElseGet(() -> Key.get(GroupProvider.class)))
                     .to(CachingGroupProvider.class)
                     .in(Scopes.SINGLETON);
+            binder.bind(GroupCacheInvalidationListener.class)
+                    .to(CachingGroupProvider.class)
+                    .in(Scopes.SINGLETON);
         }
     }
 
@@ -125,6 +131,9 @@ public void configure(Binder binder)
                             .orElseGet(() -> Key.get(GroupProvider.class)))
                     .to(Key.get(GroupProvider.class, ForCachingGroupProvider.class))
                     .in(Scopes.SINGLETON);
+            binder.bind(GroupCacheInvalidationListener.class)
+                    .to(NoOpGroupCacheInvalidationListener.class)
+                    .in(Scopes.SINGLETON);
         }
     }
 

diff --git a/...ugin-toolkit/src/main/java/io/trino/plugin/base/group/GroupCacheInvalidationListener.java b/...ugin-toolkit/src/main/java/io/trino/plugin/base/group/GroupCacheInvalidationListener.java
@@ -0,0 +1,21 @@
+/*
+ * 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.plugin.base.group;
+
+public interface GroupCacheInvalidationListener
+{
+    void invalidate(String user);
+
+    void invalidateAll();
+}
diff --git a/...-toolkit/src/main/java/io/trino/plugin/base/group/NoOpGroupCacheInvalidationListener.java b/...-toolkit/src/main/java/io/trino/plugin/base/group/NoOpGroupCacheInvalidationListener.java
@@ -0,0 +1,28 @@
+/*
+ * 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.plugin.base.group;
+
+public class NoOpGroupCacheInvalidationListener
+        implements GroupCacheInvalidationListener
+{
+    @Override
+    public void invalidate(String user)
+    {
+    }
+
+    @Override
+    public void invalidateAll()
+    {
+    }
+}
diff --git a/...ino-plugin-toolkit/src/test/java/io/trino/plugin/base/group/TestCachingGroupProvider.java b/...ino-plugin-toolkit/src/test/java/io/trino/plugin/base/group/TestCachingGroupProvider.java
@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.BindingAnnotation;
+import com.google.inject.Inject;
 import com.google.inject.Injector;
 import com.google.inject.Key;
 import io.airlift.bootstrap.Bootstrap;
@@ -47,7 +48,7 @@ public void testWithOutCaching()
         CountingGroupProvider countingGroupProvider = new CountingGroupProvider();
         ImmutableMap<String, String> properties = ImmutableMap.of(
                 "cache.enabled", "false");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.empty()).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.empty()).create(properties);
 
         innerTestWithOutCaching(countingGroupProvider, groupProvider);
     }
@@ -58,7 +59,7 @@ public void testWithOutCachingWithBindingAnnotation()
         CountingGroupProvider countingGroupProvider = new CountingGroupProvider();
         ImmutableMap<String, String> properties = ImmutableMap.of(
                 "cache.enabled", "false");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.of(ForTesting.class)).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.of(ForTesting.class)).create(properties);
 
         innerTestWithOutCaching(countingGroupProvider, groupProvider);
     }
@@ -69,7 +70,7 @@ public void testWithOutCachingWithPrefix()
         CountingGroupProvider countingGroupProvider = new CountingGroupProvider();
         ImmutableMap<String, String> properties = ImmutableMap.of(
                 "group-provider.cache.enabled", "false");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.empty()).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.empty()).create(properties);
 
         innerTestWithOutCaching(countingGroupProvider, groupProvider);
     }
@@ -80,18 +81,42 @@ public void testWithOutCachingWithPrefixWithBindingAnnotation()
         CountingGroupProvider countingGroupProvider = new CountingGroupProvider();
         ImmutableMap<String, String> properties = ImmutableMap.of(
                 "group-provider.cache.enabled", "false");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.of(ForTesting.class)).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.of(ForTesting.class)).create(properties);
 
         innerTestWithOutCaching(countingGroupProvider, groupProvider);
     }
 
-    private static void innerTestWithOutCaching(CountingGroupProvider countingGroupProvider, GroupProvider groupProvider)
+    private static void innerTestWithOutCaching(CountingGroupProvider countingGroupProvider, TestingGroupProvider groupProvider)
     {
         assertThat(countingGroupProvider.getCount()).isEqualTo(0);
-        assertThat(groupProvider.getGroups("testUser")).isEqualTo(ImmutableSet.of("test", "testUser"));
+
+        // first batch
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
         assertThat(countingGroupProvider.getCount()).isEqualTo(1);
-        assertThat(groupProvider.getGroups("testUser")).isEqualTo(ImmutableSet.of("test", "testUser"));
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
         assertThat(countingGroupProvider.getCount()).isEqualTo(2);
+
+        // second batch
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(3);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(4);
+
+        // invalidate user
+        groupProvider.invalidate("testUser1");
+        // no effect:
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(5);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(6);
+
+        // invalidate all
+        groupProvider.invalidateAll();
+        // no effect:
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(7);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(8);
     }
 
     @Test
@@ -101,7 +126,7 @@ public void testWithCaching()
         Map<String, String> properties = ImmutableMap.of(
                 "cache.enabled", "true",
                 "cache.ttl", "1 h");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.empty()).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.empty()).create(properties);
 
         innerTestWithCaching(countingGroupProvider, groupProvider);
     }
@@ -113,7 +138,7 @@ public void testWithCachingWithBindingAnnotation()
         Map<String, String> properties = ImmutableMap.of(
                 "cache.enabled", "true",
                 "cache.ttl", "1 h");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.of(ForTesting.class)).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.empty(), Optional.of(ForTesting.class)).create(properties);
 
         innerTestWithCaching(countingGroupProvider, groupProvider);
     }
@@ -125,7 +150,7 @@ public void testWithCachingWithPrefix()
         Map<String, String> properties = ImmutableMap.of(
                 "group-provider.cache.enabled", "true",
                 "group-provider.cache.ttl", "1 h");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.empty()).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.empty()).create(properties);
 
         innerTestWithCaching(countingGroupProvider, groupProvider);
     }
@@ -137,19 +162,42 @@ public void testWithCachingWithPrefixWithBindingAnnotation()
         Map<String, String> properties = ImmutableMap.of(
                 "group-provider.cache.enabled", "true",
                 "group-provider.cache.ttl", "1 h");
-        GroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.of(ForTesting.class)).create(properties);
+        TestingGroupProvider groupProvider = new TestingGroupProviderFactory(countingGroupProvider, Optional.of("group-provider"), Optional.of(ForTesting.class)).create(properties);
 
         innerTestWithCaching(countingGroupProvider, groupProvider);
     }
 
-    private static void innerTestWithCaching(CountingGroupProvider countingGroupProvider, GroupProvider groupProvider)
+    private static void innerTestWithCaching(CountingGroupProvider countingGroupProvider, TestingGroupProvider groupProvider)
     {
         assertThat(countingGroupProvider.getCount()).isEqualTo(0);
-        assertThat(groupProvider.getGroups("testUser")).isEqualTo(ImmutableSet.of("test", "testUser"));
-        assertThat(countingGroupProvider.getCount()).isEqualTo(1);
-        assertThat(groupProvider.getGroups("testUser")).isEqualTo(ImmutableSet.of("test", "testUser"));
-        // second call is handled by the cache so delegate not invoked
+
+        // first batch
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
         assertThat(countingGroupProvider.getCount()).isEqualTo(1);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(2);
+
+        // second batch is handled by the cache so delegate not invoked
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(2);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(2);
+
+        // invalidate user
+        groupProvider.invalidate("testUser1");
+        // effect on testUser1 only:
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(3);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(3);
+
+        // invalidate all
+        groupProvider.invalidateAll();
+        // effect on both:
+        assertThat(groupProvider.getGroups("testUser1")).isEqualTo(ImmutableSet.of("test", "testUser1"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(4);
+        assertThat(groupProvider.getGroups("testUser2")).isEqualTo(ImmutableSet.of("test", "testUser2"));
+        assertThat(countingGroupProvider.getCount()).isEqualTo(5);
     }
 
     private static class CountingGroupProvider
@@ -191,24 +239,60 @@ public String getName()
         }
 
         @Override
-        public GroupProvider create(Map<String, String> config)
+        public TestingGroupProvider create(Map<String, String> config)
         {
             CachingGroupProviderModule.Builder moduleBuilder = CachingGroupProviderModule.builder();
             prefix.ifPresent(moduleBuilder::withPrefix);
             bindingAnnotation.ifPresent(moduleBuilder::withBindingAnnotation);
 
             Bootstrap app = new Bootstrap(
                     moduleBuilder.build(),
-                    binder -> binder.bind(Key.get(GroupProvider.class, ForCachingGroupProvider.class))
-                            .toInstance(groupProvider));
+                    binder -> {
+                        binder.bind(Key.get(GroupProvider.class, ForCachingGroupProvider.class))
+                                .toInstance(groupProvider);
+                        bindingAnnotation.ifPresent(bindingAnnotation ->
+                                binder.bind(GroupProvider.class).to(Key.get(GroupProvider.class, bindingAnnotation)));
+                        binder.bind(TestingGroupProvider.class);
+                    });
 
             Injector injector = app
                     .doNotInitializeLogging()
                     .setRequiredConfigurationProperties(config)
                     .initialize();
 
-            return bindingAnnotation.map(bindingAnnotation -> injector.getInstance(Key.get(GroupProvider.class, bindingAnnotation)))
-                    .orElseGet(() -> injector.getInstance(GroupProvider.class));
+            return injector.getInstance(TestingGroupProvider.class);
+        }
+    }
+
+    private static class TestingGroupProvider
+            implements GroupProvider, GroupCacheInvalidationListener
+    {
+        private final GroupProvider delegate;
+        private final GroupCacheInvalidationListener invalidationListener;
+
+        @Inject
+        public TestingGroupProvider(GroupProvider delegate, GroupCacheInvalidationListener invalidationListener)
+        {
+            this.delegate = delegate;
+            this.invalidationListener = invalidationListener;
+        }
+
+        @Override
+        public Set<String> getGroups(String user)
+        {
+            return delegate.getGroups(user);
+        }
+
+        @Override
+        public void invalidate(String user)
+        {
+            invalidationListener.invalidate(user);
+        }
+
+        @Override
+        public void invalidateAll()
+        {
+            invalidationListener.invalidateAll();
         }
     }