Test delete entry in remote Infinispan and check removals are replica…

…ted to the second site anyway Signed-off-by: Michal Hajas <[email protected]>
keycloak · May 20, 2024 · d556770 · d556770
1 parent 77b7849
commit d556770
Show file tree

Hide file tree

Showing 10 changed files with 267 additions and 32 deletions.
diff --git a/dataset/src/main/java/org/keycloak/benchmark/cache/CacheResource.java b/dataset/src/main/java/org/keycloak/benchmark/cache/CacheResource.java
@@ -84,4 +84,12 @@ public boolean contains(@PathParam("id") String id) {
             return false;
         }
     }
+
+    @GET
+    @Path("/size")
+    @NoCache
+    @Produces(MediaType.APPLICATION_JSON)
+    public int size() {
+        return cache.size();
+    }
 }
diff --git a/dataset/src/main/java/org/keycloak/benchmark/cache/CacheResourceProvider.java b/dataset/src/main/java/org/keycloak/benchmark/cache/CacheResourceProvider.java
@@ -29,6 +29,7 @@
 
 import org.infinispan.Cache;
 import org.infinispan.client.hotrod.RemoteCache;
+import org.infinispan.context.Flag;
 import org.infinispan.manager.EmbeddedCacheManager;
 import org.jboss.resteasy.reactive.NoCache;
 import org.keycloak.benchmark.dataset.TaskResponse;

diff --git a/dataset/src/main/java/org/keycloak/benchmark/cache/RemoteCacheResource.java b/dataset/src/main/java/org/keycloak/benchmark/cache/RemoteCacheResource.java
@@ -20,13 +20,16 @@
 
 import java.util.UUID;
 
+import jakarta.ws.rs.DefaultValue;
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.NotFoundException;
 import jakarta.ws.rs.Path;
 import jakarta.ws.rs.PathParam;
 import jakarta.ws.rs.Produces;
 
+import jakarta.ws.rs.QueryParam;
 import org.infinispan.Cache;
+import org.infinispan.client.hotrod.Flag;
 import org.infinispan.client.hotrod.RemoteCache;
 import org.infinispan.commons.CacheConfigurationException;
 import org.jboss.logging.Logger;
@@ -94,4 +97,30 @@ public boolean contains(@PathParam("id") String id) {
             return false;
         }
     }
+
+    @GET
+    @Path("/remove/{id}")
+    @NoCache
+    @Produces(MediaType.APPLICATION_JSON)
+    public boolean remove(@PathParam("id") String id, @QueryParam("skipListeners") @DefaultValue("false") boolean skipListeners) {
+        if (decorateCacheForRemovalAndSkipListenersIfTrue(skipListeners).remove(id) != null) {
+            return true;
+        } else if (id.length() == 36) {
+            try {
+                UUID uuid = UUID.fromString(id);
+                return decorateCacheForRemovalAndSkipListenersIfTrue(skipListeners).remove(uuid) != null;
+            } catch (IllegalArgumentException iae) {
+                logger.warnf("Given string %s not an UUID", id);
+                return false;
+            }
+        } else {
+            return false;
+        }
+    }
+
+    public RemoteCache decorateCacheForRemovalAndSkipListenersIfTrue(boolean skipListeners) {
+        return skipListeners
+                ? remoteCache.withFlags(Flag.SKIP_LISTENER_NOTIFICATION, Flag.FORCE_RETURN_VALUE)
+                : remoteCache.withFlags(Flag.FORCE_RETURN_VALUE);
+    }
 }
diff --git a/...hmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/AbstractCrossDCTest.java b/...hmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/AbstractCrossDCTest.java
@@ -1,17 +1,6 @@
 package org.keycloak.benchmark.crossdc;
 
-import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.keycloak.benchmark.crossdc.util.HttpClientUtils.MOCK_COOKIE_MANAGER;
-import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.DISTRIBUTED_CACHES;
-
-import java.io.IOException;
-import java.net.URISyntaxException;
-import java.net.UnknownHostException;
-import java.net.http.HttpClient;
-import java.util.Collections;
-import java.util.List;
-
+import jakarta.ws.rs.NotFoundException;
 import org.jboss.logging.Logger;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
@@ -28,8 +17,6 @@
 import org.keycloak.representations.idm.RealmRepresentation;
 import org.keycloak.representations.idm.UserRepresentation;
 
