keycloak-uncached

client stat improvement

1/31/2018 4:05:13 PM

Details

diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java
index c55013e..14167a3 100755
--- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/InfinispanUserSessionProvider.java
@@ -20,6 +20,7 @@ package org.keycloak.models.sessions.infinispan;
 import org.infinispan.Cache;
 import org.infinispan.client.hotrod.RemoteCache;
 import org.infinispan.context.Flag;
+import org.infinispan.stream.CacheCollectors;
 import org.jboss.logging.Logger;
 import org.keycloak.cluster.ClusterProvider;
 import org.keycloak.common.util.Time;
@@ -68,7 +69,10 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -296,16 +300,6 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
         return getUserSessions(realm, client, firstResult, maxResults, false);
     }
 
-    @Override
-    public List<UserSessionModel> getOfflineUserSessions(RealmModel realm) {
-        return getOfflineUserSessions(realm, -1, -1);
-    }
-
-    @Override
-    public List<UserSessionModel> getOfflineUserSessions(RealmModel realm, int first, int max) {
-        return getUserSessions(realm, first, max, true);
-    }
-
     protected List<UserSessionModel> getUserSessions(final RealmModel realm, ClientModel client, int firstResult, int maxResults, final boolean offline) {
         final String clientUuid = client.getId();
         UserSessionPredicate predicate = UserSessionPredicate.create(realm.getId()).client(clientUuid);
@@ -313,22 +307,6 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
         return getUserSessionModels(realm, firstResult, maxResults, offline, predicate);
     }
 
-    @Override
-    public List<UserSessionModel> getUserSessions(RealmModel realm) {
-        return getUserSessions(realm, -1, -1);
-    }
-
-    @Override
-    public List<UserSessionModel> getUserSessions(RealmModel realm, int firstResult, int maxResults) {
-        return getUserSessions(realm, firstResult, maxResults, false);
-    }
-
-    protected List<UserSessionModel> getUserSessions(final RealmModel realm, int firstResult, int maxResults, final boolean offline) {
-        UserSessionPredicate predicate = UserSessionPredicate.create(realm.getId());
-
-        return getUserSessionModels(realm, firstResult, maxResults, offline, predicate);
-    }
-
     protected List<UserSessionModel> getUserSessionModels(RealmModel realm, int firstResult, int maxResults, boolean offline, UserSessionPredicate predicate) {
         Cache<String, SessionEntityWrapper<UserSessionEntity>> cache = getCache(offline);
         cache = CacheDecorators.skipCacheLoaders(cache);
@@ -428,6 +406,25 @@ public class InfinispanUserSessionProvider implements UserSessionProvider {
         return getUserSessionsCount(realm, client, false);
     }
 
+    @Override
+    public Map<String, Long> getActiveClientSessionStats(RealmModel realm, boolean offline) {
+        Cache<String, SessionEntityWrapper<UserSessionEntity>> cache = getCache(offline);
+        cache = CacheDecorators.skipCacheLoaders(cache);
+        return cache.entrySet().stream()
+                .filter(UserSessionPredicate.create(realm.getId()))
+                .map(Mappers.authClientSessionSetMapper())
+                .flatMap(Mappers::toStream)
+                .collect(
+                        countingGroupingCollector()
+                );
+    }
+
+    public static Collector<String, ?, Map<String, Long>> countingGroupingCollector() {
+        return CacheCollectors.serializableCollector(
+                () -> Collectors.groupingBy(Function.identity(), Collectors.counting())
+        );
+    }
+
     protected long getUserSessionsCount(RealmModel realm, ClientModel client, boolean offline) {
         Cache<String, SessionEntityWrapper<UserSessionEntity>> cache = getCache(offline);
         cache = CacheDecorators.skipCacheLoaders(cache);
diff --git a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/stream/Mappers.java b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/stream/Mappers.java
index 177fd23..79761d0 100644
--- a/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/stream/Mappers.java
+++ b/model/infinispan/src/main/java/org/keycloak/models/sessions/infinispan/stream/Mappers.java
@@ -25,9 +25,13 @@ import org.keycloak.models.sessions.infinispan.entities.SessionEntity;
 import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity;
 
 import java.io.Serializable;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.function.Function;
+import java.util.stream.Stream;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
@@ -125,4 +129,21 @@ public class Mappers {
         }
     }
 
+    private static class AuthClientSessionSetMapper implements Function<Map.Entry<String, SessionEntityWrapper<UserSessionEntity>>, Set<String>>, Serializable {
+
+        @Override
+        public Set<String> apply(Map.Entry<String, SessionEntityWrapper<UserSessionEntity>> entry) {
+            UserSessionEntity entity = entry.getValue().getEntity();
+            return entity.getAuthenticatedClientSessions().keySet();
+        }
+    }
+
+    public static <T> Stream<T> toStream(Collection<T> collection) {
+        return collection.stream();
+    }
+
+    public static Function<Map.Entry<String, SessionEntityWrapper<UserSessionEntity>>, Set<String>> authClientSessionSetMapper() {
+        return new AuthClientSessionSetMapper();
+    }
+
 }
