From 2df3a1b2581e6a9ec5584ea1d16d426b978705c2 Mon Sep 17 00:00:00 2001
From: Murilo Queiroz <muriloq@gmail.com>
Date: Thu, 30 Sep 2010 17:17:51 -0300
Subject: [PATCH] Added support for key tags.

See http://code.google.com/p/redis/wiki/FAQ#I'm_using_some_form_of_key_hashing_for_partitioning,_but_wh
---
 src/main/java/redis/clients/util/Sharded.java | 37 ++++++++++++++++++-
 .../clients/jedis/tests/ShardedJedisTest.java | 24 ++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/main/java/redis/clients/util/Sharded.java b/src/main/java/redis/clients/util/Sharded.java
index 1f56a57..21fb4a0 100644
--- a/src/main/java/redis/clients/util/Sharded.java
+++ b/src/main/java/redis/clients/util/Sharded.java
@@ -4,11 +4,19 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.TreeMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 public class Sharded<R, S extends ShardInfo<R>> {
     public static final int DEFAULT_WEIGHT = 1;
     private TreeMap<Long, S> nodes;
     private final Hashing algo;
+    
+    /** 
+     * The default pattern used for extracting a key tag. 
+     * The pattern must have a group (between parenthesis), which delimits the tag to be hashed.
+     */
+    private Pattern tagPattern = Pattern.compile("\\{(.+?)\\}"); // the tag is anything between {} 
 
     public Sharded(List<S> shards) {
         this(shards, Hashing.MURMUR_HASH); // MD5 is really not good as we works with 64-bits not 128
@@ -19,6 +27,16 @@ public class Sharded<R, S extends ShardInfo<R>> {
         initialize(shards);
     }
 
+    public Sharded(List<S> shards, Pattern tagPattern) {
+        this(shards, Hashing.MURMUR_HASH, tagPattern); // MD5 is really not good as we works with 64-bits not 128
+    }
+
+    public Sharded(List<S> shards, Hashing algo, Pattern tagPattern) {
+        this.algo = algo;
+        this.tagPattern = tagPattern; 
+        initialize(shards);
+    }
+
     private void initialize(List<S> shards) {
         nodes = new TreeMap<Long, S>();
 
@@ -40,11 +58,26 @@ public class Sharded<R, S extends ShardInfo<R>> {
     }
 
     public R getShard(String key) {
-        return nodes.floorEntry(algo.hash(key)).getValue().getResource();
+        return nodes.floorEntry(algo.hash(getKeyTag(key))).getValue().getResource();
     }
 
     public S getShardInfo(String key) {
-        return nodes.floorEntry(algo.hash(key)).getValue();
+        return nodes.floorEntry(algo.hash(getKeyTag(key))).getValue();
+    }
+    
+    /**
+     * A key tag is a special pattern inside a key that, if preset, is the only part of the key hashed 
+     * in order to select the server for this key.
+     * 
+     * @see http://code.google.com/p/redis/wiki/FAQ#I'm_using_some_form_of_key_hashing_for_partitioning,_but_wh
+     * @param key
+     * @return The tag if it exists, or the original key 
+     */
+    public String getKeyTag(String key){
+    	Matcher m = tagPattern.matcher(key);
+    	if (m.find()) 
+    		return m.group(1); 
+    	return key; 
     }
 
     public Collection<S> getAllShards() {
diff --git a/src/test/java/redis/clients/jedis/tests/ShardedJedisTest.java b/src/test/java/redis/clients/jedis/tests/ShardedJedisTest.java
index f70b2f2..e15451e 100644
--- a/src/test/java/redis/clients/jedis/tests/ShardedJedisTest.java
+++ b/src/test/java/redis/clients/jedis/tests/ShardedJedisTest.java
@@ -82,5 +82,29 @@ public class ShardedJedisTest extends Assert {
 	assertEquals("bar1", j.get("b"));
 	j.disconnect();
     }
+    
+    
+    /** @author muriloq@gmail.com */
+    @Test
+    public void checkKeyTags(){
+    	List<JedisShardInfo> shards = new ArrayList<JedisShardInfo>();
+    	shards.add(new JedisShardInfo(redis1.host, redis1.port));
+    	shards.add(new JedisShardInfo(redis2.host, redis2.port));
+    	ShardedJedis jedis = new ShardedJedis(shards);
+    	
+    	assertEquals(jedis.getKeyTag("foo"),"foo");
+    	assertEquals(jedis.getKeyTag("foo{bar}"),"bar");
+    	assertEquals(jedis.getKeyTag("foo{bar}}"),"bar"); // default pattern is non greedy
+    	assertEquals(jedis.getKeyTag("{bar}foo"),"bar"); // Key tag may appear anywhere
+    	assertEquals(jedis.getKeyTag("f{bar}oo"),"bar"); // Key tag may appear anywhere
+    	
+    	ShardInfo s1 = jedis.getShardInfo("a{bar}");
+    	ShardInfo s2 = jedis.getShardInfo("b{bar}");
+    	assertSame(s1, s2);
+    	
+    	ShardInfo s3 = jedis.getShardInfo("a");
+    	ShardInfo s4 = jedis.getShardInfo("b");
+    	assertNotSame(s3, s4);
+    }
 
 }
\ No newline at end of file