[#3035][YSQL] Table colocation drop flow

Summary:
Currently, table colocation is supported for the create flow: `CREATE
DATABASE`, `CREATE TABLE`, and `CREATE INDEX` including opt-out.  Work
on supporting the drop flow.

* Add support for `DROP DATABASE` on colocated databases.
* Add support for `DROP TABLE` on colocated tables with the caveat that
  data for the table does not get deleted.  In the meantime, to work
  around this issue, execute `DELETE FROM <table>` before `DROP TABLE
  <table>`.
* Omit support for `TRUNCATE`.  For now, executing `TRUNCATE` on a
  colocated table will delete **all data** (including those of other
  tables) from its parent tablet.
* Update tests.

Test Plan:
* `./yb_build.sh --cxx-test pg_libpq-test --gtest_filter
  PgLibPqTest.TableColocation`
* `./yb_build.sh --java-test org.yb.pgsql.TestPgRegressBetaFeatures`

Reviewers: neha

Reviewed By: neha

Subscribers: yql, bogdan

Differential Revision: https://phabricator.dev.yugabyte.com/D7737

src/postgres/src/test/regress/expected/yb_feature_colocation.out

@@ -11,7 +11,7 @@ CREATE TABLE foo2 (a INT, b INT, PRIMARY KEY (a ASC));
 -- opt out of using colocated tablet
 CREATE TABLE foo3 (a INT) WITH (colocated = false);
 -- multi column primary key table
-CREATE TABLE foo4(a INT, b INT, PRIMARY KEY (a ASC, b DESC));
+CREATE TABLE foo4 (a INT, b INT, PRIMARY KEY (a ASC, b DESC));
 CREATE TABLE foo5 (a INT, PRIMARY KEY (a ASC)) WITH (colocated = true);
 INSERT INTO foo1 (a) VALUES (0), (1), (2);
 INSERT INTO foo1 (a, b) VALUES (0, '0'); -- fail
@@ -191,3 +191,98 @@ SELECT * FROM bar3;
  3 | 3
 (4 rows)
 
+-- non-colocated table with index
+CREATE TABLE bar4 (a INT, b INT, PRIMARY KEY (a ASC)) WITH (colocated = false);
+CREATE INDEX ON bar4 (a);
+INSERT INTO bar4 (a, b) VALUES (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5);
+EXPLAIN SELECT * FROM bar4 WHERE a = 1;
+                              QUERY PLAN
+-----------------------------------------------------------------------
+ Index Scan using bar4_a_idx on bar4  (cost=0.00..4.11 rows=1 width=8)
+   Index Cond: (a = 1)
+(2 rows)
+
+SELECT * FROM bar4 WHERE a = 1;
+ a | b
+---+---
+ 1 | 1
+(1 row)
+
+UPDATE bar4 SET b = b + 1 WHERE a > 3;
+SELECT * FROM bar4;
+ a | b
+---+---
+ 0 | 0
+ 1 | 1
+ 2 | 2
+ 3 | 3
+ 4 | 5
+ 5 | 6
+(6 rows)
+
+DELETE FROM bar4 WHERE a > 3;
+SELECT * FROM bar4;
+ a | b
+---+---
+ 0 | 0
+ 1 | 1
+ 2 | 2
+ 3 | 3
+(4 rows)
+
+-- drop table
+DROP TABLE foo1;
+SELECT * FROM foo1;
+ERROR:  relation "foo1" does not exist
+LINE 1: SELECT * FROM foo1;
+                      ^
+DROP TABLE foo3;
+SELECT * FROM foo3;
+ERROR:  relation "foo3" does not exist
+LINE 1: SELECT * FROM foo3;
+                      ^
+DROP TABLE bar1;
+SELECT * FROM bar1;
+ERROR:  relation "bar1" does not exist
+LINE 1: SELECT * FROM bar1;
+                      ^
+-- drop index
+DROP INDEX bar4_a_idx;
+EXPLAIN SELECT * FROM bar4 WHERE a = 1;
+                              QUERY PLAN
+----------------------------------------------------------------------
+ Index Scan using bar4_pkey on bar4  (cost=0.00..4.11 rows=1 width=8)
+   Index Cond: (a = 1)
+(2 rows)
+
+\dt
+        List of relations
+ Schema | Name | Type  |  Owner
+--------+------+-------+----------
+ public | bar2 | table | yugabyte
+ public | bar3 | table | yugabyte
+ public | bar4 | table | yugabyte
+ public | foo2 | table | yugabyte
+ public | foo4 | table | yugabyte
+ public | foo5 | table | yugabyte
+(6 rows)
+
+\di
+               List of relations
+ Schema |    Name    | Type  |  Owner   | Table
+--------+------------+-------+----------+-------
+ public | bar2_a_idx | index | yugabyte | bar2
+ public | bar2_pkey  | index | yugabyte | bar2
+ public | bar3_a_idx | index | yugabyte | bar3
+ public | bar3_pkey  | index | yugabyte | bar3
+ public | bar4_pkey  | index | yugabyte | bar4
+ public | foo2_pkey  | index | yugabyte | foo2
+ public | foo4_pkey  | index | yugabyte | foo4
+ public | foo5_pkey  | index | yugabyte | foo5
+(8 rows)
+
+-- drop database
+\c yugabyte
+DROP DATABASE colocation_test;
+\c colocation_test
+\connect: FATAL:  database "colocation_test" does not exist

src/postgres/src/test/regress/sql/yb_feature_colocation.sql

@@ -12,7 +12,7 @@ CREATE TABLE foo2 (a INT, b INT, PRIMARY KEY (a ASC));
 -- opt out of using colocated tablet
 CREATE TABLE foo3 (a INT) WITH (colocated = false);
 -- multi column primary key table
-CREATE TABLE foo4(a INT, b INT, PRIMARY KEY (a ASC, b DESC));
+CREATE TABLE foo4 (a INT, b INT, PRIMARY KEY (a ASC, b DESC));
 CREATE TABLE foo5 (a INT, PRIMARY KEY (a ASC)) WITH (colocated = true);
 
 INSERT INTO foo1 (a) VALUES (0), (1), (2);
@@ -62,3 +62,34 @@ UPDATE bar3 SET b = b + 1 WHERE a > 3;
 SELECT * FROM bar3;
 DELETE FROM bar3 WHERE a > 3;
 SELECT * FROM bar3;
+
+-- non-colocated table with index
+CREATE TABLE bar4 (a INT, b INT, PRIMARY KEY (a ASC)) WITH (colocated = false);
+CREATE INDEX ON bar4 (a);
+INSERT INTO bar4 (a, b) VALUES (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5);
+EXPLAIN SELECT * FROM bar4 WHERE a = 1;
+SELECT * FROM bar4 WHERE a = 1;
+UPDATE bar4 SET b = b + 1 WHERE a > 3;
+SELECT * FROM bar4;
+DELETE FROM bar4 WHERE a > 3;
+SELECT * FROM bar4;
+
+-- drop table
+DROP TABLE foo1;
+SELECT * FROM foo1;
+DROP TABLE foo3;
+SELECT * FROM foo3;
+DROP TABLE bar1;
+SELECT * FROM bar1;
+
+-- drop index
+DROP INDEX bar4_a_idx;
+EXPLAIN SELECT * FROM bar4 WHERE a = 1;
+
+\dt
+\di
+
+-- drop database
+\c yugabyte
+DROP DATABASE colocation_test;
+\c colocation_test

src/yb/integration-tests/create-table-itest.cc

@@ -346,7 +346,7 @@ TEST_F(CreateTableITest, TestNoAllocBlacklist) {
   ASSERT_EQ(inspect_->ListTabletsOnTS(1).size(), 0);
 }
 
-TEST_F(CreateTableITest, TabletColocationRemoteBootstrapTest) {
+TEST_F(CreateTableITest, TableColocationRemoteBootstrapTest) {
   const int kNumReplicas = 3;
   vector<string> ts_flags;
   vector<string> master_flags;

src/yb/master/async_rpc_tasks.cc

@@ -990,5 +990,42 @@ bool AsyncAddTableToTablet::SendRequest(int attempt) {
   return true;
 }
 
+// ============================================================================
+//  Class AsyncRemoveTableFromTablet.
+// ============================================================================
+AsyncRemoveTableFromTablet::AsyncRemoveTableFromTablet(
+    Master* master, ThreadPool* callback_pool, const scoped_refptr<TabletInfo>& tablet,
+    const scoped_refptr<TableInfo>& table)
+    : RetryingTSRpcTask(
+          master, callback_pool, gscoped_ptr<TSPicker>(new PickLeaderReplica(tablet)), table.get()),
+      table_(table),
+      tablet_(tablet),
+      tablet_id_(tablet->tablet_id()) {
+  req_.set_tablet_id(tablet->id());
+  req_.set_remove_table_id(table->id());
+}
+
+string AsyncRemoveTableFromTablet::description() const {
+  return Substitute("RemoveTableFromTablet RPC ($0) ($1)", table_->ToString(), tablet_->ToString());
+}
+
+void AsyncRemoveTableFromTablet::HandleResponse(int attempt) {
+  if (!rpc_.status().ok()) {
+    AbortTask();
+    LOG(WARNING) << Substitute(
+        "Got error when removing table $0 from tablet $1, attempt $2 and error $3",
+        table_->ToString(), tablet_->ToString(), attempt, rpc_.status().ToString());
+    return;
+  }
+  TransitionToTerminalState(MonitoredTaskState::kRunning, MonitoredTaskState::kComplete);
+}
+
+bool AsyncRemoveTableFromTablet::SendRequest(int attempt) {
+  ts_admin_proxy_->RemoveTableFromTabletAsync(req_, &resp_, &rpc_, BindRpcCallback());
+  VLOG(1) << "Send RemoveTableFromTablet request (attempt " << attempt << "):\n"
+          << req_.DebugString();
+  return true;
+}
+
 }  // namespace master
 }  // namespace yb

src/yb/master/async_rpc_tasks.h

@@ -549,6 +549,33 @@ class AsyncAddTableToTablet : public RetryingTSRpcTask {
   tserver::AddTableToTabletResponsePB resp_;
 };
 
+// Task to remove a table from a tablet. Catalog Manager uses this task to send the request to the
+// tserver admin service.
+class AsyncRemoveTableFromTablet : public RetryingTSRpcTask {
+ public:
+  AsyncRemoveTableFromTablet(
+      Master* master, ThreadPool* callback_pool, const scoped_refptr<TabletInfo>& tablet,
+      const scoped_refptr<TableInfo>& table);
+
+  Type type() const override { return ASYNC_REMOVE_TABLE_FROM_TABLET; }
+
+  std::string type_name() const override { return "Remove Table from Tablet"; }
+
+  std::string description() const override;
+
+ private:
+  TabletId tablet_id() const override { return tablet_id_; }
+
+  bool SendRequest(int attempt) override;
+  void HandleResponse(int attempt) override;
+
+  const scoped_refptr<TableInfo> table_;
+  const scoped_refptr<TabletInfo> tablet_;
+  const TabletId tablet_id_;
+  tserver::RemoveTableFromTabletRequestPB req_;
+  tserver::RemoveTableFromTabletResponsePB resp_;
+};
+
 } // namespace master
 } // namespace yb
 