-import jakarta.ws.rs.NotFoundException;
-
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.net.UnknownHostException;
@@ -41,7 +28,9 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.keycloak.benchmark.crossdc.util.HttpClientUtils.MOCK_COOKIE_MANAGER;
+import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.CLIENT_SESSIONS;
 import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.DISTRIBUTED_CACHES;
+import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.SESSIONS;
 
 public abstract class AbstractCrossDCTest {
     private static final Logger LOG = Logger.getLogger(AbstractCrossDCTest.class);
@@ -119,6 +108,9 @@ public void setUpTestEnvironment() throws UnknownHostException {
         user.setCredentials(Collections.singletonList(credential));
 
         realmResource.users().create(user).close();
+
+        assertCacheSize(SESSIONS, 0);
+        assertCacheSize(CLIENT_SESSIONS, 0);
     }
 
     @AfterEach
@@ -161,11 +153,11 @@ private static void failbackHealthChecks() throws URISyntaxException, IOExceptio
 
     protected void assertCacheSize(String cache, int size) {
         // Embedded caches
-        assertEquals(DC_1.kc().embeddedIspn().cache(cache).size(), size, () -> "Embedded cache " + cache + " in DC1 has " + DC_1.ispn().cache(cache).size() + " entries");
-        assertEquals(DC_2.kc().embeddedIspn().cache(cache).size(), size, () -> "Embedded cache " + cache + " in DC2 has " + DC_2.ispn().cache(cache).size() + " entries");
+        assertEquals(size, DC_1.kc().embeddedIspn().cache(cache).size(), () -> "Embedded cache " + cache + " in DC1 has " + DC_1.ispn().cache(cache).size() + " entries");
+        assertEquals(size, DC_2.kc().embeddedIspn().cache(cache).size(), () -> "Embedded cache " + cache + " in DC2 has " + DC_2.ispn().cache(cache).size() + " entries");
 
         // External caches
-        assertEquals(DC_1.ispn().cache(cache).size(), size, () -> "External cache " + cache + " in DC1 has " + DC_1.ispn().cache(cache).size() + " entries");
-        assertEquals(DC_2.ispn().cache(cache).size(), size, () -> "External cache " + cache + " in DC2 has " + DC_2.ispn().cache(cache).size() + " entries");
+        assertEquals(size, DC_1.ispn().cache(cache).size(), () -> "External cache " + cache + " in DC1 has " + DC_1.ispn().cache(cache).size() + " entries");
+        assertEquals(size, DC_2.ispn().cache(cache).size(), () -> "External cache " + cache + " in DC2 has " + DC_2.ispn().cache(cache).size() + " entries");
     }
 }
diff --git a/...benchmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/LoginLogoutTest.java b/...benchmark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/LoginLogoutTest.java
@@ -1,13 +1,15 @@
 package org.keycloak.benchmark.crossdc;
 
 import org.junit.jupiter.api.Test;
+import org.keycloak.benchmark.crossdc.util.InfinispanUtils;
 
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.net.http.HttpResponse;
 import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.CLIENT_SESSIONS;
 import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.SESSIONS;
@@ -80,4 +82,70 @@ public void testRefreshTokenRevocation() throws Exception {
             assertCacheSize(CLIENT_SESSIONS, 0);
         }
     }
