make pregel tasks GET API load-balancing aware

arangodb · jsteemann · Jul 1, 2021 · Jun 1, 2021 · Jun 28, 2021 · Jun 29, 2021
commit 11e7d47dd0259b9a08c90b408d89c1ae8437dc89
diff --git a/arangod/Pregel/Conductor.cpp b/arangod/Pregel/Conductor.cpp
@@ -899,7 +899,7 @@ void Conductor::toVelocyPack(VPackBuilder& result) const {
     result.add("algorithm", VPackValue(_algorithm->name()));
   }
   result.add("created", VPackValue(timepointToString(_created)));
-  if (_state != ExecutionState::DEFAULT) {
+  if (_expires != std::chrono::system_clock::time_point{}) {
     result.add("expires", VPackValue(timepointToString(_expires)));
   }
   result.add("ttl", VPackValue(_ttl.count()));
@@ -919,7 +919,10 @@ void Conductor::toVelocyPack(VPackBuilder& result) const {
     result.add("vertexCount", VPackValue(_totalVerticesCount));
     result.add("edgeCount", VPackValue(_totalEdgesCount));
   }
-  result.add("parallelism", _userParams.slice().get(Utils::parallelismKey));
+  VPackSlice p = _userParams.slice().get(Utils::parallelismKey);
+  if (!p.isNone()) {
+    result.add("parallelism", p);
+  }
   if (_masterContext) {
     VPackObjectBuilder ob(&result, "masterContext");
     _masterContext->serializeValues(result);

diff --git a/arangod/Pregel/PregelFeature.cpp b/arangod/Pregel/PregelFeature.cpp
@@ -33,6 +33,9 @@
 #include "Cluster/ClusterInfo.h"
 #include "Cluster/ServerState.h"
 #include "FeaturePhases/V8FeaturePhase.h"
+#include "GeneralServer/AuthenticationFeature.h"
+#include "Network/Methods.h"
+#include "Network/NetworkFeature.h"
 #include "Pregel/AlgoRegistry.h"
 #include "Pregel/Conductor.h"
 #include "Pregel/Recovery.h"
@@ -44,6 +47,9 @@
 #include "VocBase/LogicalCollection.h"
 #include "VocBase/ticks.h"
 
+using namespace arangodb;
+using namespace arangodb::pregel;
+
 namespace {
 bool authorized(std::string const& user) {
   auto const& exec = arangodb::ExecContext::current();
@@ -53,6 +59,17 @@ bool authorized(std::string const& user) {
   return (user == exec.user());
 }
 
+network::Headers buildHeaders() {
+  auto auth = AuthenticationFeature::instance();
+
+  network::Headers headers;
+  if (auth != nullptr && auth->isActive()) {
+    headers.try_emplace(StaticStrings::Authorization,
+                        "bearer " + auth->tokenCache().jwtToken());
+  }
+  return headers;
+}
+
 /// @brief custom deleter for the PregelFeature.
 /// it does nothing, i.e. doesn't delete it. This is because the ApplicationServer
 /// is managing the PregelFeature, but we need a shared_ptr to it here as well. The
@@ -64,9 +81,6 @@ struct NonDeleter {
 
 }  // namespace
 
-using namespace arangodb;
-using namespace arangodb::pregel;
-
 std::pair<Result, uint64_t> PregelFeature::startExecution(
     TRI_vocbase_t& vocbase, std::string algorithm,
     std::vector<std::string> const& vertexCollections,
@@ -499,7 +513,11 @@ uint64_t PregelFeature::numberOfActiveConductors() const {
   return nr;
 }
 
-void PregelFeature::toVelocyPack(arangodb::velocypack::Builder& result) const {
+Result PregelFeature::toVelocyPack(TRI_vocbase_t& vocbase,
+                                   arangodb::velocypack::Builder& result, 
+                                   bool allDatabases, bool fanout) const {
+  Result res;
+
   result.openArray();
   {
     MUTEX_LOCKER(guard, _mutex);
@@ -513,5 +531,62 @@ void PregelFeature::toVelocyPack(arangodb::velocypack::Builder& result) const {
       ce.conductor->toVelocyPack(result);
     }
   }
+
+  if (ServerState::instance()->isCoordinator() && fanout) {
+    // coordinator case, fan out to other coordinators!
+    NetworkFeature const& nf = vocbase.server().getFeature<NetworkFeature>();
+    network::ConnectionPool* pool = nf.pool();
+    if (pool == nullptr) {
+      THROW_ARANGO_EXCEPTION(TRI_ERROR_SHUTTING_DOWN);
+    }
+
+    std::vector<network::FutureRes> futures;
+
+    network::RequestOptions options;
+    options.timeout = network::Timeout(30.0);
+    options.database = vocbase.name();
+    options.param("local", "true");
+    options.param("all", allDatabases ? "true" : "false");
+
+    std::string const url = "/_api/control_pregel";
+
+    auto& ci = vocbase.server().getFeature<ClusterFeature>().clusterInfo();
+    for (auto const& coordinator : ci.getCurrentCoordinators()) {
+      if (coordinator == ServerState::instance()->getId()) {
+        // ourselves!
+        continue;
+      }
+
      auto f = network::sendRequestRetry(pool, "server:" + coordinator, fuerte::RestVerb::Get,
+                                    url, VPackBuffer<uint8_t>{}, options, ::buildHeaders());
+      futures.emplace_back(std::move(f));
+    }
+
+    if (!futures.empty()) {
+      auto responses = futures::collectAll(futures).get();
+      for (auto const& it : responses) {
+        auto& resp = it.get();
+        res.reset(resp.combinedResult());
+        if (res.is(TRI_ERROR_ARANGO_DATABASE_NOT_FOUND)) {
+          // it is expected in a multi-coordinator setup that a coordinator is not
+          // aware of a database that was created very recently.
+          res.reset();
+        }
+        if (res.fail()) {
+          break;
+        }
+        auto slice = resp.slice();
+        // copy results from other coordinators
+        if (slice.isArray()) {
+          for (auto const& entry : VPackArrayIterator(slice)) {
+            result.add(entry);
+          }
+        }
+      }
+    }
+  } 
+
   result.close();
+
+  return res;
 }
diff --git a/arangod/Pregel/PregelFeature.h b/arangod/Pregel/PregelFeature.h
@@ -95,7 +95,8 @@ class PregelFeature final : public application_features::ApplicationFeature {
     _softShutdownOngoing.store(true, std::memory_order_relaxed);
   }
 
-  void toVelocyPack(arangodb::velocypack::Builder& result) const;
+  Result toVelocyPack(TRI_vocbase_t& vocbase, arangodb::velocypack::Builder& result,
+                      bool allDatabases, bool fanout) const;
 
  private:
   void scheduleGarbageCollection();

diff --git a/arangod/RestHandler/RestControlPregelHandler.cpp b/arangod/RestHandler/RestControlPregelHandler.cpp
@@ -190,8 +190,11 @@ void RestControlPregelHandler::getExecutionStatus() {
   std::vector<std::string> const& suffixes = _request->decodedSuffixes();
 
   if (suffixes.empty()) {
+    bool const allDatabases = _request->parsedValue("all", false);
+    bool const fanout = ServerState::instance()->isCoordinator() && !_request->parsedValue("local", false);
+
     VPackBuilder builder;
-    _pregel.toVelocyPack(builder);
+    _pregel.toVelocyPack(_vocbase, builder, allDatabases, fanout);
     generateResult(rest::ResponseCode::OK, builder.slice());
     return;
   }

diff --git a/arangod/V8Server/v8-collection.cpp b/arangod/V8Server/v8-collection.cpp
@@ -1716,7 +1716,7 @@ static void JS_PregelStatus(v8::FunctionCallbackInfo<v8::Value> const& args) {
   // check the arguments
   uint32_t const argLength = args.Length();
   if (argLength == 0) {
-    pregel.toVelocyPack(builder);
+    pregel.toVelocyPack(vocbase, builder, false, true);
     TRI_V8_RETURN(TRI_VPackToV8(isolate, builder.slice()));
     return;
   }

diff --git a/tests/js/client/load-balancing/load-balancing-pregel-noauth-cluster.js b/tests/js/client/load-balancing/load-balancing-pregel-noauth-cluster.js
@@ -1,5 +1,5 @@
 /* jshint globalstrict:true, strict:true, maxlen: 5000 */
-/* global assertTrue, assertFalse, assertEqual, require*/
+/* global assertTrue, assertFalse, assertEqual, assertNotEqual, require*/
 
 ////////////////////////////////////////////////////////////////////////////////
 /// DISCLAIMER
@@ -117,6 +117,66 @@ function PregelSuite () {
       cs = [];
       coordinators = [];
     },
+
+    testPregelAllJobs: function() {
+      let url = baseUrl;
+      const task = {
+        algorithm: "pagerank",
+        vertexCollections: ["vertices"],
+        edgeCollections: ["edges"],
+        params: {
+          resultField: "result",
+          store: false
+        },
+        store: false
+      };
+      let ids = [];
+      for (let i = 0; i < 5; ++i) {
+        let result = sendRequest('POST', url, task, i % 2 === 0);
+
+        assertFalse(result === undefined || result === {});
+        assertEqual(result.status, 200);
+        assertTrue(result.body > 1);
+        ids.push(result.body);
+      }
+      assertEqual(5, ids.length);
+
+      let result = sendRequest('GET', url, {}, false);
+      assertEqual(result.status, 200);
+      assertEqual(5, result.body.length);
+      result.body.forEach((task) => {
+        assertEqual("PageRank", task.algorithm);
+        assertEqual(db._name(), task.database);
+        assertNotEqual(-1, ids.indexOf(task.id));
+      });
+
+      ids.forEach((id) => {
+        const taskId = id;
+        url = `${baseUrl}/${taskId}`;
+        result = sendRequest('GET', url, {}, false);
+
+        assertFalse(result === undefined || result === {});
+        assertEqual(result.status, 200);
+        assertTrue(result.body.state === 'running' || result.body.state === 'done');
+      });
+
+      require('internal').wait(5.0, false);
+
+      ids.forEach((id) => {
+        const taskId = id;
+        url = `${baseUrl}/${taskId}`;
+        result = sendRequest('DELETE', url, {}, {}, false);
+
+        assertFalse(result === undefined || result === {});
+        assertFalse(result.body.error);
+        assertEqual(result.status, 200);
+      });
+
+      url = baseUrl;
+      result = sendRequest('GET', url, {}, false);
+      assertEqual(result.status, 200);
+      assertEqual([], result.body);
+    },
 
     testPregelForwarding: function() {
       let url = baseUrl;