hg: add add/removeStore to cuniondatapackstore

Summary:
In a future diff we'll need the ability to modify the union store on
the fly, so let's add addstore and removestore apis.

Reviewed By: ryanmce

Differential Revision: D7051102

fbshipit-source-id: 901a50720bfdf4e5c59714d092830e65edccdfce

hgext/extlib/cstore/py-datapackstore.h

@@ -29,6 +29,7 @@ extern "C" {
 #include "hgext/extlib/cstore/pythonkeyiterator.h"
 #include "hgext/extlib/cstore/pythonutil.h"
 #include "hgext/extlib/cstore/uniondatapackstore.h"
+#include "hgext/extlib/cstore/util.h"
 
 // --------- DatapackStore Implementation ---------
 
@@ -206,6 +207,41 @@ static PyTypeObject datapackstoreType = {
 
 // --------- UnionDatapackStore Implementation ---------
 
+static void addStore(py_uniondatapackstore *self, PythonObj store)
+{
+  PyObject *rawStore = (PyObject *)store;
+
+  // Record the substore references, so:
+  // A) We can decref them in case of an error.
+  // B) They don't get GC'd while the uniondatapackstore holds on to them.
+  int iscdatapack =
+      PyObject_IsInstance(rawStore, (PyObject *)&datapackstoreType);
+
+  switch (iscdatapack) {
+  case 1:
+    // Store is C datapack
+    {
+      self->cstores.push_back(store);
+      py_datapackstore *subStore = (py_datapackstore *)rawStore;
+      self->uniondatapackstore->addStore(&subStore->datapackstore);
+    }
+    break;
+  case 0:
+    // Store is PythonDataStore, it's memory management
+    // is performed by py_uniondatapackstore
+    {
+      std::shared_ptr<PythonDataStore> pystore =
+          std::make_shared<PythonDataStore>(store);
+      self->pystores.push_back(pystore);
+      self->uniondatapackstore->addStore(pystore.get());
+    }
+    break;
+  default:
+    // Error
+    throw std::logic_error("invalid store type passed to addStore");
+  }
+}
+
 /*
  * Initializes the contents of a uniondatapackstore
  */
@@ -217,54 +253,19 @@ static int uniondatapackstore_init(py_uniondatapackstore *self, PyObject *args)
   }
 
   try {
-    std::vector<DataStore *> stores;
-    std::vector<PythonObj> cSubStores;
-    std::vector<std::shared_ptr<PythonDataStore>> pySubStores;
+    // We have to manually call the member constructor, since the provided
+    // 'self' is just zerod out memory.
+    new (&self->uniondatapackstore)
+        std::shared_ptr<UnionDatapackStore>(new UnionDatapackStore());
+    new (&self->cstores) std::vector<PythonObj>();
+    new (&self->pystores) std::vector<std::shared_ptr<PythonDataStore>>();
 
     PyObject *item;
     PythonObj inputIterator = PyObject_GetIter(storeList);
     while ((item = PyIter_Next(inputIterator)) != NULL) {
-      // Record the substore references, so:
-      // A) We can decref them in case of an error.
-      // B) They don't get GC'd while the uniondatapackstore holds on to them.
-      int iscdatapack =
-          PyObject_IsInstance(item, (PyObject *)&datapackstoreType);
-
       PythonObj store(item);
-      switch (iscdatapack) {
-      case 1:
-        // Store is C datapack
-        {
-          cSubStores.push_back(store);
-          py_datapackstore *subStore = (py_datapackstore *)item;
-          stores.push_back(&subStore->datapackstore);
-        }
-        break;
-      case 0:
-        // Store is PythonDataStore, it's memory management
-        // is performed by py_uniondatapackstore
-        {
-          std::shared_ptr<PythonDataStore> pystore =
-              std::make_shared<PythonDataStore>(store);
-          pySubStores.push_back(pystore);
-          stores.push_back(pystore.get());
-        }
-        break;
-      default:
-        // Error
-        return -1;
-      }
+      addStore(self, store);
     }
-
-    // We have to manually call the member constructor, since the provided
-    // 'self' is just zerod out memory.
-    new (&self->uniondatapackstore)
-        std::shared_ptr<UnionDatapackStore>(new UnionDatapackStore(stores));
-    new (&self->cstores) std::vector<PythonObj>();
-    new (&self->pystores) std::vector<std::shared_ptr<PythonDataStore>>();
-
-    self->cstores = cSubStores;
-    self->pystores = pySubStores;
   } catch (const std::exception &ex) {
     PyErr_SetString(PyExc_RuntimeError, ex.what());
     return -1;
@@ -281,6 +282,73 @@ static void uniondatapackstore_dealloc(py_uniondatapackstore *self)
   PyObject_Del(self);
 }
 
+static PyObject *uniondatapackstore_addStore(py_uniondatapackstore *self,
+                                             PyObject *storeObj)
+{
+  try {
+    PythonObj store(storeObj);
+    Py_INCREF(storeObj);
+
+    addStore(self, store);
+    Py_RETURN_NONE;
+  } catch (const pyexception &ex) {
+    return NULL;
+  } catch (const std::exception &ex) {
+    PyErr_SetString(PyExc_RuntimeError, ex.what());
+    return NULL;
+  }
+}
+
+static PyObject *uniondatapackstore_removeStore(py_uniondatapackstore *self,
+                                                PyObject *storeObj)
+{
+  try {
+    PythonObj store(storeObj);
+    Py_INCREF(storeObj);
+
+    // Record the substore references, so:
+    // A) We can decref them in case of an error.
+    // B) They don't get GC'd while the uniondatapackstore holds on to them.
+    int iscdatapack =
+        PyObject_IsInstance(storeObj, (PyObject *)&datapackstoreType);
+
+    switch (iscdatapack) {
+    case 1:
+      // Store is C datapack
+      {
+        removeFromVector(self->cstores, store);
+        py_datapackstore *subStore = (py_datapackstore *)storeObj;
+        self->uniondatapackstore->removeStore(&subStore->datapackstore);
+      }
+      break;
+    case 0:
+      // Store is PythonDataStore, it's memory management
+      // is performed by py_uniondatapackstore
+      {
+        auto pystores = self->pystores;
+        for (auto it = pystores.begin(); it != pystores.end(); ++it) {
+          if ((*it)->getStore() == store) {
+            self->uniondatapackstore->removeStore((*it).get());
+            removeFromVector(self->pystores, *it);
+            break;
+          }
+        }
+      }
+      break;
+    default:
+      // Error
+      throw std::logic_error("invalid store type");
+    }
+
+    Py_RETURN_NONE;
+  } catch (const pyexception &ex) {
+    return NULL;
+  } catch (const std::exception &ex) {
+    PyErr_SetString(PyExc_RuntimeError, ex.what());
+    return NULL;
+  }
+}
+
 static PyObject *uniondatapackstore_get(py_uniondatapackstore *self,
                                         PyObject *args)
 {
@@ -405,6 +473,8 @@ static PyObject *uniondatapackstore_getmetrics(py_uniondatapackstore *self)
 
 static PyMethodDef uniondatapackstore_methods[] = {
     {"get", (PyCFunction)uniondatapackstore_get, METH_VARARGS, ""},
+    {"addstore", (PyCFunction)uniondatapackstore_addStore, METH_O, ""},
+    {"removestore", (PyCFunction)uniondatapackstore_removeStore, METH_O, ""},
     {"getdeltachain", (PyCFunction)uniondatapackstore_getdeltachain,
      METH_VARARGS, ""},
     {"getmissing", (PyCFunction)uniondatapackstore_getmissing, METH_O, ""},

hgext/extlib/cstore/pythondatastore.cpp

@@ -122,3 +122,8 @@ bool PythonDataStore::contains(const Key &key)
   std::shared_ptr<KeyIterator> it = getMissing(iter);
   return (!it->next());
 }
+
+PythonObj PythonDataStore::getStore()
+{
+  return this->_store;
+}

hgext/extlib/cstore/pythondatastore.h

@@ -73,6 +73,8 @@ public:
 
   ~PythonDataStore() = default;
 
+  PythonObj getStore();
+
   DeltaChainIterator getDeltaChain(const Key &key);
 
   std::shared_ptr<KeyIterator> getMissing(KeyIterator &missing);

hgext/extlib/cstore/pythonutil.cpp

@@ -44,6 +44,11 @@ PythonObj &PythonObj::operator=(const PythonObj &other)
   return *this;
 }
 
+bool PythonObj::operator==(const PythonObj &other) const
+{
+  return this->obj == other.obj;
+}
+
 PythonObj::operator PyObject *() const
 {
   return this->obj;

hgext/extlib/cstore/pythonutil.h

@@ -59,6 +59,8 @@ public:
 
   PythonObj &operator=(const PythonObj &other);
 
+  bool operator==(const PythonObj &other) const;
+
   operator PyObject *() const;
 
   operator bool() const;

hgext/extlib/cstore/uniondatapackstore.cpp

@@ -8,6 +8,7 @@
 // no-check-code
 
 #include "hgext/extlib/cstore/uniondatapackstore.h"
+#include "hgext/extlib/cstore/util.h"
 
 #include <algorithm>
 #include <memory>
@@ -16,8 +17,12 @@ extern "C" {
 #include "mercurial/mpatch.h"
 }
 
-UnionDatapackStore::UnionDatapackStore(std::vector<DataStore *> stores)
-    : _stores(stores)
+UnionDatapackStore::UnionDatapackStore()
+{
+}
+
+UnionDatapackStore::UnionDatapackStore(std::vector<DataStore*> &stores)
+  : _stores(stores)
 {
 }
 
@@ -150,3 +155,13 @@ void UnionDatapackStore::markForRefresh()
     substore->markForRefresh();
   }
 }
