From 0bf50078b1227275ed764ec2007f0579326bce5e Mon Sep 17 00:00:00 2001
From: Sergei Politov <spolitov@users.noreply.github.com>
Date: Fri, 27 Apr 2018 12:51:18 +0300
Subject: [PATCH] ENG-3197: Fail GetTableLocations when some tablets are not
 running

Summary:
Added flag to GetTableLocations, that requires that all involved tablets should be running.
Otherwise method fails.
This flag is used during open table.

Also created test that catches this issue.

Test Plan: ybd --cxx-test ql-dml-test --gtest_filter *.OpenRecentlyCreatedTable

Reviewers: mihnea, bogdan, mikhail

Reviewed By: mikhail

Subscribers: ybase, bharat

Differential Revision: https://phabricator.dev.yugabyte.com/D4698
---
 src/yb/client/ql-dml-test.cc     | 35 ++++++++++++++++++++++++++++++++
 src/yb/client/table-internal.cc  |  1 +
 src/yb/master/catalog_manager.cc |  8 +++++++-
 src/yb/master/master.proto       |  2 ++
 4 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/yb/client/ql-dml-test.cc b/src/yb/client/ql-dml-test.cc
index 0fbdf4148342..ba5e633f2da4 100644
--- a/src/yb/client/ql-dml-test.cc
+++ b/src/yb/client/ql-dml-test.cc
@@ -22,6 +22,7 @@
 #include "yb/tserver/mini_tablet_server.h"
 #include "yb/tserver/tablet_server.h"
 
+#include "yb/util/backoff_waiter.h"
 #include "yb/util/curl_util.h"
 #include "yb/util/jsonreader.h"
 #include "yb/util/random.h"
@@ -1059,5 +1060,39 @@ TEST_F(QLDmlTest, TestSimultaneousReadAndWrite) {
   }
 }
 
+TEST_F(QLDmlTest, OpenRecentlyCreatedTable) {
+  constexpr int kNumIterations = 10;
+  constexpr int kNumKeys = 100;
+  const auto kOpenTimeout = 30s;
+  const auto kMaxWait = 5s;
+
+  for (int i = 0; i != kNumIterations; ++i) {
+    client::YBTableName table_name(kTableName.namespace_name(), Format("table_$0", i));
+    std::thread table_creation_thread([this, table_name] {
+      YBSchemaBuilder builder;
+      builder.AddColumn("k")->Type(INT32)->HashPrimaryKey()->NotNull();
+      builder.AddColumn("v")->Type(INT32);
+      TableHandle table;
+      ASSERT_OK(table.Create(table_name, 9, client_.get(), &builder));
+    });
+    TableHandle table;
+    BackoffWaiter waiter(std::chrono::steady_clock::now() + kOpenTimeout, kMaxWait);
+    while (!table.Open(table_name, client_.get()).ok()) {
+      ASSERT_TRUE(waiter.Wait());
+    }
+    auto session = NewSession();
+    ASSERT_OK(session->SetFlushMode(YBSession::MANUAL_FLUSH));
+    for (int k = 0; k != kNumKeys; ++k) {
+      auto op = table.NewWriteOp(QLWriteRequestPB::QL_STMT_INSERT);
+      auto* const req = op->mutable_request();
+      QLAddInt32HashValue(req, k);
+      table.AddInt32ColumnValue(req, "v", -k);
+      ASSERT_OK(session->Apply(op));
+    }
+    ASSERT_OK(session->Flush());
+    table_creation_thread.join();
+  }
+}
+
 }  // namespace client
 }  // namespace yb
diff --git a/src/yb/client/table-internal.cc b/src/yb/client/table-internal.cc
index a62819964c98..548af51409bf 100644
--- a/src/yb/client/table-internal.cc
+++ b/src/yb/client/table-internal.cc
@@ -94,6 +94,7 @@ Status YBTable::Data::Open() {
   deadline.AddDelta(client_->default_admin_operation_timeout());
 
   req.mutable_table()->set_table_id(info_.table_id);
+  req.set_require_tablets_running(true);
   Status s;
   // TODO: replace this with Async RPC-retrier based RPC in the next revision,
   // adding exponential backoff and allowing this to be used safely in a
diff --git a/src/yb/master/catalog_manager.cc b/src/yb/master/catalog_manager.cc
index 84089cdbd38d..dac177aa9985 100644
--- a/src/yb/master/catalog_manager.cc
+++ b/src/yb/master/catalog_manager.cc
@@ -5017,9 +5017,15 @@ Status CatalogManager::GetTableLocations(const GetTableLocationsRequestPB* req,
   vector<scoped_refptr<TabletInfo>> tablets_in_range;
   table->GetTabletsInRange(req, &tablets_in_range);
 
+  bool require_tablets_runnings = req->require_tablets_running();
   for (const scoped_refptr<TabletInfo>& tablet : tablets_in_range) {
-    if (!BuildLocationsForTablet(tablet, resp->add_tablet_locations()).ok()) {
+    auto status = BuildLocationsForTablet(tablet, resp->add_tablet_locations());
+    if (!status.ok()) {
       // Not running.
+      if (require_tablets_runnings) {
+        resp->mutable_tablet_locations()->Clear();
+        return SetupError(resp->mutable_error(), MasterErrorPB::TABLE_NOT_FOUND, status);
+      }
       resp->mutable_tablet_locations()->RemoveLast();
     }
   }
diff --git a/src/yb/master/master.proto b/src/yb/master/master.proto
index 10bbd3303974..282f9a5786f9 100644
--- a/src/yb/master/master.proto
+++ b/src/yb/master/master.proto
@@ -692,6 +692,8 @@ message GetTableLocationsRequestPB {
   optional bytes partition_key_end = 4;
 
   optional uint32 max_returned_locations = 5 [ default = 10 ];
+
+  optional bool require_tablets_running = 6;
 }
 
 message GetTableLocationsResponsePB {