src/yb/master/catalog_entity_info.cc

@@ -492,6 +492,16 @@ void TableInfo::GetAllTablets(TabletInfos *ret) const {
   }
 }
 
+TabletInfoPtr TableInfo::GetColocatedTablet() const {
+  shared_lock<decltype(lock_)> l(lock_);
+  if (colocated()) {
+    for (const TableInfo::TabletInfoMap::value_type& e : tablet_map_) {
+      return make_scoped_refptr(e.second);
+    }
+  }
+  return nullptr;
+}
+
 IndexInfo TableInfo::GetIndexInfo(const TableId& index_id) const {
   auto l = LockForRead();
   for (const auto& index_info_pb : l->data().pb.indexes()) {

src/yb/master/catalog_entity_info.h

@@ -361,12 +361,16 @@ class TableInfo : public RefCountedThreadSafe<TableInfo>,
   // This only returns tablets which are in RUNNING state.
   void GetTabletsInRange(const GetTableLocationsRequestPB* req, TabletInfos *ret) const;
 
+  // Get all tablets of the table.
   void GetAllTablets(TabletInfos *ret) const;
 
+  // Get the tablet of the table.  The table must be colocated.
+  TabletInfoPtr GetColocatedTablet() const;
+
   // Get info of the specified index.
   IndexInfo GetIndexInfo(const TableId& index_id) const;
 
-  // Returns true if the table creation is in-progress.
+  // Returns true if all tablets of the table are deleted.
   bool AreAllTabletsDeleted() const;
 
   // Returns true if the table creation is in-progress.

src/yb/master/catalog_manager.cc

@@ -2905,9 +2905,16 @@ Status CatalogManager::DeleteTable(const DeleteTableRequestPB* req,
     SleepFor(MonoDelta::FromMilliseconds(FLAGS_catalog_manager_inject_latency_in_delete_table_ms));
   }
 
-  for (int i = 0; i < tables.size(); i++) {
+  for (const scoped_refptr<TableInfo> &table : tables) {
     // Send a DeleteTablet() request to each tablet replica in the table.
-    DeleteTabletsAndSendRequests(tables[i]);
+    DeleteTabletsAndSendRequests(table);
+    // Send a RemoveTableFromTablet() request to each colocated parent tablet replica in the table.
+    if (IsColocatedUserTable(*table)) {
+      auto call = std::make_shared<AsyncRemoveTableFromTablet>(
+          master_, worker_pool_.get(), table->GetColocatedTablet(), table);
+      table->AddTask(call);
+      WARN_NOT_OK(call->Run(), "Failed to send RemoveTableFromTablet request");
+    }
   }
 
   // If there are any permissions granted on this table find them and delete them. This is necessary
@@ -3068,7 +3075,9 @@ void CatalogManager::CleanUpDeletedTables() {
           // gotten to point 3, which would add further tasks for the deletes.
           //
           // However, HasTasks is cheaper than AreAllTabletsDeleted...
-          if (table->AreAllTabletsDeleted() || IsSystemTableUnlocked(*table)) {
+          if (table->AreAllTabletsDeleted() ||
+              IsSystemTableUnlocked(*table) ||
+              IsColocatedUserTable(*table)) {
             tables_to_delete.push_back(table);
             // TODO(bogdan): uncomment this once we also untangle catalog loader logic.
             // Since we have lock_, this table cannot be in the map AND be DELETED.
@@ -3716,6 +3725,10 @@ bool CatalogManager::IsColocatedParentTable(const TableInfo& table) const {
   return table.id().find(kColocatedParentTableIdSuffix) != std::string::npos;
 }
 
+bool CatalogManager::IsColocatedUserTable(const TableInfo& table) const {
+  return table.colocated() && !IsColocatedParentTable(table);
+}
+
 bool CatalogManager::IsSequencesSystemTable(const TableInfo& table) const {
   if (table.GetTableType() == PGSQL_TABLE_TYPE) {
     // This case commonly occurs during unit testing. Avoid unnecessary assert within Get().
@@ -5275,6 +5288,10 @@ void CatalogManager::DeleteTabletsAndSendRequests(const scoped_refptr<TableInfo>
       return;
     }
   }
+  // Do not delete the tablet of a colocated table.
+  if (IsColocatedUserTable(*table)) {
+    return;
+  }
 
   vector<scoped_refptr<TabletInfo>> tablets;
   table->GetAllTablets(&tablets);
@@ -5289,6 +5306,10 @@ void CatalogManager::DeleteTabletsAndSendRequests(const scoped_refptr<TableInfo>
     CHECK_OK(sys_catalog_->UpdateItem(tablet.get(), leader_ready_term_));
     tablet_lock->Commit();
   }
+  if (IsColocatedParentTable(*table)) {
+    SharedLock<LockType> catalog_lock(lock_);
+    colocated_tablet_ids_map_.erase(table->namespace_id());
+  }
 }
 
 void CatalogManager::SendDeleteTabletRequest(

src/yb/master/catalog_manager.h

@@ -438,6 +438,9 @@ class CatalogManager : public tserver::TabletPeerLookupIf {
   // Is the table a table created for colocated database?
   bool IsColocatedParentTable(const TableInfo& table) const;
 
+  // Is the table a table created in a colocated database?
+  bool IsColocatedUserTable(const TableInfo& table) const;
+
   // Is the table created by user?
   // Note that table can be regular table or index in this case.
   bool IsUserCreatedTable(const TableInfo& table) const;
@@ -917,8 +920,8 @@ class CatalogManager : public tserver::TabletPeerLookupIf {
   // Request tablet servers to delete all replicas of the tablet.
   void DeleteTabletReplicas(const TabletInfo* tablet, const std::string& msg);
 
-  // Marks each of the tablets in the given table as deleted and triggers requests
-  // to the tablet servers to delete them.
+  // Marks each of the tablets in the given table as deleted and triggers requests to the tablet
+  // servers to delete them.  The table parameter is expected to be given "write locked".
   void DeleteTabletsAndSendRequests(const scoped_refptr<TableInfo>& table);
 
   // Send the "delete tablet request" to the specified TS/tablet.

src/yb/server/monitored_task.h

@@ -74,6 +74,7 @@ class MonitoredTask : public std::enable_shared_from_this<MonitoredTask> {
     ASYNC_COPARTITION_TABLE,
     ASYNC_FLUSH_TABLETS,
     ASYNC_ADD_TABLE_TO_TABLET,
+    ASYNC_REMOVE_TABLE_FROM_TABLET,
   };
 
   virtual Type type() const = 0;