+
+    @Test
+    public void testRemoteStoreDiscrepancyMissingSessionInPrimaryRemoteISPN() throws URISyntaxException, IOException, InterruptedException {
+        // Create a new user session
+        Map<String, Object> tokensMap = LOAD_BALANCER_KEYCLOAK.passwordGrant(REALM_NAME, CLIENTID, USERNAME, MAIN_PASSWORD);
+
+        // Make sure all ISPNs can see the entry in the cache
+        assertTrue(DC_1.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_1.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+
+        // Remove the session from the remote store in DC1 only
+        try (var close = InfinispanUtils.withBackupDisabled(DC_1.ispn().cache(SESSIONS), DC_2.ispn().siteName())) {
+            assertFalse(DC_1.ispn().cache(SESSIONS).isBackupOnline(DC_2.ispn().siteName()));
+            DC_1.ispn().cache(SESSIONS).remove((String) tokensMap.get("session_state"));
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        assertTrue(DC_1.ispn().cache(SESSIONS).isBackupOnline(DC_2.ispn().siteName()));
+
+        assertTrue(DC_1.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_1.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+
+        LOAD_BALANCER_KEYCLOAK.logout(REALM_NAME, (String) tokensMap.get("id_token"), CLIENTID);
+
+        assertFalse(DC_1.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_1.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_2.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_2.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+    }
+
+    @Test
+    public void testRemoteStoreDiscrepancyMissingSessionInBackupRemoteISPN() throws URISyntaxException, IOException, InterruptedException {
+        // Create a new user session
+        Map<String, Object> tokensMap = LOAD_BALANCER_KEYCLOAK.passwordGrant(REALM_NAME, CLIENTID, USERNAME, MAIN_PASSWORD);
+
+        // Make sure all ISPNs can see the entry in the cache
+        assertTrue(DC_1.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_1.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+
+        // Remove the session from the remote store in DC1 only
+        try (var close = InfinispanUtils.withBackupDisabled(DC_2.ispn().cache(SESSIONS), DC_1.ispn().siteName())) {
+            assertFalse(DC_2.ispn().cache(SESSIONS).isBackupOnline(DC_1.ispn().siteName()));
+            DC_2.ispn().cache(SESSIONS).remove((String) tokensMap.get("session_state"));
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        assertTrue(DC_2.ispn().cache(SESSIONS).isBackupOnline(DC_1.ispn().siteName()));
+
+        assertTrue(DC_1.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_1.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertTrue(DC_2.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_2.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+
+        LOAD_BALANCER_KEYCLOAK.logout(REALM_NAME, (String) tokensMap.get("id_token"), CLIENTID);
+
+        assertFalse(DC_1.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_1.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_2.kc().embeddedIspn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+        assertFalse(DC_2.ispn().cache(SESSIONS).contains((String) tokensMap.get("session_state")));
+    }
 }
diff --git a/...ark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/client/DatacenterInfo.java b/...ark-crossdc-tests/src/test/java/org/keycloak/benchmark/crossdc/client/DatacenterInfo.java
@@ -14,7 +14,7 @@ public class DatacenterInfo {
 
     public DatacenterInfo(HttpClient httpClient, String keycloakServerURL, String infinispanServerURL) {
         this.keycloak = new KeycloakClient(httpClient, keycloakServerURL);
-        this.infinispan = new ExternalInfinispanClient(httpClient, infinispanServerURL, AbstractCrossDCTest.ISPN_USERNAME, AbstractCrossDCTest.MAIN_PASSWORD);
+        this.infinispan = new ExternalInfinispanClient(httpClient, infinispanServerURL, AbstractCrossDCTest.ISPN_USERNAME, AbstractCrossDCTest.MAIN_PASSWORD, keycloakServerURL);
 
         this.keycloakServerURL = keycloakServerURL;
         this.infinispanServerURL = infinispanServerURL;

diff --git a/...c-tests/src/test/java/org/keycloak/benchmark/crossdc/client/ExternalInfinispanClient.java b/...c-tests/src/test/java/org/keycloak/benchmark/crossdc/client/ExternalInfinispanClient.java
@@ -2,6 +2,7 @@
 
 import org.apache.http.client.utils.URIBuilder;
 import org.keycloak.benchmark.crossdc.util.InfinispanUtils;
+import org.keycloak.util.JsonSerialization;
 
 import java.io.IOException;
 import java.net.URI;
@@ -10,6 +11,7 @@
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.regex.Pattern;
@@ -20,24 +22,44 @@
 import static org.keycloak.benchmark.crossdc.AbstractCrossDCTest.ISPN_USERNAME;
 import static org.keycloak.benchmark.crossdc.AbstractCrossDCTest.MAIN_PASSWORD;
 import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.getBasicAuthenticationHeader;
+import static org.keycloak.benchmark.crossdc.util.InfinispanUtils.getNestedValue;
 
-public class ExternalInfinispanClient implements InfinispanClient {
+public class ExternalInfinispanClient implements InfinispanClient<InfinispanClient.ExternalCache> {
     private final HttpClient httpClient;
     private final String infinispanUrl;
     private final String username;
     private final String password;
+    private final String keycloakServerURL;
+    private final String siteName;
 
     Pattern UUID_REGEX = Pattern.compile("[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}");
 
-    public ExternalInfinispanClient(HttpClient httpClient, String infinispanUrl, String username, String password) {
+    public ExternalInfinispanClient(HttpClient httpClient, String infinispanUrl, String username, String password, String keycloakServerURL) {
         assertNotNull(infinispanUrl, "Infinispan URL cannot be null");
         this.httpClient = httpClient;
         this.infinispanUrl = infinispanUrl;
         this.username = username;
         this.password = password;
+        this.keycloakServerURL = keycloakServerURL;
+
+        HttpResponse<String> stringHttpResponse = sendRequestWithAction(infinispanUrl + "/rest/v2/cache-managers/default", "GET", null);
+        assertEquals(200, stringHttpResponse.statusCode());
+
+        Map<String, Object> returnedValues;
+        try {
+            returnedValues = JsonSerialization.readValue(stringHttpResponse.body(), Map.class);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+
+        this.siteName = (String) returnedValues.get("local_site");
     }
 
-    public class ExternalCache implements InfinispanClient.Cache {
+    public String siteName() {
+        return siteName;
+    }
+
+    public class ExternalCache implements InfinispanClient.ExternalCache {
 
         private final String cacheName;
 
@@ -102,8 +124,39 @@ public void clear() {
         }
 
         @Override
-        public boolean contains(String key) {
-            return keys().contains(key);
+        public boolean contains(String key) throws URISyntaxException, IOException, InterruptedException {
+            URI uri = new URIBuilder( keycloakServerURL + "/realms/master/remote-cache/" + cacheName + "/contains/" + key).build();
+            HttpRequest request = HttpRequest.newBuilder()
+                    .uri(uri)
+                    .GET()
+                    .build();
+
+            HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+            return Boolean.parseBoolean(response.body());
+        }
+
+        @Override
+        public boolean remove(String key) {
+            URI uri = null;
+            try {
+                uri = new URIBuilder( keycloakServerURL + "/realms/master/remote-cache/" + cacheName + "/remove/" + key + "?skipListeners=true").build();
+            } catch (URISyntaxException e) {
+                throw new RuntimeException(e);
+            }
+
+            HttpRequest request = HttpRequest.newBuilder()
+                    .uri(uri)
+                    .GET()
+                    .build();
+
+
+            HttpResponse<String> response = null;
+            try {
+                response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+            } catch (IOException | InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+            return Boolean.parseBoolean(response.body());
         }
 
         @Override
@@ -144,10 +197,60 @@ public Set<String> keys() {
                 throw new RuntimeException(e);
             }
         }
+
+        @Override
+        public void takeOffline(String backupSiteName) {
+            HttpResponse<String> stringHttpResponse = sendRequestWithAction(infinispanUrl + "/rest/v2/caches/" + cacheName + "/x-site/backups/" + backupSiteName, "POST", "take-offline");
+            assertEquals(200, stringHttpResponse.statusCode());
+        }
+
+        @Override
+        public void bringOnline(String backupSiteName) {
+            HttpResponse<String> stringHttpResponse = sendRequestWithAction(infinispanUrl + "/rest/v2/caches/" + cacheName + "/x-site/backups/" + backupSiteName, "POST", "bring-online");
+            assertEquals(200, stringHttpResponse.statusCode());
+        }
+
+        @Override
+        public boolean isBackupOnline(String backupSiteName) throws IOException {
+            String response = sendRequestWithAction(infinispanUrl + "/rest/v2/caches/" + cacheName + "/x-site/backups/", "GET", null).body();
+            Map<String, Object> returnedValues = JsonSerialization.readValue(response, Map.class);
+
+            String status = getNestedValue(returnedValues, backupSiteName, "status");
+            return "online".equals(status);
+        }
+    }
+
+    private HttpResponse<String> sendRequestWithAction(String url, String method, String action) {
+        URI uri = null;
+        try {
+            var uriBuilder = new URIBuilder(url);
+
+            if (action != null) {
+                uriBuilder.addParameter("action", action);
+            }
+
+            uri = uriBuilder.build();
+        } catch (URISyntaxException e) {
+            throw new RuntimeException(e);
+        }
+
+        var builder = HttpRequest.newBuilder().uri(uri);
+
+        if (method.equals("POST")) {
+            builder = builder.method("POST", HttpRequest.BodyPublishers.noBody());
+        }
+
+        HttpRequest request = builder.header("Authorization", getBasicAuthenticationHeader(ISPN_USERNAME, MAIN_PASSWORD))
+                .build();
+        try {
+            return httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+        } catch (IOException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
     }
 
     @Override
-    public Cache cache(String name) {
+    public ExternalCache cache(String name) {
         return new ExternalCache(name);
     }
 }