+
+void UnionDatapackStore::addStore(DataStore *store)
+{
+  _stores.push_back(store);
+}
+
+void UnionDatapackStore::removeStore(DataStore *store)
+{
+  removeFromVector<DataStore *>(_stores, store);
+}

hgext/extlib/cstore/uniondatapackstore.h

@@ -59,7 +59,9 @@ class UnionDatapackStore : public Store
 public:
   std::vector<DataStore *> _stores;
 
-  UnionDatapackStore(std::vector<DataStore *> stores);
+  UnionDatapackStore();
+
+  UnionDatapackStore(std::vector<DataStore*> &stores);
 
   ~UnionDatapackStore() override;
 
@@ -72,6 +74,9 @@ public:
   UnionDatapackStoreKeyIterator getMissing(KeyIterator &missing);
 
   void markForRefresh();
+
+  void addStore(DataStore *store);
+  void removeStore(DataStore *store);
 };
 
 #endif // FBHGEXT_CSTORE_UNIONDATAPACKSTORE_H

hgext/extlib/cstore/util.h

@@ -0,0 +1,9 @@
+template <class T> void removeFromVector(std::vector<T> &vec, T &item)
+{
+  for (auto it = vec.begin(); it != vec.end(); ++it) {
+    if (*it == item) {
+      vec.erase(it);
+      break;
+    }
+  }
+}

