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.

Author: kangelov

openid-connect-common/src/main/java/org/mitre/oauth2/model/AuthenticationHolderEntity.java

@@ -33,7 +33,8 @@
 @Entity
 @Table(name = "authentication_holder")
 @NamedQueries ({
-@NamedQuery(name = "AuthenticationHolderEntity.getByAuthentication", query = "select a from AuthenticationHolderEntity a where a.authentication = :authentication")
+@NamedQuery(name = "AuthenticationHolderEntity.getByAuthentication", query = "select a from AuthenticationHolderEntity a where a.authentication = :authentication"),
+@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 {
 

openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2AccessTokenEntity.java

@@ -58,6 +58,7 @@
 @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"),

openid-connect-common/src/main/java/org/mitre/oauth2/model/OAuth2RefreshTokenEntity.java

@@ -50,6 +50,7 @@
 @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")

openid-connect-common/src/main/java/org/mitre/oauth2/repository/AuthenticationHolderRepository.java

@@ -16,6 +16,8 @@
  ******************************************************************************/
 package org.mitre.oauth2.repository;
 
+import java.util.List;
+
 import org.mitre.oauth2.model.AuthenticationHolderEntity;
 import org.springframework.security.oauth2.provider.OAuth2Authentication;
 
@@ -30,5 +32,8 @@ public interface AuthenticationHolderRepository {
 public void remove(AuthenticationHolderEntity a);
 
 public AuthenticationHolderEntity save(AuthenticationHolderEntity a);
+
+public List<AuthenticationHolderEntity> getOrphanedAuthenticationHolders();
+
 
 }

openid-connect-common/src/main/java/org/mitre/oauth2/repository/OAuth2TokenRepository.java

@@ -57,5 +57,9 @@ public interface OAuth2TokenRepository {
 public Set<OAuth2AccessTokenEntity> getAllAccessTokens();
 
 public Set<OAuth2RefreshTokenEntity> getAllRefreshTokens();
+
+public Set<OAuth2AccessTokenEntity> getAllExpiredAccessTokens();
+
+public Set<OAuth2RefreshTokenEntity> getAllExpiredRefreshTokens();
 
 }

openid-connect-server/src/main/java/org/mitre/oauth2/repository/impl/JpaAuthenticationHolderRepository.java

@@ -16,6 +16,8 @@
  ******************************************************************************/
 package org.mitre.oauth2.repository.impl;
 
+import java.util.List;
+
 import javax.persistence.EntityManager;
 import javax.persistence.PersistenceContext;
 import javax.persistence.TypedQuery;
@@ -31,6 +33,8 @@
 @Transactional
 public class JpaAuthenticationHolderRepository implements AuthenticationHolderRepository {
 
+private static final int MAXEXPIREDRESULTS = 1000;
+
 @PersistenceContext
 private EntityManager manager;
 
@@ -73,5 +77,14 @@ public void remove(AuthenticationHolderEntity a) {
 public AuthenticationHolderEntity save(AuthenticationHolderEntity a) {
 return JpaUtil.saveOrUpdate(a.getId(), manager, a);
 }
+
+@Override
+@Transactional
+public List<AuthenticationHolderEntity> getOrphanedAuthenticationHolders() {
+TypedQuery<AuthenticationHolderEntity> query = manager.createNamedQuery("AuthenticationHolderEntity.getUnusedAuthenticationHolders", AuthenticationHolderEntity.class);
+query.setMaxResults(MAXEXPIREDRESULTS);
+List<AuthenticationHolderEntity> unusedAuthenticationHolders = query.getResultList();
+return unusedAuthenticationHolders;
+}
 
 }

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 @@
 @Repository
 public class JpaOAuth2TokenRepository implements OAuth2TokenRepository {
 
+private static final int MAXEXPIREDRESULTS = 1000;
+
 @PersistenceContext
 private EntityManager manager;
 
@@ -178,5 +181,21 @@ public OAuth2AccessTokenEntity getAccessTokenForIdToken(OAuth2AccessTokenEntity
 List<OAuth2AccessTokenEntity> accessTokens = queryA.getResultList();
 return JpaUtil.getSingleResult(accessTokens);
 }
+
+@Override
+public Set<OAuth2AccessTokenEntity> getAllExpiredAccessTokens() {
+TypedQuery<OAuth2AccessTokenEntity> query = manager.createNamedQuery("OAuth2AccessTokenEntity.getAllExpiredByDate", OAuth2AccessTokenEntity.class);
+query.setParameter("date", new Date());
+query.setMaxResults(MAXEXPIREDRESULTS);
+return new LinkedHashSet<OAuth2AccessTokenEntity>(query.getResultList());
+}
+
+@Override
+public Set<OAuth2RefreshTokenEntity> getAllExpiredRefreshTokens() {
+TypedQuery<OAuth2RefreshTokenEntity> query = manager.createNamedQuery("OAuth2RefreshTokenEntity.getAllExpiredByDate", OAuth2RefreshTokenEntity.class);
+query.setParameter("date", new Date());
+query.setMaxResults(MAXEXPIREDRESULTS);
+return new LinkedHashSet<OAuth2RefreshTokenEntity>(query.getResultList());
+}
 
 }

openid-connect-server/src/main/java/org/mitre/oauth2/service/impl/DefaultOAuth2ProviderTokenService.java

@@ -399,36 +399,38 @@ public void clearExpiredTokens() {
 Collection<OAuth2AccessTokenEntity> 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<OAuth2RefreshTokenEntity> refreshTokens = getExpiredRefreshTokens();
 logger.info("Found " + refreshTokens.size() + " expired refresh tokens");
 for (OAuth2RefreshTokenEntity oAuth2RefreshTokenEntity : refreshTokens) {
 revokeRefreshToken(oAuth2RefreshTokenEntity);
 }
-}
-
-private Predicate<OAuth2AccessTokenEntity> isAccessTokenExpired = new Predicate<OAuth2AccessTokenEntity>() {
-@Override
-public boolean apply(OAuth2AccessTokenEntity input) {
-return (input != null && input.isExpired());
+
+Collection<AuthenticationHolderEntity> authHolders = getOrphanedAuthenticationHolders();
+logger.info("Found " + authHolders.size() + " orphaned authentication holders");
+for(AuthenticationHolderEntity authHolder : authHolders) {
+authenticationHolderRepository.remove(authHolder);
 }
-};
-
-private Predicate<OAuth2RefreshTokenEntity> isRefreshTokenExpired = new Predicate<OAuth2RefreshTokenEntity>() {
-@Override
-public boolean apply(OAuth2RefreshTokenEntity input) {
-return (input != null && input.isExpired());
-}
-};
+}
 
 private Collection<OAuth2AccessTokenEntity> getExpiredAccessTokens() {
-return Collections2.filter(tokenRepository.getAllAccessTokens(), isAccessTokenExpired);
+return Sets.newHashSet(tokenRepository.getAllExpiredAccessTokens());
 }
 
 private Collection<OAuth2RefreshTokenEntity> getExpiredRefreshTokens() {
-return Collections2.filter(tokenRepository.getAllRefreshTokens(), isRefreshTokenExpired);
+return Sets.newHashSet(tokenRepository.getAllExpiredRefreshTokens());
+}
+
+private Collection<AuthenticationHolderEntity> getOrphanedAuthenticationHolders() {
+return Sets.newHashSet(authenticationHolderRepository.getOrphanedAuthenticationHolders());
 }
 
 /* (non-Javadoc)