From 287dce5c0241773c978532cbce4d3f35c96d1d58 Mon Sep 17 00:00:00 2001 From: kangelov Date: Tue, 8 Apr 2014 08:14:20 -0400 Subject: [PATCH] Performance improvement of token cleanup: an alternative token cleanup mechanism designed to maintain a very compact memory footprint while performing cleanup in consecutive runs of the cleanup thread. This serves to address OutOfMemoryException issues of the original token cleanup mechanism when process is under load. Also, added cleanup of the authentication_holder table. Conflicts: openid-connect-common/src/main/java/org/mitre/oauth2/model/AuthenticationHolderEntity.java openid-connect-common/src/main/java/org/mitre/oauth2/repository/AuthenticationHolderRepository.java openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaAuthenticationHolderRepository.java openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java --- .../model/AuthenticationHolderEntity.java | 3 +- .../oauth2/model/OAuth2AccessTokenEntity.java | 1 + .../model/OAuth2RefreshTokenEntity.java | 1 + .../AuthenticationHolderRepository.java | 5 +++ .../repository/OAuth2TokenRepository.java | 4 ++ .../JpaAuthenticationHolderRepository.java | 13 +++++++ .../impl/JpaOAuth2TokenRepository.java | 19 ++++++++++ .../DefaultOAuth2ProviderTokenService.java | 38 ++++++++++--------- 8 files changed, 65 insertions(+), 19 deletions(-) diff --git a/openid-connect-common/src/main/java/org/mitre/oauth2/model/AuthenticationHolderEntity.java b/openid-connect-common/src/main/java/org/mitre/oauth2/model/AuthenticationHolderEntity.java index be0c04959..3a2ab37bc 100644 --- a/openid-connect-common/src/main/java/org/mitre/oauth2/model/AuthenticationHolderEntity.java +++ b/openid-connect-common/src/main/java/org/mitre/oauth2/model/AuthenticationHolderEntity.java @@ -34,7 +34,8 @@ import org.springframework.security.oauth2.provider.OAuth2Authentication; @Table(name = "authentication_holder") @NamedQueries ({ @NamedQuery(name = "AuthenticationHolderEntity.getByAuthentication", query = "select a from AuthenticationHolderEntity a where a.authentication = :authentication"), - @NamedQuery(name = "AuthenticationHolderEntity.getAll", query = "select a from AuthenticationHolderEntity a") + @NamedQuery(name = "AuthenticationHolderEntity.getAll", query = "select a from AuthenticationHolderEntity a"), + @NamedQuery(name = "AuthenticationHolderEntity.getUnusedAuthenticationHolders", query = "select a from AuthenticationHolderEntity a where a.id not in (select t.authenticationHolderId from OAuth2AccessTokenEntity t) and a.id not in (select r.authenticationHolderId from OAuth2RefreshTokenEntity r)") }) public class AuthenticationHolderEntity { diff --git a/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2AccessTokenEntity.java b/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2AccessTokenEntity.java index ff0a3182b..07dd6e3a1 100644 --- a/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2AccessTokenEntity.java +++ b/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2AccessTokenEntity.java @@ -58,6 +58,7 @@ import com.nimbusds.jwt.JWTParser; @Table(name = "access_token") @NamedQueries({ @NamedQuery(name = "OAuth2AccessTokenEntity.getAll", query = "select a from OAuth2AccessTokenEntity a"), + @NamedQuery(name = "OAuth2AccessTokenEntity.getAllExpiredByDate", query = "select a from OAuth2AccessTokenEntity a where a.expiration <= :date"), @NamedQuery(name = "OAuth2AccessTokenEntity.getByRefreshToken", query = "select a from OAuth2AccessTokenEntity a where a.refreshToken = :refreshToken"), @NamedQuery(name = "OAuth2AccessTokenEntity.getByClient", query = "select a from OAuth2AccessTokenEntity a where a.client = :client"), @NamedQuery(name = "OAuth2AccessTokenEntity.getByAuthentication", query = "select a from OAuth2AccessTokenEntity a where a.authenticationHolder.authentication = :authentication"), diff --git a/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2RefreshTokenEntity.java b/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2RefreshTokenEntity.java index f05c1d7ab..234fc3ba6 100644 --- a/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2RefreshTokenEntity.java +++ b/openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2RefreshTokenEntity.java @@ -50,6 +50,7 @@ import com.nimbusds.jwt.JWTParser; @Table(name = "refresh_token") @NamedQueries({ @NamedQuery(name = "OAuth2RefreshTokenEntity.getAll", query = "select r from OAuth2RefreshTokenEntity r"), + @NamedQuery(name = "OAuth2RefreshTokenEntity.getAllExpiredByDate", query = "select r from OAuth2RefreshTokenEntity r where r.expiration <= :date"), @NamedQuery(name = "OAuth2RefreshTokenEntity.getByClient", query = "select r from OAuth2RefreshTokenEntity r where r.client = :client"), @NamedQuery(name = "OAuth2RefreshTokenEntity.getByTokenValue", query = "select r from OAuth2RefreshTokenEntity r where r.value = :tokenValue"), @NamedQuery(name = "OAuth2RefreshTokenEntity.getByAuthentication", query = "select r from OAuth2RefreshTokenEntity r where r.authenticationHolder.authentication = :authentication") diff --git a/openid-connect-common/src/main/java/org/mitre/oauth2/repository/AuthenticationHolderRepository.java b/openid-connect-common/src/main/java/org/mitre/oauth2/repository/AuthenticationHolderRepository.java index f90a93758..5e51666f1 100644 --- a/openid-connect-common/src/main/java/org/mitre/oauth2/repository/AuthenticationHolderRepository.java +++ b/openid-connect-common/src/main/java/org/mitre/oauth2/repository/AuthenticationHolderRepository.java @@ -17,6 +17,8 @@ package org.mitre.oauth2.repository; import java.util.Collection; +import java.util.List; + import org.mitre.oauth2.model.AuthenticationHolderEntity; import org.springframework.security.oauth2.provider.OAuth2Authentication; @@ -33,5 +35,8 @@ public interface AuthenticationHolderRepository { public void remove(AuthenticationHolderEntity a); public AuthenticationHolderEntity save(AuthenticationHolderEntity a); + + public List getOrphanedAuthenticationHolders(); + } diff --git a/openid-connect-common/src/main/java/org/mitre/oauth2/repository/OAuth2TokenRepository.java b/openid-connect-common/src/main/java/org/mitre/oauth2/repository/OAuth2TokenRepository.java index 785586d37..a06914d14 100644 --- a/openid-connect-common/src/main/java/org/mitre/oauth2/repository/OAuth2TokenRepository.java +++ b/openid-connect-common/src/main/java/org/mitre/oauth2/repository/OAuth2TokenRepository.java @@ -57,5 +57,9 @@ public interface OAuth2TokenRepository { public Set getAllAccessTokens(); public Set getAllRefreshTokens(); + + public Set getAllExpiredAccessTokens(); + + public Set getAllExpiredRefreshTokens(); } diff --git a/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaAuthenticationHolderRepository.java b/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaAuthenticationHolderRepository.java index be72dea08..7cbed75ed 100644 --- a/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaAuthenticationHolderRepository.java +++ b/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaAuthenticationHolderRepository.java @@ -17,6 +17,8 @@ package org.mitre.oauth2.repository.impl; import java.util.Collection; +import java.util.List; + import javax.persistence.EntityManager; import javax.persistence.PersistenceContext; import javax.persistence.TypedQuery; @@ -32,6 +34,8 @@ import org.springframework.transaction.annotation.Transactional; @Transactional public class JpaAuthenticationHolderRepository implements AuthenticationHolderRepository { + private static final int MAXEXPIREDRESULTS = 1000; + @PersistenceContext private EntityManager manager; @@ -80,5 +84,14 @@ public class JpaAuthenticationHolderRepository implements AuthenticationHolderRe public AuthenticationHolderEntity save(AuthenticationHolderEntity a) { return JpaUtil.saveOrUpdate(a.getId(), manager, a); } + + @Override + @Transactional + public List getOrphanedAuthenticationHolders() { + TypedQuery query = manager.createNamedQuery("AuthenticationHolderEntity.getUnusedAuthenticationHolders", AuthenticationHolderEntity.class); + query.setMaxResults(MAXEXPIREDRESULTS); + List unusedAuthenticationHolders = query.getResultList(); + return unusedAuthenticationHolders; + } } diff --git a/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaOAuth2TokenRepository.java b/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaOAuth2TokenRepository.java index e381f66bd..37ce82c3b 100644 --- a/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaOAuth2TokenRepository.java +++ b/openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaOAuth2TokenRepository.java @@ -16,6 +16,7 @@ ******************************************************************************/ package org.mitre.oauth2.repository.impl; +import java.util.Date; import java.util.LinkedHashSet; import java.util.List; import java.util.Set; @@ -36,6 +37,8 @@ import org.springframework.transaction.annotation.Transactional; @Repository public class JpaOAuth2TokenRepository implements OAuth2TokenRepository { + private static final int MAXEXPIREDRESULTS = 1000; + @PersistenceContext private EntityManager manager; @@ -178,5 +181,21 @@ public class JpaOAuth2TokenRepository implements OAuth2TokenRepository { List accessTokens = queryA.getResultList(); return JpaUtil.getSingleResult(accessTokens); } + + @Override + public Set getAllExpiredAccessTokens() { + TypedQuery query = manager.createNamedQuery("OAuth2AccessTokenEntity.getAllExpiredByDate", OAuth2AccessTokenEntity.class); + query.setParameter("date", new Date()); + query.setMaxResults(MAXEXPIREDRESULTS); + return new LinkedHashSet(query.getResultList()); + } + + @Override + public Set getAllExpiredRefreshTokens() { + TypedQuery query = manager.createNamedQuery("OAuth2RefreshTokenEntity.getAllExpiredByDate", OAuth2RefreshTokenEntity.class); + query.setParameter("date", new Date()); + query.setMaxResults(MAXEXPIREDRESULTS); + return new LinkedHashSet(query.getResultList()); + } } diff --git a/openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java b/openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java index d65e03361..22a83e67c 100644 --- a/openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java +++ b/openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java @@ -367,7 +367,13 @@ public class DefaultOAuth2ProviderTokenService implements OAuth2TokenEntityServi Collection accessTokens = getExpiredAccessTokens(); logger.info("Found " + accessTokens.size() + " expired access tokens"); for (OAuth2AccessTokenEntity oAuth2AccessTokenEntity : accessTokens) { - revokeAccessToken(oAuth2AccessTokenEntity); + try { + revokeAccessToken(oAuth2AccessTokenEntity); + } catch (IllegalArgumentException e) { + //An ID token is deleted with its corresponding access token, but then the ID token is on the list of expired tokens as well and there is + //nothing in place to distinguish it from any other. + //An attempt to delete an already deleted token returns an error, stopping the cleanup dead. We need it to keep going. + } } Collection refreshTokens = getExpiredRefreshTokens(); @@ -375,28 +381,24 @@ public class DefaultOAuth2ProviderTokenService implements OAuth2TokenEntityServi for (OAuth2RefreshTokenEntity oAuth2RefreshTokenEntity : refreshTokens) { revokeRefreshToken(oAuth2RefreshTokenEntity); } + + Collection authHolders = getOrphanedAuthenticationHolders(); + logger.info("Found " + authHolders.size() + " orphaned authentication holders"); + for(AuthenticationHolderEntity authHolder : authHolders) { + authenticationHolderRepository.remove(authHolder); + } } - - private Predicate isAccessTokenExpired = new Predicate() { - @Override - public boolean apply(OAuth2AccessTokenEntity input) { - return (input != null && input.isExpired()); - } - }; - - private Predicate isRefreshTokenExpired = new Predicate() { - @Override - public boolean apply(OAuth2RefreshTokenEntity input) { - return (input != null && input.isExpired()); - } - }; - + private Collection getExpiredAccessTokens() { - return Sets.filter(Sets.newHashSet(tokenRepository.getAllAccessTokens()), isAccessTokenExpired); + return Sets.newHashSet(tokenRepository.getAllExpiredAccessTokens()); } private Collection getExpiredRefreshTokens() { - return Sets.filter(Sets.newHashSet(tokenRepository.getAllRefreshTokens()), isRefreshTokenExpired); + return Sets.newHashSet(tokenRepository.getAllExpiredRefreshTokens()); + } + + private Collection getOrphanedAuthenticationHolders() { + return Sets.newHashSet(authenticationHolderRepository.getOrphanedAuthenticationHolders()); } /* (non-Javadoc)