From 77d7600b649ef1c109933796167d20dd2eb4577d Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Mon, 18 Feb 2019 16:34:43 -0500
Subject: [PATCH] Add ability for proportional rollout of the OCI data model

---
 app.py                                     |  3 ++-
 data/registry_model/__init__.py            | 10 +++++-----
 data/registry_model/modelsplitter.py       | 16 ++++++++++++++--
 data/registry_model/test/test_interface.py |  2 +-
 4 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/app.py b/app.py
index 9926ee766..312ab8f7e 100644
--- a/app.py
+++ b/app.py
@@ -93,7 +93,8 @@ app.config.update(environ_config)
 
 # Split the registry model based on config.
 # TODO(jschorr): Remove once we are fully on the OCI data model.
-registry_model.setup_split(app.config.get('OCI_NAMESPACE_WHITELIST') or set(),
+registry_model.setup_split(app.config.get('OCI_NAMESPACE_PROPORTION') or 0,
+                           app.config.get('OCI_NAMESPACE_WHITELIST') or set(),
                            app.config.get('V22_NAMESPACE_WHITELIST') or set())
 
 # Allow user to define a custom storage preference for the local instance.
diff --git a/data/registry_model/__init__.py b/data/registry_model/__init__.py
index fead18217..8fe882f1e 100644
--- a/data/registry_model/__init__.py
+++ b/data/registry_model/__init__.py
@@ -12,15 +12,15 @@ class RegistryModelProxy(object):
   def __init__(self):
     self._model = oci_model if os.getenv('OCI_DATA_MODEL') == 'true' else pre_oci_model
 
-  def setup_split(self, oci_whitelist, v22_whitelist):
+  def setup_split(self, oci_model_proportion, oci_whitelist, v22_whitelist):
     if os.getenv('OCI_DATA_MODEL') == 'true':
       return
-      
+
     logger.info('===============================')
-    logger.info('Enabling split registry model with OCI whitelist `%s` and V22 whitelist `%s`',
-                oci_whitelist, v22_whitelist)
+    logger.info('Split registry model: OCI %s proportion and whitelist `%s` and V22 whitelist `%s`',
+                oci_model_proportion, oci_whitelist, v22_whitelist)
     logger.info('===============================')
-    self._model = SplitModel(oci_whitelist, v22_whitelist)
+    self._model = SplitModel(oci_model_proportion, oci_whitelist, v22_whitelist)
 
   def set_for_testing(self, use_oci_model):
     self._model = oci_model if use_oci_model else pre_oci_model
diff --git a/data/registry_model/modelsplitter.py b/data/registry_model/modelsplitter.py
index 8c61a7b35..349d8e8cf 100644
--- a/data/registry_model/modelsplitter.py
+++ b/data/registry_model/modelsplitter.py
@@ -1,5 +1,6 @@
 import inspect
 import logging
+import hashlib
 
 from data.database import DerivedStorageForImage, TagManifest, Manifest, Image
 from data.registry_model.registry_oci_model import oci_model
@@ -11,12 +12,14 @@ logger = logging.getLogger(__name__)
 
 
 class SplitModel(object):
-  def __init__(self, oci_namespace_whitelist, v22_namespace_whitelist):
+  def __init__(self, oci_model_proportion, oci_namespace_whitelist, v22_namespace_whitelist):
     self.v22_namespace_whitelist = set(v22_namespace_whitelist)
 
     self.oci_namespace_whitelist = set(oci_namespace_whitelist)
     self.oci_namespace_whitelist.update(v22_namespace_whitelist)
 
+    self.oci_model_proportion = oci_model_proportion
+
   def supports_schema2(self, namespace_name):
     """ Returns whether the implementation of the data interface supports schema 2 format
         manifests. """
@@ -75,7 +78,16 @@ class SplitModel(object):
         args_dict = {argnames[index + 1]: value for index, value in enumerate(args)}
 
       namespace_name = self._namespace_from_kwargs(args_dict)
-      if namespace_name in self.oci_namespace_whitelist:
+      use_oci = namespace_name in self.oci_namespace_whitelist
+      if not use_oci and self.oci_model_proportion:
+        # Hash the namespace name and see if it falls into the proportion bucket.
+        bucket = (int(hashlib.md5(namespace_name).hexdigest(), 16) % 100)
+        if bucket <= int(self.oci_model_proportion * 100):
+          logger.debug('Enabling OCI for namespace `%s` in proportional bucket',
+                       namespace_name)
+          use_oci = True
+
+      if use_oci:
         logger.debug('Calling method `%s` under OCI data model for namespace `%s`',
                      attr, namespace_name)
         return getattr(oci_model, attr)(*args, **kwargs)
diff --git a/data/registry_model/test/test_interface.py b/data/registry_model/test/test_interface.py
index c20bb28ff..747a21eb1 100644
--- a/data/registry_model/test/test_interface.py
+++ b/data/registry_model/test/test_interface.py
@@ -36,7 +36,7 @@ from test.fixtures import *
 
 
 @pytest.fixture(params=[PreOCIModel(), OCIModel(),
-                        SplitModel({'devtable'}, {'buynlarge'})])
+                        SplitModel(0, {'devtable'}, {'buynlarge'})])
 def registry_model(request, initialized_db):
   return request.param