colocation: cache db colocated in postgres (#4101)

Summary: In an effort to bring awareness of colocation to the postgres backend, cache whether the database is colocated upon connecting to it, as suggested by Neha. Add postgres global variable `MyDatabaseColocated` to cache the value; add protobufs `GetNamespaceInfoRequestPB` and `GetNamespaceInfoResponsePB` with master service RPC `GetNamespaceInfo` to retrieve whether the database is colocated from master; add various functions which are checked in the following call tree. * [ ] InitPostgres * [x] YBCPgIsDatabaseColocated * [x] PgApiImpl::IsDatabaseColocated * [x] PgSession::IsDatabaseColocated * [x] YBClient::GetNamespaceInfo * [x] YBClient::Data::SyncLeaderMasterRpc<GetNamespaceInfo...> * ... * [x] MasterServiceImpl::GetNamespaceInfo * [x] CatalogManager::GetNamespaceInfo Close: #4101 Test Plan: * Manual test for `postgres`'s `MyDatabaseColocated` variable: 1. Connect to a non-colocated database (e.g. `./bin/ysqlsh -c yugabyte`) 1. `gdb -p <pid_of_that_postgres_process>` 1. `print MyDatabaseColocated` (within GDB) should be `false` 1. `quit` (within GDB) 1. Create and connect to a colocated database (e.g. `CREATE DATABASE co COLOCATED true`, `\c co`) 1. `gdb -p <pid_of_that_postgres_process>` (should be different) 1. `print MyDatabaseColocated` (within GDB) should be `true` * `./yb_build.sh --cxx-test client_client-test --gtest_filter ClientTest.GetNamespaceInfo` Reviewers: mihnea, neha Reviewed By: neha Subscribers: yql, bogdan Differential Revision: https://phabricator.dev.yugabyte.com/D8207
yugabyte · Apr 1, 2020 · f008209 · f008209
1 parent 4b24697
commit f008209
Show file tree

Hide file tree

Showing 20 changed files with 158 additions and 9 deletions.
diff --git a/src/postgres/src/backend/access/ybc/ybcam.c b/src/postgres/src/backend/access/ybc/ybcam.c
@@ -963,7 +963,7 @@ static void ybcSetupTargets(Relation relation,
  * Begin a scan for
  *   SELECT <Targets> FROM <Relation relation> USING <Relation index>
  * NOTES:
- * - "relation" is the table beging SELECTed.
+ * - "relation" is the table being SELECTed.
  * - "index" identify the INDEX that will be used for scaning.
  * - "nkeys" and "key" identify which key columns are provided in the SELECT WHERE clause.
  *   nkeys = Number of key.

diff --git a/src/postgres/src/backend/utils/init/globals.c b/src/postgres/src/backend/utils/init/globals.c
@@ -85,6 +85,8 @@ Oid			MyDatabaseId = InvalidOid;
 
 Oid			MyDatabaseTableSpace = InvalidOid;
 
+bool		MyDatabaseColocated = false;
+
 /*
  * DatabasePath is the path (relative to DataDir) of my database's
  * primary directory, ie, its directory in the default tablespace.

diff --git a/src/postgres/src/backend/utils/init/postinit.c b/src/postgres/src/backend/utils/init/postinit.c
@@ -1030,6 +1030,15 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 	}
 	RelationCacheInitializePhase3();
 
+	/*
+	 * Also cache whather the database is colocated for optimization purposes.
+	 */
+	if (IsYugaByteEnabled() && !IsBootstrapProcessingMode())
+	{
+		HandleYBStatus(YBCPgIsDatabaseColocated(MyDatabaseId,
+												&MyDatabaseColocated));
+	}
+
 	/* set up ACL framework (so CheckMyDatabase can check permissions) */
 	initialize_acl();
 

diff --git a/src/postgres/src/backend/utils/misc/pg_yb_utils.c b/src/postgres/src/backend/utils/misc/pg_yb_utils.c
@@ -236,8 +236,8 @@ void
 HandleYBStatus(YBCStatus status)
 {
 	if (!status) {
-    return;
-  }
+		return;
+	}
 	/* Copy the message to the current memory context and free the YBCStatus. */
 	const uint32_t pg_err_code = YBCStatusPgsqlError(status);
 	char* msg_buf = DupYBStatusMessage(status, pg_err_code == ERRCODE_UNIQUE_VIOLATION);

diff --git a/src/postgres/src/include/miscadmin.h b/src/postgres/src/include/miscadmin.h
@@ -185,6 +185,8 @@ extern PGDLLIMPORT Oid MyDatabaseId;
 
 extern PGDLLIMPORT Oid MyDatabaseTableSpace;
 
+extern PGDLLIMPORT bool MyDatabaseColocated;
+
 /*
  * Date/Time Configuration
  *

diff --git a/src/yb/client/client-internal.cc b/src/yb/client/client-internal.cc
@@ -270,6 +270,7 @@ YB_CLIENT_SPECIALIZE_SIMPLE(CreateNamespace);
 YB_CLIENT_SPECIALIZE_SIMPLE(DeleteNamespace);
 YB_CLIENT_SPECIALIZE_SIMPLE(AlterNamespace);
 YB_CLIENT_SPECIALIZE_SIMPLE(ListNamespaces);
+YB_CLIENT_SPECIALIZE_SIMPLE(GetNamespaceInfo);
 YB_CLIENT_SPECIALIZE_SIMPLE(ReservePgsqlOids);
 YB_CLIENT_SPECIALIZE_SIMPLE(GetYsqlCatalogConfig);
 YB_CLIENT_SPECIALIZE_SIMPLE(CreateUDType);

diff --git a/src/yb/client/client-test.cc b/src/yb/client/client-test.cc
@@ -130,6 +130,7 @@ using base::subtle::NoBarrier_AtomicIncrement;
 using base::subtle::NoBarrier_Load;
 using base::subtle::NoBarrier_Store;
 using master::CatalogManager;
+using master::GetNamespaceInfoResponsePB;
 using master::GetTableLocationsRequestPB;
 using master::GetTableLocationsResponsePB;
 using master::TabletLocationsPB;
@@ -2210,5 +2211,35 @@ TEST_F(ClientTest, FlushTable) {
 }
 #endif  // !defined(__clang__)
 
+TEST_F(ClientTest, GetNamespaceInfo) {
+  const std::string kPgsqlKeyspaceID = "1234";
+  const std::string kPgsqlKeyspaceName = "psql" + kKeyspaceName;
+  GetNamespaceInfoResponsePB resp;
+
+  // Setup.
+  ASSERT_OK(client_->CreateNamespace(kPgsqlKeyspaceName,
+                                     YQLDatabase::YQL_DATABASE_PGSQL,
+                                     "" /* creator_role_name */,
+                                     kPgsqlKeyspaceID,
+                                     "" /* source_namespace_id */,
+                                     boost::none /* next_pg_oid */,
+                                     true /* colocated */));
+
+  // CQL non-colocated.
+  ASSERT_OK(client_->GetNamespaceInfo(
+        "" /* namespace_id */, kKeyspaceName, YQL_DATABASE_CQL, &resp));
+  ASSERT_EQ(resp.namespace_().name(), kKeyspaceName);
+  ASSERT_EQ(resp.namespace_().database_type(), YQL_DATABASE_CQL);
+  ASSERT_FALSE(resp.colocated());
+
+  // SQL colocated.
+  ASSERT_OK(client_->GetNamespaceInfo(
+        kPgsqlKeyspaceID, "" /* namespace_name */, YQL_DATABASE_PGSQL, &resp));
+  ASSERT_EQ(resp.namespace_().id(), kPgsqlKeyspaceID);
+  ASSERT_EQ(resp.namespace_().name(), kPgsqlKeyspaceName);
+  ASSERT_EQ(resp.namespace_().database_type(), YQL_DATABASE_PGSQL);
+  ASSERT_TRUE(resp.colocated());
+}
+
 }  // namespace client
 }  // namespace yb
