From 8ebc7048197d28a8c0443bdc9338e6948ec2b703 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A1nos=20Benjamin=20Antal?= <benjamin.antal@memgraph.io>
Date: Tue, 25 Oct 2022 14:37:18 +0200
Subject: [PATCH] Fix profile queries with ScanAll

---
 src/query/v2/plan/operator.cpp | 36 ++++++++++++++++++----------------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/src/query/v2/plan/operator.cpp b/src/query/v2/plan/operator.cpp
index 3c36559c9..62a927cfd 100644
--- a/src/query/v2/plan/operator.cpp
+++ b/src/query/v2/plan/operator.cpp
@@ -370,27 +370,29 @@ class DistributedScanAllAndFilterCursor : public Cursor {
   bool Pull(Frame &frame, ExecutionContext &context) override {
     SCOPED_PROFILE_OP(op_name_);
     auto &shard_manager = *context.shard_request_manager;
-    if (MustAbort(context)) {
-      throw HintedAbortError();
-    }
-    using State = msgs::ExecutionState<msgs::ScanVerticesRequest>;
-
-    if (request_state_.state == State::INITIALIZING) {
-      if (!input_cursor_->Pull(frame, context)) {
-        return false;
+    while (true) {
+      if (MustAbort(context)) {
+        throw HintedAbortError();
       }
-    }
+      using State = msgs::ExecutionState<msgs::ScanVerticesRequest>;
 
-    if (current_vertex_it == current_batch.end()) {
-      if (request_state_.state == State::COMPLETED || !MakeRequest(shard_manager)) {
-        ResetExecutionState();
-        return Pull(frame, context);
+      if (request_state_.state == State::INITIALIZING) {
+        if (!input_cursor_->Pull(frame, context)) {
+          return false;
+        }
       }
-    }
 
-    frame[output_symbol_] = TypedValue(std::move(*current_vertex_it));
-    ++current_vertex_it;
-    return true;
+      if (current_vertex_it == current_batch.end()) {
+        if (request_state_.state == State::COMPLETED || !MakeRequest(shard_manager)) {
+          ResetExecutionState();
+          continue;
+        }
+      }
+
+      frame[output_symbol_] = TypedValue(std::move(*current_vertex_it));
+      ++current_vertex_it;
+      return true;
+    }
   }
 
   void Shutdown() override { input_cursor_->Shutdown(); }