TransitionalCache.transform_memcache_key: Don't prepend prefix to key

When doing an operation with multiple keys a prefix can be specified, but the
returned value should exclude the prefix and just include the bare keys.

r2/r2/lib/app_globals.py

@@ -882,19 +882,24 @@ class Globals(object):
             self.maincache = CacheChain((localcache_cls(), self.mcrouter))
         cache_chains.update(maincache=self.maincache)
 
-        def rewrite_subreddit_key(key, prefix=''):
+        def get_new_subreddit_prefix_and_key(key, prefix=''):
             old_prefix = "Subreddit_"
             new_prefix = "sr:"
-            key = prefix + str(key)
-            assert key.startswith(old_prefix)
-            sr_id = key[len(old_prefix):]
-            return new_prefix + sr_id
+
+            if prefix:
+                assert prefix == old_prefix
+                return new_prefix, key
+            else:
+                key = str(key)
+                assert key.startswith(old_prefix)
+                sr_id = key[len(old_prefix):]
+                return '', new_prefix + sr_id
 
         self.transitionalcache = TransitionalCache(
             original_cache=self.cache,
             replacement_cache=self.maincache,
             read_original=True,
-            key_transform=rewrite_subreddit_key,
+            key_transform=get_new_subreddit_prefix_and_key,
         )
 
         if stalecaches:

r2/r2/lib/cache.py

@@ -522,39 +522,88 @@ class TransitionalCache(CacheUtils):
             return self.replacement.caches
 
     def transform_memcache_key(self, args, kwargs):
+        """Use key_transform to transform keys and prefix.
+
+        key_transform() returns (new_prefix, new_key)
+
+        If "prefix" is specified in kwargs, the transformation will look like:
+        key_transform("key", "old_prefix_") --> "new_prefix_", "key"
+
+        If "prefix" is not specified in kwargs, it must already be part of the
+        key, and the transformation looks like:
+        key_transform("old_prefix_key") --> "", "new_prefix_key"
+
+        We don't currently handle multiple gets or sets where the prefix is
+        already prepended to the keys because the return values are different:
+        get(["old_prefix_A", "old_prefix_B"])
+        old:
+            {"old_prefix_A": val, "old_prefix_B": val}
+        new:
+            {"new_prefix_A": val, "new_prefix_B": val}
+
+        They must be looked up with a prefix:
+        get(["A", "B"], prefix="old_prefix_")
+        old:
+            {"A": val, "B": val}
+        new (translated to get(["A", "B"], prefix="new_prefix_"):
+            {"A": val, "B": val}
+
+        The special case of the above is for a single item lookup, where the
+        return value does not include the key.
+
+        We could handle the general multiple key case by maintaining a mapping
+        of {old_key: new_key} and using that to transform the return value.
+
+        """
+
         if self.key_transform:
-            old_key = args[0]
             prefix = kwargs.get("prefix", "")
+            new_kwargs = copy(kwargs)
 
-            if isinstance(old_key, dict):  # multiget passes a dict
-                new_key = {
-                    self.key_transform(k, prefix): v
-                    for k, v in old_key.iteritems()
-                }
-            elif isinstance(old_key, (list, set, tuple)):
-                new_key = [self.key_transform(k, prefix) for k in old_key]
+            if isinstance(args[0], dict):
+                assert prefix, "must include prefix"
+                new_prefixes = []
+                old_key_dict = args[0]
+                new_key_dict = {}
+
+                for old_key, val in old_key_dict.iteritems():
+                    new_prefix, new_key = self.key_transform(old_key, prefix)
+                    new_key_dict[new_key] = val
+                    new_prefixes.append(new_prefix)
+
+                assert all(p == new_prefixes[0] for p in new_prefixes[1:])
+                new_kwargs["prefix"] = new_prefixes[0]
+                new_args = (new_key_dict,) + args[1:]
+            elif isinstance(args[0], (list, set, tuple)):
+                assert prefix, "must include prefix"
+                new_prefixes = []
+                old_key_list = args[0]
+                new_key_list = []
+
+                for old_key in old_key_list:
+                    new_prefix, new_key = self.key_transform(old_key, prefix)
+                    new_key_list.append(new_key)
+                    new_prefixes.append(new_prefix)
+
+                assert all(p == new_prefixes[0] for p in new_prefixes[1:])
+                new_kwargs["prefix"] = new_prefixes[0]
+                new_args = (new_key_list,) + args[1:]
             else:
-                new_key = self.key_transform(old_key, prefix)
+                # single keys can't specify a prefix
+                _, new_key = self.key_transform(args[0])
+                new_args = (new_key,) + args[1:]
 
-            return (new_key,) + args[1:]
+            return new_args, new_kwargs
         else:
-            return args
+            return args, kwargs
 
     def make_get_fn(fn_name):
         def transitional_cache_get_fn(self, *args, **kwargs):
             if self.read_original:
                 return getattr(self.original, fn_name)(*args, **kwargs)
             else:
-                args = self.transform_memcache_key(args, kwargs)
-
-                # remove the prefix argument because the transform fn has
-                # already tacked that onto the keys
-                try:
-                    kwargs.pop("prefix")
-                except KeyError:
-                    pass
-
-                return getattr(self.replacement, fn_name)(*args, **kwargs)
+                new_args, new_kwargs = self.transform_memcache_key(args, kwargs)
+                return getattr(self.replacement, fn_name)(*new_args, **new_kwargs)
         return transitional_cache_get_fn
 
     get = make_get_fn("get")
@@ -564,16 +613,10 @@ class TransitionalCache(CacheUtils):
     def make_set_fn(fn_name):
         def transitional_cache_set_fn(self, *args, **kwargs):
             ret_original = getattr(self.original, fn_name)(*args, **kwargs)
-            args = self.transform_memcache_key(args, kwargs)
 
-            # remove the prefix argument because the transform fn has
-            # already tacked that onto the keys
-            try:
-                kwargs.pop("prefix")
-            except KeyError:
-                pass
+            new_args, new_kwargs = self.transform_memcache_key(args, kwargs)
+            ret_replacement = getattr(self.replacement, fn_name)(*new_args, **new_kwargs)
 
-            ret_replacement = getattr(self.replacement, fn_name)(*args, **kwargs)
             if self.read_original:
                 return ret_original
             else: