KNOX-2650 - Loading token management page can be slow if there are lo…

…ts of tokens (apache#486)

gateway-server/src/main/java/org/apache/knox/gateway/services/token/impl/JDBCTokenStateService.java

@@ -22,7 +22,6 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
-import java.util.TreeSet;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
@@ -297,16 +296,11 @@ public TokenMetadata getTokenMetadata(String tokenId) throws UnknownTokenExcepti
 
   @Override
   public Collection<KnoxToken> getTokens(String userName) {
-    final Collection<KnoxToken> tokens = new TreeSet<>();
     try {
-      tokens.addAll(tokenDatabase.getTokens(userName));
-      for (KnoxToken token : tokens) {
-        token.setMetadata(tokenDatabase.getTokenMetadata(token.getTokenId()));
-      }
+      return tokenDatabase.getTokens(userName);
     } catch (SQLException e) {
       log.errorFetchingTokensForUserFromDatabase(userName, e.getMessage(), e);
+      return Collections.emptyList();
     }
-    return tokens;
   }
-
 }

gateway-server/src/main/java/org/apache/knox/gateway/services/token/impl/TokenStateDatabase.java

@@ -29,11 +29,10 @@
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
-import java.util.TreeSet;
-
 import javax.sql.DataSource;
 
 import org.apache.commons.codec.binary.Base64;
@@ -57,9 +56,9 @@ public class TokenStateDatabase {
   private static final String ADD_METADATA_SQL = "INSERT INTO " + TOKEN_METADATA_TABLE_NAME + "(token_id, md_name, md_value) VALUES(?, ?, ?)";
   private static final String UPDATE_METADATA_SQL = "UPDATE " + TOKEN_METADATA_TABLE_NAME + " SET md_value = ? WHERE token_id = ? AND md_name = ?";
   private static final String GET_METADATA_SQL = "SELECT md_name, md_value FROM " + TOKEN_METADATA_TABLE_NAME + " WHERE token_id = ?";
-  private static final String GET_TOKENS_BY_USER_NAME_SQL = "SELECT kt.token_id, kt.issue_time, kt.expiration, kt.max_lifetime FROM " + TOKENS_TABLE_NAME
-      + " kt, " + TOKEN_METADATA_TABLE_NAME + " ktm WHERE kt.token_id = ktm.token_id AND ktm.md_name = '" + TokenMetadata.USER_NAME
-      + "' AND ktm.md_value = ? ORDER BY kt.issue_time";
+  private static final String GET_TOKENS_BY_USER_NAME_SQL = "SELECT kt.token_id, kt.issue_time, kt.expiration, kt.max_lifetime, ktm.md_name, ktm.md_value FROM " + TOKENS_TABLE_NAME
+      + " kt, " + TOKEN_METADATA_TABLE_NAME + " ktm WHERE kt.token_id = ktm.token_id AND kt.token_id IN (SELECT token_id FROM " + TOKEN_METADATA_TABLE_NAME + " WHERE md_name = '" + TokenMetadata.USER_NAME + "' AND md_value = ? )"
+      + " ORDER BY kt.issue_time";
 
   private final DataSource dataSource;
 
@@ -192,24 +191,34 @@ TokenMetadata getTokenMetadata(String tokenId) throws SQLException {
         final Map<String, String> metadataMap = new HashMap<>();
         while (rs.next()) {
           String metadataName = rs.getString(1);
-          metadataMap.put(metadataName, metadataName.equals(TokenMetadata.PASSCODE) ? new String(Base64.decodeBase64(rs.getString(2).getBytes(UTF_8)), UTF_8) : rs.getString(2));
+          metadataMap.put(metadataName, decodeMetadata(metadataName, rs.getString(2)));
         }
         return metadataMap.isEmpty() ? null : new TokenMetadata(metadataMap);
       }
     }
   }
 
+  private static String decodeMetadata(String metadataName, String metadataValue) {
+    return metadataName.equals(TokenMetadata.PASSCODE) ? new String(Base64.decodeBase64(metadataValue.getBytes(UTF_8)), UTF_8) : metadataValue;
+  }
+
   Collection<KnoxToken> getTokens(String userName) throws SQLException {
-    final Collection<KnoxToken> tokens = new TreeSet<>();
+    Map<String, KnoxToken> tokenMap = new LinkedHashMap<>();
     try (Connection connection = dataSource.getConnection(); PreparedStatement getTokenIdsStatement = connection.prepareStatement(GET_TOKENS_BY_USER_NAME_SQL)) {
       getTokenIdsStatement.setString(1, userName);
       try (ResultSet rs = getTokenIdsStatement.executeQuery()) {
         while(rs.next()) {
-          tokens.add(new KnoxToken(rs.getString(1), rs.getLong(2), rs.getLong(3), rs.getLong(4)));
+          String tokenId = rs.getString(1);
+          long issueTime = rs.getLong(2);
+          long expiration = rs.getLong(3);
+          long maxLifeTime = rs.getLong(4);
+          String metaName = rs.getString(5);
+          String metaValue = rs.getString(6);
+          KnoxToken token = tokenMap.computeIfAbsent(tokenId, id -> new KnoxToken(tokenId, issueTime, expiration, maxLifeTime));
+          token.addMetadata(metaName, decodeMetadata(metaName, metaValue));
         }
-        return tokens;
+        return tokenMap.values();
       }
     }
   }
-
 }

gateway-server/src/test/java/org/apache/knox/gateway/services/token/impl/JDBCTokenStateServiceTest.java

@@ -28,6 +28,11 @@
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.text.SimpleDateFormat;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -39,6 +44,7 @@
 import org.apache.derby.drda.NetworkServerControl;
 import org.apache.knox.gateway.config.GatewayConfig;
 import org.apache.knox.gateway.services.security.AliasService;
+import org.apache.knox.gateway.services.security.token.KnoxToken;
 import org.apache.knox.gateway.services.security.token.TokenMetadata;
 import org.apache.knox.gateway.services.security.token.UnknownTokenException;
 import org.apache.knox.gateway.services.security.token.impl.TokenMAC;
@@ -130,6 +136,55 @@ public void testAddToken() throws Exception {
     assertEquals(issueTime + maxLifetimeDuration, getLongTokenAttributeFromDatabase(tokenId, TokenStateDatabase.GET_MAX_LIFETIME_SQL));
   }
 
+  @Test
+  public void testAddTokensForMultipleUsers() throws Exception {
+    String user1 = "user1";
+    String user2 = "user2";
+    String id1 = "token1";
+    String id2 = "token2";
+    String id3 = "token3";
+
+    long issueTime1 = 1;
+    long expiration1 = 1;
+    String comment1 = "comment1";
+
+    long issueTime2 = 2;
+    long expiration2 = 2;
+    String comment2 = "comment2";
+
+    long issueTime3 = 3;
+    long expiration3 = 3;
+    String comment3 = "comment3";
+
+    truncateDatabase();
+
+    saveToken(user1, id1, issueTime1, expiration1, comment1);
+    saveToken(user1, id2, issueTime2, expiration2, comment2);
+    saveToken(user2, id3, issueTime3, expiration3, comment3);
+
+    List<KnoxToken> user1Tokens = new ArrayList<>(jdbcTokenStateService.getTokens(user1));
+    assertEquals(2, user1Tokens.size());
+    assertToken(user1Tokens.get(0), id1, expiration1, comment1, issueTime1);
+    assertToken(user1Tokens.get(1), id2, expiration2, comment2, issueTime2);
+
+    List<KnoxToken> user2Tokens = new ArrayList<>(jdbcTokenStateService.getTokens(user2));
+    assertEquals(1, user2Tokens.size());
+    assertToken(user2Tokens.get(0), id3, expiration3, comment3, issueTime3);
+  }
+
+  private void assertToken(KnoxToken knoxToken, String tokenId, long expiration, String comment, long issueTime) {
+    SimpleDateFormat df = new SimpleDateFormat(KnoxToken.DATE_FORMAT, Locale.getDefault());
+    assertEquals(tokenId, knoxToken.getTokenId());
+    assertEquals(df.format(new Date(issueTime)), knoxToken.getIssueTime());
+    assertEquals(df.format(new Date(expiration)), knoxToken.getExpiration());
+    assertEquals(comment, knoxToken.getMetadata().getComment());
+  }
+
+  private void saveToken(String user, String tokenId, long issueTime, long expiration, String comment) {
+    jdbcTokenStateService.addToken(tokenId, issueTime, expiration);
+    jdbcTokenStateService.addMetadata(tokenId, new TokenMetadata(user, comment));
+  }
+
   @Test(expected = UnknownTokenException.class)
   public void testRemoveToken() throws Exception {
     truncateDatabase();

gateway-spi/src/main/java/org/apache/knox/gateway/services/security/token/KnoxToken.java

@@ -19,13 +19,15 @@
 
 import java.text.DateFormat;
 import java.text.SimpleDateFormat;
+import java.util.Collections;
 import java.util.Date;
 import java.util.Locale;
 
-public class KnoxToken implements Comparable<KnoxToken> {
+public class KnoxToken implements Comparable<KnoxToken>{
+  public static final String DATE_FORMAT = "yyyy-MM-dd'T'HH:mm:ss.SSSZ";
   // SimpleDateFormat is not thread safe must use as a ThreadLocal
   private static final ThreadLocal<DateFormat> KNOX_TOKEN_TS_FORMAT = ThreadLocal
-      .withInitial(() -> new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSZ", Locale.getDefault()));
+      .withInitial(() -> new SimpleDateFormat(DATE_FORMAT, Locale.getDefault()));
 
   private final String tokenId;
   private final long issueTime;
@@ -34,7 +36,7 @@ public class KnoxToken implements Comparable<KnoxToken> {
   private TokenMetadata metadata;
 
   public KnoxToken(String tokenId, long issueTime, long expiration, long maxLifetimeDuration) {
-    this(tokenId, issueTime, expiration, maxLifetimeDuration, null);
+    this(tokenId, issueTime, expiration, maxLifetimeDuration, new TokenMetadata(Collections.emptyMap()));
   }
 
   public KnoxToken(String tokenId, long issueTime, long expiration, long maxLifetime, TokenMetadata metadata) {
@@ -85,4 +87,8 @@ public void setMetadata(TokenMetadata metadata) {
   public int compareTo(KnoxToken other) {
     return Long.compare(this.issueTime, other.issueTime);
   }
+
+  public void addMetadata(String name, String value) {
+    metadata.add(name, value);
+  }
 }

gateway-spi/src/main/java/org/apache/knox/gateway/services/security/token/TokenMetadata.java

@@ -117,4 +117,8 @@ public boolean equals(Object obj) {
   public int hashCode() {
     return HashCodeBuilder.reflectionHashCode(this);
   }
+
+  public void add(String name, String value) {
+    metadataMap.put(name, value);
+  }
 }