hgext/remotefilelog/contentstore.py

@@ -35,7 +35,7 @@ class ChainIndicies(object):
 
 class unioncontentstore(object):
     def __init__(self, *args, **kwargs):
-        self.stores = args
+        self.stores = list(args)
         self.writestore = kwargs.get('writestore')
 
         # If allowincomplete==True then the union store can return partial
@@ -154,6 +154,12 @@ class unioncontentstore(object):
             if util.safehasattr(store, 'markforrefresh'):
                 store.markforrefresh()
 
+    def addstore(self, store):
+        self.stores.append(store)
+
+    def removestore(self, store):
+        self.stores.remove(store)
+
 class remotefilelogcontentstore(basestore.basestore):
     def __init__(self, *args, **kwargs):
         super(remotefilelogcontentstore, self).__init__(*args, **kwargs)

tests/test-fb-hgext-cstore-uniondatapackstore.py

@@ -155,6 +155,36 @@ class uniondatapackstoretests(unittest.TestCase):
         self.assertEquals(set([("foo", missinghash1), ("foo2", missinghash2)]),
                           set(missing))
 
+    def testAddRemoveStore(self):
+        packdir = self.makeTempDir()
+        revisions = [("foo", self.getFakeHash(), nullid, "content")]
+        store = self.createPackStore(packdir, revisions=revisions)
+
+        packdir2 = self.makeTempDir()
+        revisions2 = [("foo2", self.getFakeHash(), nullid, "content2")]
+        store2 = self.createPackStore(packdir2, revisions=revisions2)
+
+        unionstore = uniondatapackstore([store])
+        unionstore.addstore(store2)
+
+        # Fetch from store2
+        result = unionstore.get('foo2', revisions2[0][1])
+        self.assertEquals(result, revisions2[0][3])
+
+        # Drop the store
+        unionstore.removestore(store2)
+
+        # Fetch from store1
+        result = unionstore.get('foo', revisions[0][1])
+        self.assertEquals(result, revisions[0][3])
+
+        # Fetch from missing store2
+        try:
+            unionstore.get('foo2', revisions2[0][1])
+            self.asserFalse(True, "get should've thrown")
+        except KeyError:
+            pass
+
 class uniondatastorepythontests(uniondatapackstoretests):
     def createPackStore(self, packdir, revisions=None):
         if revisions is None: