From 0d68cb34f76ad61b33690a71192647478b590b26 Mon Sep 17 00:00:00 2001
From: Rodrigo Kumpera <kumpera@fb.com>
Date: Wed, 26 Apr 2023 18:10:10 +0000
Subject: [PATCH] [Store] Introduce Store v2 API.

V2 Store extends with the following methods: append, multi_get and multi_set.
Capabiltities can be probed at runtime and compile time.
---
 README.md                       |  7 ++++++
 gloo/rendezvous/prefix_store.cc | 40 +++++++++++++++++++++++++++++++++
 gloo/rendezvous/prefix_store.h  |  6 +++++
 gloo/rendezvous/store.h         | 25 +++++++++++++++++++++
 4 files changed, 78 insertions(+)

diff --git a/README.md b/README.md
index 777874f0e..3d64d0a51 100644
--- a/README.md
+++ b/README.md
@@ -54,6 +54,13 @@ Ubuntu, this version ships with version 17.10 and up. If you run an
 older version, you'll have to install Google Test yourself, and set
 the `GTEST_ROOT` CMake variable.
 
+You can install Google Test using conda with:
+``` shell
+conda install -c anaconda gmock gtest
+```
+Be carefull that you might need to fish for a package that works with your glibc
+
+
 To build the tests, run:
 
 ``` shell
diff --git a/gloo/rendezvous/prefix_store.cc b/gloo/rendezvous/prefix_store.cc
index 4cdd1d6d5..8065356b0 100644
--- a/gloo/rendezvous/prefix_store.cc
+++ b/gloo/rendezvous/prefix_store.cc
@@ -43,5 +43,45 @@ void PrefixStore::wait(
   store_.wait(joinedKeys, timeout);
 }
 
+bool PrefixStore::has_v2_support() {
+  return store_.has_v2_support();
+}
+
+std::vector<std::vector<char>> PrefixStore::multi_get(const std::vector<std::string>& keys) {
+  if (!store_.has_v2_support()) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("underlying store doesn't support multi_get");
+  }
+  std::vector<std::string> prefixed_keys;
+  for(auto& key : keys) {
+    prefixed_keys.push_back(joinKey(key));
+  }
+  return store_.multi_get(prefixed_keys);
+}
+
+void PrefixStore::multi_set(const std::vector<std::string>& keys, const std::vector<std::vector<char>>& values) {
+  if (!store_.has_v2_support()) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("underlying store doesn't support multi_set");
+  }
+  std::vector<std::string> prefixed_keys;
+  for(auto& key : keys) {
+    prefixed_keys.push_back(joinKey(key));
+  }
+  return store_.multi_set(prefixed_keys, values);
+}
+
+void PrefixStore::append(const std::string& key, const std::vector<char>& data) {
+  if (!store_.has_v2_support()) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("underlying store doesn't support append");
+  }
+  store_.append(joinKey(key), data);
+}
+
+int64_t PrefixStore::add(const std::string& key, int64_t value) {
+  if (!store_.has_v2_support()) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("underlying store doesn't support append");
+  }
+  return store_.add(joinKey(key), value);
+  }
+
 } // namespace rendezvous
 } // namespace gloo
diff --git a/gloo/rendezvous/prefix_store.h b/gloo/rendezvous/prefix_store.h
index 3c4271fcc..aa5bfeb87 100644
--- a/gloo/rendezvous/prefix_store.h
+++ b/gloo/rendezvous/prefix_store.h
@@ -34,6 +34,12 @@ class PrefixStore : public Store {
       const std::vector<std::string>& keys,
       const std::chrono::milliseconds& timeout) override;
 
+  virtual bool has_v2_support() override;
+  virtual std::vector<std::vector<char>> multi_get(const std::vector<std::string>& keys) override;
+  virtual void multi_set(const std::vector<std::string>& keys, const std::vector<std::vector<char>>& values) override;
+  virtual void append(const std::string& key, const std::vector<char>& data) override;
+  virtual int64_t add(const std::string& key, int64_t value) override;
+
  protected:
   const std::string prefix_;
   Store& store_;
diff --git a/gloo/rendezvous/store.h b/gloo/rendezvous/store.h
index f1286a789..da315a964 100644
--- a/gloo/rendezvous/store.h
+++ b/gloo/rendezvous/store.h
@@ -13,6 +13,10 @@
 #include <vector>
 
 #include "gloo/common/logging.h"
+#include "gloo/common/error.h"
+
+//can be used by upstream users to know whether this is available or not.
+#define GLOO_STORE_HAS_STORE_V2 1
 
 namespace gloo {
 namespace rendezvous {
@@ -39,6 +43,27 @@ class Store {
     wait(keys);
   }
 
+  virtual bool has_v2_support() {
+    // If True, the following operations are guaranteed to be efficiently and correclty implemented.
+    return false;
+  }
+
+  virtual std::vector<std::vector<char>> multi_get(const std::vector<std::string>& /*keys*/) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("this store doesn't support multi_get");
+  }
+
+  virtual void multi_set(const std::vector<std::string>& /*keys*/, const std::vector<std::vector<char>>& /*values*/) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("this store doesn't support multi_set");
+  }
+
+  virtual void append(const std::string& key, const std::vector<char>& /*data*/) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("this store doesn't support append");
+  }
+
+  virtual int64_t add(const std::string& key, int64_t value) {
+    GLOO_THROW_INVALID_OPERATION_EXCEPTION("this store doesn't support add");
+  }
+
 };
 
 } // namespace rendezvous