- Change dbSize implements into using scan

src/main/java/org/crazycake/shiro/BaseRedisManager.java

@@ -105,11 +105,20 @@ public void del(byte[] key) {
      * Return the size of redis db.
      */
     @Override
-    public Long dbSize() {
-        Long dbSize = 0L;
+    public Long dbSize(byte[] pattern) {
+        long dbSize = 0L;
         Jedis jedis = getJedis();
         try {
-            dbSize = jedis.dbSize();
+            ScanParams params = new ScanParams();
+            params.count(count);
+            params.match(pattern);
+            byte[] cursor = ScanParams.SCAN_POINTER_START_BINARY;
+            ScanResult<byte[]> scanResult;
+            do {
+                scanResult = jedis.scan(cursor, params);
+                dbSize++;
+                cursor = scanResult.getCursorAsBytes();
+            } while (scanResult.getStringCursor().compareTo(ScanParams.SCAN_POINTER_START) > 0);
         } finally {
             jedis.close();
         }

src/main/java/org/crazycake/shiro/IRedisManager.java

@@ -33,7 +33,7 @@ public interface IRedisManager {
     /**
      * size
      */
-    Long dbSize();
+    Long dbSize(byte[] pattern);
 
     /**
      * keys

src/main/java/org/crazycake/shiro/RedisCache.java

@@ -185,7 +185,12 @@ public void clear() throws CacheException {
 
 	@Override
 	public int size() {
-		Long longSize = new Long(redisManager.dbSize());
+		Long longSize = 0L;
+		try {
+			longSize = new Long(redisManager.dbSize(keySerializer.serialize(this.keyPrefix + "*")));
+		} catch (SerializationException e) {
+			logger.error("get keys error", e);
+		}
 		return longSize.intValue();
 	}
 

src/main/java/org/crazycake/shiro/RedisClusterManager.java

@@ -61,13 +61,15 @@ protected JedisCluster getJedisCluster() {
         return jedisCluster;
     }
 
+    @Override
     public byte[] get(byte[] key) {
         if (key == null) {
             return null;
         }
         return getJedisCluster().get(key);
     }
 
+    @Override
     public byte[] set(byte[] key, byte[] value, int expireTime) {
         if (key == null) {
             return null;
@@ -79,30 +81,31 @@ public byte[] set(byte[] key, byte[] value, int expireTime) {
         return value;
     }
 
+    @Override
     public void del(byte[] key) {
         if (key == null) {
             return;
         }
         getJedisCluster().del(key);
     }
 
-    public Long dbSize() {
+    @Override
+    public Long dbSize(byte[] pattern) {
         Long dbSize = 0L;
         Map<String, JedisPool> clusterNodes = getJedisCluster().getClusterNodes();
-        for (String k : clusterNodes.keySet()) {
-            JedisPool jp = clusterNodes.get(k);
-            Jedis connection = jp.getResource();
-            try {
-                dbSize += connection.dbSize();
-            } catch (Exception e) {
-                e.printStackTrace();
-            } finally {
-                connection.close();
+        Iterator<Map.Entry<String, JedisPool>> nodeIt = clusterNodes.entrySet().iterator();
+        while (nodeIt.hasNext()) {
+            Map.Entry<String, JedisPool> node = nodeIt.next();
+            long nodeDbSize = getDbSizeFromClusterNode(node.getValue(), pattern);
+            if (nodeDbSize == 0L) {
+                continue;
             }
+            dbSize += nodeDbSize;
         }
         return dbSize;
     }
 
+    @Override
     public Set<byte[]> keys(byte[] pattern) {
         Set<byte[]> keys = new HashSet<byte[]>();
         Map<String, JedisPool> clusterNodes = getJedisCluster().getClusterNodes();
@@ -140,6 +143,27 @@ private Set<byte[]> getKeysFromClusterNode(JedisPool jedisPool, byte[] pattern)
         return keys;
     }
 
+    private long getDbSizeFromClusterNode(JedisPool jedisPool, byte[] pattern) {
+        long dbSize = 0L;
+        Jedis jedis = jedisPool.getResource();
+
+        try {
+            ScanParams params = new ScanParams();
+            params.count(count);
+            params.match(pattern);
+            byte[] cursor = ScanParams.SCAN_POINTER_START_BINARY;
+            ScanResult<byte[]> scanResult;
+            do {
+                scanResult = jedis.scan(cursor, params);
+                dbSize++;
+                cursor = scanResult.getCursorAsBytes();
+            } while (scanResult.getStringCursor().compareTo(ScanParams.SCAN_POINTER_START) > 0);
+        } finally {
+            jedis.close();
+        }
+        return dbSize;
+    }
+
     public int getMaxAttempts() {
         return maxAttempts;
     }

src/test/java/org/crazycake/shiro/RedisCacheTest.java

@@ -82,8 +82,8 @@ public void testRedisCache() {
     }
 
     @Test
-    public void testSize() {
-        when(redisManager.dbSize()).thenReturn(2L);
+    public void testSize() throws SerializationException {
+        when(redisManager.dbSize(keySerializer.serialize(testPrefix + "*"))).thenReturn(2L);
         assertThat(redisCache.size(), is(2));
     }
 

src/test/java/org/crazycake/shiro/RedisManagerTest.java

@@ -15,7 +15,6 @@
 import java.util.Set;
 
 import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.CoreMatchers.nullValue;
 import static org.junit.Assert.assertThat;
 import static org.mockito.Mockito.*;
 
@@ -85,17 +84,6 @@ public void testDel() throws SerializationException {
         verify(jedis, times(1)).del(any((new byte[0]).getClass()));
     }
 
-    @Test
-    public void testDbSize() {
-        when(jedis.dbSize()).thenReturn(3L);
-        Long actualDbSize = redisManager.dbSize();
-        assertThat(actualDbSize, is(3L));
-
-        when(jedis.dbSize()).thenReturn(null);
-        actualDbSize = redisManager.dbSize();
-        assertThat(actualDbSize, is(nullValue()));
-    }
-
     @Test
     public void testKeys() throws SerializationException {
         ScanResult<byte[]> scanResult = mock(ScanResult.class);