diff --git a/server-spi/src/main/java/org/keycloak/models/UserSessionProvider.java b/server-spi/src/main/java/org/keycloak/models/UserSessionProvider.java
index 00827e1..ac1c224 100755
--- a/server-spi/src/main/java/org/keycloak/models/UserSessionProvider.java
+++ b/server-spi/src/main/java/org/keycloak/models/UserSessionProvider.java
@@ -20,6 +20,7 @@ package org.keycloak.models;
 import org.keycloak.provider.Provider;
 
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 import java.util.function.Predicate;
 
@@ -35,11 +36,9 @@ public interface UserSessionProvider extends Provider {
     UserSessionModel createUserSession(RealmModel realm, UserModel user, String loginUsername, String ipAddress, String authMethod, boolean rememberMe, String brokerSessionId, String brokerUserId);
     UserSessionModel createUserSession(String id, RealmModel realm, UserModel user, String loginUsername, String ipAddress, String authMethod, boolean rememberMe, String brokerSessionId, String brokerUserId);
     UserSessionModel getUserSession(RealmModel realm, String id);
-    List<UserSessionModel> getUserSessions(RealmModel realm);
     List<UserSessionModel> getUserSessions(RealmModel realm, UserModel user);
     List<UserSessionModel> getUserSessions(RealmModel realm, ClientModel client);
     List<UserSessionModel> getUserSessions(RealmModel realm, ClientModel client, int firstResult, int maxResults);
-    List<UserSessionModel> getUserSessions(RealmModel realm, int firstResult, int maxResults);
     List<UserSessionModel> getUserSessionByBrokerUserId(RealmModel realm, String brokerUserId);
     UserSessionModel getUserSessionByBrokerSessionId(RealmModel realm, String brokerSessionId);
 
@@ -51,6 +50,15 @@ public interface UserSessionProvider extends Provider {
 
     long getActiveUserSessions(RealmModel realm, ClientModel client);
 
+    /**
+     * Returns a summary of client sessions key is client.getId()
+     *
+     * @param realm
+     * @param offline
+     * @return
+     */
+    Map<String, Long> getActiveClientSessionStats(RealmModel realm, boolean offline);
+
     /** This will remove attached ClientLoginSessionModels too **/
     void removeUserSession(RealmModel realm, UserSessionModel session);
     void removeUserSessions(RealmModel realm, UserModel user);
@@ -77,11 +85,9 @@ public interface UserSessionProvider extends Provider {
     /** Will automatically attach newly created offline client session to the offlineUserSession **/
     AuthenticatedClientSessionModel createOfflineClientSession(AuthenticatedClientSessionModel clientSession, UserSessionModel offlineUserSession);
     List<UserSessionModel> getOfflineUserSessions(RealmModel realm, UserModel user);
-    List<UserSessionModel> getOfflineUserSessions(RealmModel realm);
 
     long getOfflineSessionsCount(RealmModel realm, ClientModel client);
     List<UserSessionModel> getOfflineUserSessions(RealmModel realm, ClientModel client, int first, int max);
-    List<UserSessionModel> getOfflineUserSessions(RealmModel realm, int first, int max);
 
     /** Triggered by persister during pre-load. It optionally imports authenticatedClientSessions too if requested. Otherwise the imported UserSession will have empty list of AuthenticationSessionModel **/
     UserSessionModel importUserSession(UserSessionModel persistentUserSession, boolean offline, boolean importAuthenticatedClientSessions);
diff --git a/services/src/main/java/org/keycloak/services/resources/admin/RealmAdminResource.java b/services/src/main/java/org/keycloak/services/resources/admin/RealmAdminResource.java
index b5b583b..3a82baa 100644
--- a/services/src/main/java/org/keycloak/services/resources/admin/RealmAdminResource.java
+++ b/services/src/main/java/org/keycloak/services/resources/admin/RealmAdminResource.java
@@ -507,19 +507,7 @@ public class RealmAdminResource {
 
         Map<String, Map<String, String>> data = new HashMap();
         {
-            List<UserSessionModel> userSessions = session.sessions().getUserSessions(realm);
-            Map<String, Long> activeCount = new HashMap<>();
-            // we have to iterate over all realm user sessions as clients coming from client storage provider might not be reachable from getClients()
-            for (UserSessionModel userSession : userSessions) {
-                for (String id : userSession.getAuthenticatedClientSessions().keySet()) {
-                    Long number = activeCount.get(id);
-                    if (number == null) {
-                        activeCount.put(id, new Long(1));
-                    } else {
-                        activeCount.put(id, number + 1);
-                    }
-                }
-            }
+            Map<String, Long> activeCount =session.sessions().getActiveClientSessionStats(realm, false);
             for (Map.Entry<String, Long> entry : activeCount.entrySet()) {
                 Map<String, String> map = new HashMap<>();
                 ClientModel client = realm.getClientById(entry.getKey());
@@ -532,19 +520,7 @@ public class RealmAdminResource {
             }
         }
         {
-            Map<String, Long> offlineCount = new HashMap<>();
-            // we have to iterate over all realm user sessions as clients coming from client storage provider might not be reachable from getClients()
-            List<UserSessionModel> offlineSessions = session.sessions().getOfflineUserSessions(realm);
-            for (UserSessionModel userSession : offlineSessions) {
-                for (String id : userSession.getAuthenticatedClientSessions().keySet()) {
-                    Long number = offlineCount.get(id);
-                    if (number == null) {
-                        offlineCount.put(id, new Long(1));
-                    } else {
-                        offlineCount.put(id, number + 1);
-                    }
-                }
-            }
+            Map<String, Long> offlineCount = session.sessions().getActiveClientSessionStats(realm, true);
             for (Map.Entry<String, Long> entry : offlineCount.entrySet()) {
                 Map<String, String> map = data.get(entry.getKey());
                 if (map == null) {
diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/ClientStorageTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/ClientStorageTest.java
index 5656e6e..e251702 100644
--- a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/ClientStorageTest.java
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/ClientStorageTest.java
@@ -77,7 +77,9 @@ import java.io.File;
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.util.Calendar;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static java.util.Calendar.DAY_OF_WEEK;
 import static java.util.Calendar.HOUR_OF_DAY;
@@ -153,9 +155,27 @@ public class ClientStorageTest extends AbstractTestRealmKeycloakTest {
 
 
 
-    //@Test
-    public void testRunConsole() throws Exception {
-        Thread.sleep(10000000);
+    @Test
+    public void testClientStats() throws Exception {
+        testDirectGrant("hardcoded-client");
+        testDirectGrant("hardcoded-client");
+        testBrowser("test-app");
+        offlineTokenDirectGrantFlowNoRefresh();
+        List<Map<String, String>> list = adminClient.realm("test").getClientSessionStats();
+        boolean hardTested = false;
+        boolean testAppTested = false;
+        for (Map<String, String> entry : list) {
+            if (entry.get("clientId").equals("hardcoded-client")) {
+                Assert.assertEquals("3", entry.get("active"));
+                Assert.assertEquals("1", entry.get("offline"));
+                hardTested = true;
+            } else if (entry.get("clientId").equals("test-app")) {
+                Assert.assertEquals("1", entry.get("active"));
+                Assert.assertEquals("0", entry.get("offline"));
+                testAppTested = true;
+            }
+        }
+        Assert.assertTrue(hardTested && testAppTested);
     }
 
 
@@ -166,10 +186,10 @@ public class ClientStorageTest extends AbstractTestRealmKeycloakTest {
         //Thread.sleep(10000000);
     }
 
-    private void testBrowser(String clientId) {
+     private void testBrowser(String clientId) {
         oauth.clientId(clientId);
         String loginFormUrl = oauth.getLoginFormUrl();
-        log.info("loginFormUrl: " + loginFormUrl);
+        //log.info("loginFormUrl: " + loginFormUrl);
 
         //Thread.sleep(10000000);
 
@@ -405,6 +425,15 @@ public class ClientStorageTest extends AbstractTestRealmKeycloakTest {
         // Assert same token can be refreshed again
         testRefreshWithOfflineToken(token, offlineToken, offlineTokenString, token.getSessionState(), userId);
     }
+    public void offlineTokenDirectGrantFlowNoRefresh() throws Exception {
+        oauth.scope(OAuth2Constants.OFFLINE_ACCESS);
+        oauth.clientId("hardcoded-client");
+        OAuthClient.AccessTokenResponse tokenResponse = oauth.doGrantAccessTokenRequest("password", "test-user@localhost", "password");
+        Assert.assertNull(tokenResponse.getErrorDescription());
+        AccessToken token = oauth.verifyToken(tokenResponse.getAccessToken());
+        String offlineTokenString = tokenResponse.getRefreshToken();
+        RefreshToken offlineToken = oauth.verifyRefreshToken(offlineTokenString);
+    }
 
     private String testRefreshWithOfflineToken(AccessToken oldToken, RefreshToken offlineToken, String offlineTokenString,
                                                final String sessionId, String userId) {
diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/MapCollectTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/MapCollectTest.java
new file mode 100644
index 0000000..bf1e656
--- /dev/null
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/federation/storage/MapCollectTest.java
@@ -0,0 +1,110 @@
+/*
+ * Copyright 2016 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * 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 org.keycloak.testsuite.federation.storage;
+
+import org.infinispan.stream.CacheCollectors;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Function;
+import java.util.function.Predicate;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
+
+/**
+ * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
+ * @version $Revision: 1 $
+ */
+public class MapCollectTest {
+
+    public static class UserSessionObject {
+        public String id;
+        public String realm;
+        public Set<String> clients = new HashSet<>();
+
+        public UserSessionObject(String realm, String... clients) {
+            this.id = UUID.randomUUID().toString();
+            this.realm = realm;
+            for (String c : clients) this.clients.add(c);
+        }
+    }
+
+    public static class RealmFilter implements Predicate<UserSessionObject> {
+        protected String realm;
+
+        public RealmFilter(String realm) {
+            this.realm = realm;
+        }
+
+        @Override
+        public boolean test(UserSessionObject entry) {
+            return entry.realm.equals(realm);
+        }
+
+        public static RealmFilter create(String realm) {
+            return new RealmFilter(realm);
+        }
+    }
+
+    public static Set<String> clients(UserSessionObject s) {
+        return s.clients;
+    }
+
+
+    @Test
+    public void testMe() throws Exception {
+
+        List<UserSessionObject> list = Arrays.asList(
+                new UserSessionObject("realm1", "a", "b")
+                , new UserSessionObject("realm1", "a", "c")
+                , new UserSessionObject("realm1", "a", "d")
+                , new UserSessionObject("realm1", "a", "b")
+                , new UserSessionObject("realm2", "a", "b")
+                , new UserSessionObject("realm2", "a", "c")
+                , new UserSessionObject("realm2", "a", "b")
+
+        );
+
+        Map<String, Long> result = list.stream().collect(
+                Collectors.groupingBy(s -> s.realm, Collectors.summingLong(i -> 1)));
+
+        for (Map.Entry<String, Long> entry : result.entrySet()) {
+            System.out.println(entry.getKey() + ":" + entry.getValue());
+        }
+
+        result = list.stream()
+                .filter(RealmFilter.create("realm1"))
+                .map(s->s.clients)
+                .flatMap(c->c.stream())
+                .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
+
+        for (Map.Entry<String, Long> entry : result.entrySet()) {
+            System.out.println(entry.getKey() + ":" + entry.getValue());
+        }
+
+
+
+    }
+}