/*
* Copyright 2017 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.cluster.infinispan;
import java.util.HashSet;
import java.util.Set;
import org.infinispan.Cache;
import org.infinispan.client.hotrod.RemoteCache;
import org.infinispan.manager.EmbeddedCacheManager;
import org.infinispan.persistence.remote.configuration.RemoteStoreConfigurationBuilder;
import org.jboss.logging.Logger;
import org.junit.Assert;
import org.keycloak.common.util.Time;
import org.keycloak.connections.infinispan.InfinispanConnectionProvider;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.sessions.infinispan.changes.SessionEntityWrapper;
import org.keycloak.models.sessions.infinispan.entities.UserSessionEntity;
import org.keycloak.models.sessions.infinispan.initializer.SessionLoader;
import org.keycloak.models.sessions.infinispan.remotestore.RemoteCacheSessionsLoader;
import org.keycloak.models.sessions.infinispan.remotestore.RemoteCacheSessionsLoaderContext;
import org.keycloak.models.sessions.infinispan.util.InfinispanUtil;
/**
* @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
*/
public class RemoteCacheSessionsLoaderTest {
protected static final Logger logger = Logger.getLogger(RemoteCacheSessionsLoaderTest.class);
private static final int COUNT = 10000;
public static void main(String[] args) throws Exception {
String cacheName = InfinispanConnectionProvider.USER_SESSION_CACHE_NAME;
Cache cache1 = createManager(1, cacheName).getCache(cacheName);
Cache cache2 = cache1.getCacheManager().getCache("local");
RemoteCache remoteCache = InfinispanUtil.getRemoteCache(cache1);
cache1.clear();
cache2.clear();
remoteCache.clear();
try {
for (int i=0 ; i<COUNT ; i++) {
// Create initial item
UserSessionEntity session = new UserSessionEntity();
session.setId("loader-key-" + i);
session.setRealmId("master");
session.setBrokerSessionId("!23123123");
session.setBrokerUserId(null);
session.setUser("admin");
session.setLoginUsername("admin");
session.setIpAddress("123.44.143.178");
session.setStarted(Time.currentTime());
session.setLastSessionRefresh(Time.currentTime());
SessionEntityWrapper<UserSessionEntity> wrappedSession = new SessionEntityWrapper<>(session);
// Create caches, listeners and finally worker threads
remoteCache.put("loader-key-" + i, wrappedSession);
Assert.assertFalse(cache2.containsKey("loader-key-" + i));
if (i % 1000 == 0) {
logger.infof("%d sessions added", i);
}
}
// RemoteCacheSessionsLoader loader = new RemoteCacheSessionsLoader(InfinispanConnectionProvider.USER_SESSION_CACHE_NAME, 64) {
//
// @Override
// protected Cache getCache(KeycloakSession session) {
// return cache2;
// }
//
// @Override
// protected RemoteCache getRemoteCache(KeycloakSession session) {
// return remoteCache;
// }
//
// };
// Just to be able to test serializability
RemoteCacheSessionsLoader loader = new CustomLoader(cacheName, 64, cache2, remoteCache);
loader.init(null);
RemoteCacheSessionsLoaderContext ctx = loader.computeLoaderContext(null);
Assert.assertEquals(ctx.getSessionsTotal(), COUNT);
Assert.assertEquals(ctx.getIspnSegmentsCount(), 256);
//Assert.assertEquals(ctx.getSegmentsCount(), 16);
Assert.assertEquals(ctx.getSessionsPerSegment(), 64);
int totalCount = 0;
logger.infof("segmentsCount: %d", ctx.getSegmentsCount());
Set<String> visitedKeys = new HashSet<>();
for (int currentSegment=0 ; currentSegment<ctx.getSegmentsCount() ; currentSegment++) {
logger.infof("Loading segment %d", currentSegment);
loader.loadSessions(null, ctx, new SessionLoader.WorkerContext(currentSegment, currentSegment));
logger.infof("Loaded %d keys for segment %d", cache2.keySet().size(), currentSegment);
totalCount = totalCount + cache2.keySet().size();
visitedKeys.addAll(cache2.keySet());
cache2.clear();
}
Assert.assertEquals(totalCount, COUNT);
Assert.assertEquals(visitedKeys.size(), COUNT);
logger.infof("SUCCESS: Loaded %d sessions", totalCount);
} finally {
// Finish JVM
cache1.getCacheManager().stop();
}
}
private static EmbeddedCacheManager createManager(int threadId, String cacheName) {
return new TestCacheManagerFactory().createManager(threadId, cacheName, RemoteStoreConfigurationBuilder.class);
}
public static class CustomLoader extends RemoteCacheSessionsLoader {
private final transient Cache cache2;
private final transient RemoteCache remoteCache;
public CustomLoader(String cacheName, int sessionsPerSegment, Cache cache2, RemoteCache remoteCache) {
super(cacheName, sessionsPerSegment);
this.cache2 = cache2;
this.remoteCache = remoteCache;
}
@Override
protected Cache getCache(KeycloakSession session) {
return cache2;
}
@Override
protected RemoteCache getRemoteCache(KeycloakSession session) {
return remoteCache;
}
}
}