diff --git a/src/yb/client/client.cc b/src/yb/client/client.cc
@@ -80,6 +80,8 @@ using yb::master::CreateTableRequestPB;
 using yb::master::CreateTableResponsePB;
 using yb::master::DeleteTableRequestPB;
 using yb::master::DeleteTableResponsePB;
+using yb::master::GetNamespaceInfoRequestPB;
+using yb::master::GetNamespaceInfoResponsePB;
 using yb::master::GetTableSchemaRequestPB;
 using yb::master::GetTableSchemaResponsePB;
 using yb::master::GetTableLocationsRequestPB;
@@ -556,7 +558,7 @@ Status YBClient::CreateNamespace(const std::string& namespace_name,
                                  const std::string& namespace_id,
                                  const std::string& source_namespace_id,
                                  const boost::optional<uint32_t>& next_pg_oid,
-                                 bool colocated) {
+                                 const bool colocated) {
   CreateNamespaceRequestPB req;
   CreateNamespaceResponsePB resp;
   req.set_name(namespace_name);
@@ -585,15 +587,16 @@ Status YBClient::CreateNamespaceIfNotExists(const std::string& namespace_name,
                                             const std::string& creator_role_name,
                                             const std::string& namespace_id,
                                             const std::string& source_namespace_id,
-                                            const boost::optional<uint32_t>& next_pg_oid) {
+                                            const boost::optional<uint32_t>& next_pg_oid,
+                                            const bool colocated) {
   Result<bool> namespace_exists = (!namespace_id.empty() ? NamespaceIdExists(namespace_id)
                                                          : NamespaceExists(namespace_name));
   if (VERIFY_RESULT(namespace_exists)) {
     return Status::OK();
   }
 
   return CreateNamespace(namespace_name, database_type, creator_role_name, namespace_id,
-                         source_namespace_id, next_pg_oid);
+                         source_namespace_id, next_pg_oid, colocated);
 }
 
 Status YBClient::DeleteNamespace(const std::string& namespace_name,
@@ -635,6 +638,28 @@ Result<vector<master::NamespaceIdentifierPB>> YBClient::ListNamespaces(
   return result;
 }
 
+Status YBClient::GetNamespaceInfo(const std::string& namespace_id,
+                                  const std::string& namespace_name,
+                                  const boost::optional<YQLDatabase>& database_type,
+                                  master::GetNamespaceInfoResponsePB* ret) {
+  GetNamespaceInfoRequestPB req;
+  GetNamespaceInfoResponsePB resp;
+
+  if (!namespace_id.empty()) {
+    req.mutable_namespace_()->set_id(namespace_id);
+  }
+  if (!namespace_name.empty()) {
+    req.mutable_namespace_()->set_name(namespace_name);
+  }
+  if (database_type) {
+    req.mutable_namespace_()->set_database_type(*database_type);
+  }
+
+  CALL_SYNC_LEADER_MASTER_RPC(req, resp, GetNamespaceInfo);
+  ret->Swap(&resp);
+  return Status::OK();
+}
+
 Status YBClient::ReservePgsqlOids(const std::string& namespace_id,
                                   const uint32_t next_oid, const uint32_t count,
                                   uint32_t* begin_oid, uint32_t* end_oid) {

diff --git a/src/yb/client/client.h b/src/yb/client/client.h
@@ -321,7 +321,7 @@ class YBClient {
                                  const std::string& namespace_id = "",
                                  const std::string& source_namespace_id = "",
                                  const boost::optional<uint32_t>& next_pg_oid = boost::none,
-                                 bool colocated = false);
+                                 const bool colocated = false);
 
   // It calls CreateNamespace(), but before it checks that the namespace has NOT been yet
   // created. So, it prevents error 'namespace already exists'.
@@ -334,7 +334,8 @@ class YBClient {
                                             const std::string& namespace_id = "",
                                             const std::string& source_namespace_id = "",
                                             const boost::optional<uint32_t>& next_pg_oid =
-                                            boost::none);
+                                            boost::none,
+                                            const bool colocated = false);
 
   // Delete namespace with the given name.
   CHECKED_STATUS DeleteNamespace(const std::string& namespace_name,
@@ -368,6 +369,12 @@ class YBClient {
   Result<vector<master::NamespaceIdentifierPB>> ListNamespaces(
       const boost::optional<YQLDatabase>& database_type);
 
+  // Get namespace information.
+  CHECKED_STATUS GetNamespaceInfo(const std::string& namespace_id,
+                                  const std::string& namespace_name,
+                                  const boost::optional<YQLDatabase>& database_type,
+                                  master::GetNamespaceInfoResponsePB* ret);
+
   // Check if the namespace given by 'namespace_name' or 'namespace_id' exists.
   // Result value is set only on success.
   Result<bool> NamespaceExists(const std::string& namespace_name,

diff --git a/src/yb/master/catalog_manager.cc b/src/yb/master/catalog_manager.cc
@@ -4619,7 +4619,6 @@ Status CatalogManager::AlterNamespace(const AlterNamespaceRequestPB* req,
 
 Status CatalogManager::ListNamespaces(const ListNamespacesRequestPB* req,
                                       ListNamespacesResponsePB* resp) {
-
   RETURN_NOT_OK(CheckOnline());
 
   SharedLock<LockType> l(lock_);
@@ -4640,6 +4639,27 @@ Status CatalogManager::ListNamespaces(const ListNamespacesRequestPB* req,
   return Status::OK();
 }
 
+Status CatalogManager::GetNamespaceInfo(const GetNamespaceInfoRequestPB* req,
+                                        GetNamespaceInfoResponsePB* resp,
+                                        rpc::RpcContext* rpc) {
+  LOG(INFO) << __func__ << " from " << RequestorString(rpc) << ": " << req->ShortDebugString();
+  RETURN_NOT_OK(CheckOnline());
+
+  scoped_refptr<NamespaceInfo> ns;
+
+  // Look up the namespace and verify if it exists.
+  if (req->has_namespace_()) {
+    TRACE("Looking up namespace");
+    RETURN_NAMESPACE_NOT_FOUND(FindNamespace(req->namespace_(), &ns), resp);
+  }
+
+  resp->mutable_namespace_()->set_id(ns->id());
+  resp->mutable_namespace_()->set_name(ns->name());
+  resp->mutable_namespace_()->set_database_type(ns->database_type());
+  resp->set_colocated(ns->colocated());
+  return Status::OK();
+}
+
 Status CatalogManager::RedisConfigSet(
     const RedisConfigSetRequestPB* req, RedisConfigSetResponsePB* resp, rpc::RpcContext* rpc) {
   DCHECK(req->has_keyword());

diff --git a/src/yb/master/catalog_manager.h b/src/yb/master/catalog_manager.h
@@ -329,6 +329,11 @@ class CatalogManager : public tserver::TabletPeerLookupIf {
   CHECKED_STATUS ListNamespaces(const ListNamespacesRequestPB* req,
                                 ListNamespacesResponsePB* resp);
 
+  // Get information about a namespace.
+  CHECKED_STATUS GetNamespaceInfo(const GetNamespaceInfoRequestPB* req,
+                                  GetNamespaceInfoResponsePB* resp,
+                                  rpc::RpcContext* rpc);
+
   // Set Redis Config
   CHECKED_STATUS RedisConfigSet(const RedisConfigSetRequestPB* req,
                                 RedisConfigSetResponsePB* resp,

diff --git a/src/yb/master/master.proto b/src/yb/master/master.proto
@@ -1323,6 +1323,18 @@ message ListNamespacesResponsePB {
   repeated NamespaceIdentifierPB namespaces = 2;
 }
 
+message GetNamespaceInfoRequestPB {
+  optional NamespaceIdentifierPB namespace = 1;
+}
+
+message GetNamespaceInfoResponsePB {
+  // The error, if an error occurred with this request.
+  optional MasterErrorPB error = 1;
+
+  optional NamespaceIdentifierPB namespace = 2;
+  optional bool colocated = 3;
+}
+
 // ============================================================================
 //  Authentication and Authorization
 // ============================================================================
@@ -1712,6 +1724,7 @@ service MasterService {
   rpc DeleteNamespace(DeleteNamespaceRequestPB) returns (DeleteNamespaceResponsePB);
   rpc AlterNamespace(AlterNamespaceRequestPB) returns (AlterNamespaceResponsePB);
   rpc ListNamespaces(ListNamespacesRequestPB) returns (ListNamespacesResponsePB);
+  rpc GetNamespaceInfo(GetNamespaceInfoRequestPB) returns (GetNamespaceInfoResponsePB);
 
   // For Postgres:
   rpc ReservePgsqlOids(ReservePgsqlOidsRequestPB) returns (ReservePgsqlOidsResponsePB);

diff --git a/src/yb/master/master_service.cc b/src/yb/master/master_service.cc
@@ -351,6 +351,12 @@ void MasterServiceImpl::ListNamespaces(const ListNamespacesRequestPB* req,
   HandleIn(req, resp, &rpc, &CatalogManager::ListNamespaces);
 }
 
+void MasterServiceImpl::GetNamespaceInfo(const GetNamespaceInfoRequestPB* req,
+                                         GetNamespaceInfoResponsePB* resp,
+                                         RpcContext rpc) {
+  HandleIn(req, resp, &rpc, &CatalogManager::GetNamespaceInfo);
+}
+
 void MasterServiceImpl::ReservePgsqlOids(const ReservePgsqlOidsRequestPB* req,
                                          ReservePgsqlOidsResponsePB* resp,
                                          rpc::RpcContext rpc) {

diff --git a/src/yb/master/master_service.h b/src/yb/master/master_service.h
@@ -106,6 +106,9 @@ class MasterServiceImpl : public MasterServiceIf,
   void ListNamespaces(const ListNamespacesRequestPB* req,
                       ListNamespacesResponsePB* resp,
                       rpc::RpcContext rpc) override;
+  void GetNamespaceInfo(const GetNamespaceInfoRequestPB* req,
+                        GetNamespaceInfoResponsePB* resp,
+                        rpc::RpcContext rpc) override;
 
   void ReservePgsqlOids(const ReservePgsqlOidsRequestPB* req,
                         ReservePgsqlOidsResponsePB* resp,

diff --git a/src/yb/yql/pggate/pg_session.cc b/src/yb/yql/pggate/pg_session.cc
@@ -63,6 +63,7 @@ using client::YBTable;
 using client::YBTableName;
 using client::YBTableType;
 
+using yb::master::GetNamespaceInfoResponsePB;
 using yb::master::IsInitDbDoneRequestPB;
 using yb::master::IsInitDbDoneResponsePB;
 using yb::master::MasterServiceProxy;
@@ -414,6 +415,14 @@ Status PgSession::ConnectDatabase(const string& database_name) {
   return Status::OK();
 }
 
+Status PgSession::IsDatabaseColocated(const PgOid database_oid, bool *colocated) {
+  GetNamespaceInfoResponsePB resp;
+  RETURN_NOT_OK(client_->GetNamespaceInfo(
+      GetPgsqlNamespaceId(database_oid), "" /* namespace_name */, YQL_DATABASE_PGSQL, &resp));
+  *colocated = resp.colocated();
+  return Status::OK();
+}
+
 //--------------------------------------------------------------------------------------------------
 
 Status PgSession::CreateDatabase(const string& database_name,

diff --git a/src/yb/yql/pggate/pg_session.h b/src/yb/yql/pggate/pg_session.h
@@ -114,6 +114,8 @@ class PgSession : public RefCountedThreadSafe<PgSession> {
 
   CHECKED_STATUS ConnectDatabase(const std::string& database_name);
 
+  CHECKED_STATUS IsDatabaseColocated(const PgOid database_oid, bool *colocated);
+
   //------------------------------------------------------------------------------------------------
   // Operations on Database Objects.
   //------------------------------------------------------------------------------------------------

diff --git a/src/yb/yql/pggate/pggate.cc b/src/yb/yql/pggate/pggate.cc
@@ -270,6 +270,10 @@ Status PgApiImpl::ConnectDatabase(const char *database_name) {
   return pg_session_->ConnectDatabase(database_name);
 }
 
+Status PgApiImpl::IsDatabaseColocated(const PgOid database_oid, bool *colocated) {
+  return pg_session_->IsDatabaseColocated(database_oid, colocated);
+}
+
 Status PgApiImpl::NewCreateDatabase(const char *database_name,
                                     const PgOid database_oid,
                                     const PgOid source_database_oid,

diff --git a/src/yb/yql/pggate/pggate.h b/src/yb/yql/pggate/pggate.h
@@ -128,6 +128,9 @@ class PgApiImpl {
   // Connect database. Switch the connected database to the given "database_name".
   CHECKED_STATUS ConnectDatabase(const char *database_name);
 
+  // Determine whether the given database is colocated.
+  CHECKED_STATUS IsDatabaseColocated(const PgOid database_oid, bool *colocated);
+
   // Create database.
   CHECKED_STATUS NewCreateDatabase(const char *database_name,
                                    PgOid database_oid,

diff --git a/src/yb/yql/pggate/ybc_pggate.cc b/src/yb/yql/pggate/ybc_pggate.cc
@@ -128,6 +128,10 @@ YBCStatus YBCPgConnectDatabase(const char *database_name) {
   return ToYBCStatus(pgapi->ConnectDatabase(database_name));
 }
 
+YBCStatus YBCPgIsDatabaseColocated(const YBCPgOid database_oid, bool *colocated) {
+  return ToYBCStatus(pgapi->IsDatabaseColocated(database_oid, colocated));
+}
+
 YBCStatus YBCPgNewCreateDatabase(const char *database_name,
                                  const YBCPgOid database_oid,
                                  const YBCPgOid source_database_oid,

diff --git a/src/yb/yql/pggate/ybc_pggate.h b/src/yb/yql/pggate/ybc_pggate.h
@@ -63,6 +63,9 @@ YBCStatus YBCGetSharedCatalogVersion(uint64_t* catalog_version);
 // Connect database. Switch the connected database to the given "database_name".
 YBCStatus YBCPgConnectDatabase(const char *database_name);
 
+// Get whether the given database is colocated.
+YBCStatus YBCPgIsDatabaseColocated(const YBCPgOid database_oid, bool *colocated);
+
 YBCStatus YBCInsertSequenceTuple(int64_t db_oid,
                                  int64_t seq_oid,
                                  uint64_t ysql_catalog_version,
-Original file line number
+Diff line change
@@ Expand Up / @@ -185,6 +185,8 @@ extern PGDLLIMPORT Oid MyDatabaseId; @@
     extern PGDLLIMPORT Oid MyDatabaseTableSpace;
+    extern PGDLLIMPORT bool MyDatabaseColocated;
     /*
      * Date/Time Configuration
      *
@@ Expand Down @@