pg_wrapper.cc

// Copyright (c) YugaByte, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
// in compliance with the License.  You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied.  See the License for the specific language governing permissions and limitations
// under the License.

#include "yb/yql/pgwrapper/pg_wrapper.h"

#include <signal.h>

#include <vector>
#include <string>
#include <random>
#include <fstream>

#include <gflags/gflags.h>
#include <boost/algorithm/string.hpp>

#include "yb/util/errno.h"
#include "yb/util/flag_tags.h"
#include "yb/util/logging.h"
#include "yb/util/subprocess.h"
#include "yb/util/env_util.h"
#include "yb/util/net/net_util.h"
#include "yb/util/path_util.h"
#include "yb/util/scope_exit.h"

DEFINE_string(pg_proxy_bind_address, "", "Address for the PostgreSQL proxy to bind to");
DEFINE_bool(pg_transactions_enabled, true,
            "True to enable transactions in YugaByte PostgreSQL API. This should eventually "
            "be set to true by default.");
DEFINE_bool(pg_verbose_error_log, false,
            "True to enable verbose logging of errors in PostgreSQL server");
DEFINE_int32(pgsql_proxy_webserver_port, 13000, "Webserver port for PGSQL");
DECLARE_string(metric_node_name);
TAG_FLAG(pg_transactions_enabled, advanced);
TAG_FLAG(pg_transactions_enabled, hidden);

DEFINE_bool(pg_stat_statements_enabled, true,
            "True to enable statement stats in PostgreSQL server");
TAG_FLAG(pg_stat_statements_enabled, advanced);
TAG_FLAG(pg_stat_statements_enabled, hidden);

// Top-level postgres configuration flags.
DEFINE_bool(ysql_enable_auth, false,
              "True to enforce password authentication for all connections");
DEFINE_string(ysql_timezone, "",
              "Overrides the default ysql timezone for displaying and interpreting timestamps");
DEFINE_string(ysql_datestyle, "",
              "Overrides the default ysql display format for date and time values");
DEFINE_int32(ysql_max_connections, 0,
              "Overrides the maximum number of concurrent ysql connections");
DEFINE_string(ysql_default_transaction_isolation, "",
              "Overrides the default ysql transaction isolation level");
DEFINE_string(ysql_log_statement, "",
              "Sets which types of ysql statements should be logged");
DEFINE_string(ysql_log_min_messages, "",
              "Sets the lowest ysql message level to log");

// Catch-all postgres configuration flags.
DEFINE_string(ysql_pg_conf, "",
              "Comma separated list of postgres setting assignments");
DEFINE_string(ysql_hba_conf, "",
              "Comma separated list of postgres hba rules (in order)");

using std::vector;
using std::string;

using namespace std::literals;  // NOLINT

namespace yb {
namespace pgwrapper {

namespace {

Status WriteConfigFile(const string& path, const vector<string>& lines) {
  std::ofstream conf_file;
  conf_file.open(path, std::ios_base::out | std::ios_base::trunc);
  if (!conf_file) {
    return STATUS_FORMAT(
        IOError,
        "Failed to write ysql config file '%s': errno=$0: $1",
        path,
        errno,
        ErrnoToString(errno));
  }

  conf_file << "# This is an autogenerated file, do not edit manually!" << std::endl;
  for (const auto& line : lines) {
    conf_file << line << std::endl;
  }

  conf_file.close();

  return Status::OK();
}

void ReadCSVConfigValues(const string& csv, vector<string>* lines) {
  vector<string> csv_lines;
  boost::split(csv_lines, csv, boost::is_any_of(","));
  lines->insert(lines->end(), csv_lines.begin(), csv_lines.end());
}

Result<string> WritePostgresConfig(const PgProcessConf& conf) {
  // First add default configuration created by local initdb.
  string default_conf_path = JoinPathSegments(conf.data_dir, "postgresql.conf");
  std::ifstream conf_file;
  conf_file.open(default_conf_path, std::ios_base::in);
  if (!conf_file) {
    return STATUS_FORMAT(
        IOError,
        "Failed to read default postgres configuration '%s': errno=$0: $1",
        default_conf_path,
        errno,
        ErrnoToString(errno));
  }

  vector<string> lines;
  string line;
  while (std::getline(conf_file, line)) {
    lines.push_back(line);
  }

  if (!FLAGS_ysql_pg_conf.empty()) {
    ReadCSVConfigValues(FLAGS_ysql_pg_conf, &lines);
  }

  if (conf.enable_tls) {
    lines.push_back("ssl=on");
    lines.push_back(Format("ssl_cert_file='$0/node.$1.crt'",
                           conf.certs_for_client_dir,
                           conf.cert_base_name));
    lines.push_back(Format("ssl_key_file='$0/node.$1.key'",
                           conf.certs_for_client_dir,
                           conf.cert_base_name));
    lines.push_back(Format("ssl_ca_file='$0/ca.crt'", conf.certs_for_client_dir));
  }

  if (!FLAGS_ysql_timezone.empty()) {
    lines.push_back("timezone=" + FLAGS_ysql_timezone);
  }

  if (!FLAGS_ysql_datestyle.empty()) {
    lines.push_back("datestyle=" + FLAGS_ysql_datestyle);
  }

  if (FLAGS_ysql_max_connections > 0) {
    lines.push_back("max_connections=" + std::to_string(FLAGS_ysql_max_connections));
  }

  if (!FLAGS_ysql_default_transaction_isolation.empty()) {
    lines.push_back("default_transaction_isolation=" + FLAGS_ysql_default_transaction_isolation);
  }

  if (!FLAGS_ysql_log_statement.empty()) {
    lines.push_back("log_statement=" + FLAGS_ysql_log_statement);
  }

  if (!FLAGS_ysql_log_min_messages.empty()) {
    lines.push_back("log_min_messages=" + FLAGS_ysql_log_min_messages);
  }

  string conf_path = JoinPathSegments(conf.data_dir, "ysql_pg.conf");
  RETURN_NOT_OK(WriteConfigFile(conf_path, lines));
  return "config_file=" + conf_path;
}

Result<string> WritePgHbaConfig(const PgProcessConf& conf) {
  vector<string> lines;

  if (FLAGS_ysql_enable_auth || conf.enable_tls) {
    const auto host_type =  conf.enable_tls ? "hostssl" : "host";
    const auto auth_method = FLAGS_ysql_enable_auth ? (conf.enable_tls ? "md5 clientcert=1" : "md5")
                                                    : "cert";

    for (const auto addr : {"0.0.0.0/0", "::0/0"}) {
      lines.push_back(Format("$0 all all $1 $2", host_type, addr, auth_method));
    }
  }

  if (!FLAGS_ysql_hba_conf.empty()) {
    ReadCSVConfigValues(FLAGS_ysql_hba_conf, &lines);
  }

  // Enforce a default hba configuration, so users don't lock themselves out.
  if (lines.empty()) {
    lines.push_back("host all all 0.0.0.0/0 trust");
    lines.push_back("host all all ::0/0 trust");
  }

  string conf_path = JoinPathSegments(conf.data_dir, "ysql_hba.conf");
  RETURN_NOT_OK(WriteConfigFile(conf_path, lines));
  return "hba_file=" + conf_path;
}

Result<vector<string>> WritePgConfigFiles(const PgProcessConf& conf) {
  vector<string> args;
  args.push_back("-c");
  args.push_back(VERIFY_RESULT_PREPEND(WritePostgresConfig(conf),
      "Failed to write ysql pg configuration: "));
  args.push_back("-c");
  args.push_back(VERIFY_RESULT_PREPEND(WritePgHbaConfig(conf),
      "Failed to write ysql hba configuration: "));
  return args;
}

}  // namespace

string GetPostgresInstallRoot() {
  return JoinPathSegments(yb::env_util::GetRootDir("postgres"), "postgres");
}

Result<PgProcessConf> PgProcessConf::CreateValidateAndRunInitDb(
    const std::string& bind_addresses,
    const std::string& data_dir,
    const int tserver_shm_fd) {
  PgProcessConf conf;
  if (!bind_addresses.empty()) {
    auto pg_host_port = VERIFY_RESULT(HostPort::FromString(
        bind_addresses, PgProcessConf::kDefaultPort));
    conf.listen_addresses = pg_host_port.host();
    conf.pg_port = pg_host_port.port();
  }
  conf.data_dir = data_dir;
  conf.tserver_shm_fd = tserver_shm_fd;
  PgWrapper pg_wrapper(conf);
  RETURN_NOT_OK(pg_wrapper.PreflightCheck());
  RETURN_NOT_OK(pg_wrapper.InitDbLocalOnlyIfNeeded());
  return conf;
}

// ------------------------------------------------------------------------------------------------
// PgWrapper: managing one instance of a PostgreSQL child process
// ------------------------------------------------------------------------------------------------

PgWrapper::PgWrapper(PgProcessConf conf)
    : conf_(std::move(conf)) {
}

Status PgWrapper::PreflightCheck() {
  RETURN_NOT_OK(CheckExecutableValid(GetPostgresExecutablePath()));
  RETURN_NOT_OK(CheckExecutableValid(GetInitDbExecutablePath()));
  return Status::OK();
}

Status PgWrapper::Start() {
  auto postgres_executable = GetPostgresExecutablePath();
  RETURN_NOT_OK(CheckExecutableValid(postgres_executable));
  vector<string> argv {
    postgres_executable,
    "-D", conf_.data_dir,
    "-p", std::to_string(conf_.pg_port),
    "-h", conf_.listen_addresses,
    // Disable listening on a UNIX domain socket
    "-k", ""
  };

  bool log_to_file = !FLAGS_logtostderr && !FLAGS_log_dir.empty() && !conf_.force_disable_log_file;
  VLOG(1) << "Deciding whether the child postgres process should to file: "
          << EXPR_VALUE_FOR_LOG(FLAGS_logtostderr) << ", "
          << EXPR_VALUE_FOR_LOG(FLAGS_log_dir.empty()) << ", "
          << EXPR_VALUE_FOR_LOG(conf_.force_disable_log_file) << ": "
          << EXPR_VALUE_FOR_LOG(log_to_file);

  if (log_to_file) {
    argv.push_back("-c");
    argv.push_back("logging_collector=on");
    // FLAGS_log_dir should already be set by tserver during startup.
    argv.push_back("-c");
    argv.push_back("log_directory=" + FLAGS_log_dir);
  }

  argv.push_back("-c");
  // TODO: we should probably load the metrics library in a different way once we let
  // users change the shared_preload_libraries conf parameter.
  if (FLAGS_pg_stat_statements_enabled) {
    argv.push_back("shared_preload_libraries=pg_stat_statements,yb_pg_metrics");
  } else {
    argv.push_back("shared_preload_libraries=yb_pg_metrics");
  }
  argv.push_back("-c");
  argv.push_back("yb_pg_metrics.node_name=" + FLAGS_metric_node_name);
  argv.push_back("-c");
  argv.push_back("yb_pg_metrics.port=" + std::to_string(FLAGS_pgsql_proxy_webserver_port));

  auto config_file_args = CHECK_RESULT(WritePgConfigFiles(conf_));
  argv.insert(argv.end(), config_file_args.begin(), config_file_args.end());

  if (FLAGS_pg_verbose_error_log) {
    argv.push_back("-c");
    argv.push_back("log_error_verbosity=VERBOSE");
  }

  pg_proc_.emplace(postgres_executable, argv);
  pg_proc_->ShareParentStderr();
  pg_proc_->ShareParentStdout();
  pg_proc_->SetParentDeathSignal(SIGINT);
  pg_proc_->InheritNonstandardFd(conf_.tserver_shm_fd);
  SetCommonEnv(&pg_proc_.get(), /* yb_enabled */ true);
  RETURN_NOT_OK(pg_proc_->Start());
  LOG(INFO) << "PostgreSQL server running as pid " << pg_proc_->pid();
  return Status::OK();
}

void PgWrapper::Kill() {
  WARN_NOT_OK(pg_proc_->Kill(SIGQUIT), "Kill PostgreSQL server failed");
}

Status PgWrapper::InitDb(bool yb_enabled) {
  const string initdb_program_path = GetInitDbExecutablePath();
  RETURN_NOT_OK(CheckExecutableValid(initdb_program_path));
  if (!Env::Default()->FileExists(initdb_program_path)) {
    return STATUS_FORMAT(IOError, "initdb not found at: $0", initdb_program_path);
  }

  vector<string> initdb_args { initdb_program_path, "-D", conf_.data_dir, "-U", "postgres" };
  Subprocess initdb_subprocess(initdb_program_path, initdb_args);
  SetCommonEnv(&initdb_subprocess, yb_enabled);
  int exit_code = 0;
  RETURN_NOT_OK(initdb_subprocess.Start());
  RETURN_NOT_OK(initdb_subprocess.Wait(&exit_code));
  if (exit_code != 0) {
    return STATUS_FORMAT(RuntimeError, "$0 failed with exit code $1",
                         initdb_program_path,
                         exit_code);
  }

  LOG(INFO) << "initdb completed successfully. Database initialized at " << conf_.data_dir;
  return Status::OK();
}

Status PgWrapper::InitDbLocalOnlyIfNeeded() {
  if (Env::Default()->FileExists(conf_.data_dir)) {
    LOG(INFO) << "Data directory " << conf_.data_dir << " already exists, skipping initdb";
    return Status::OK();
  }
  // Do not communicate with the YugaByte cluster at all. This function is only concerned with
  // setting up the local PostgreSQL data directory on this tablet server.
  return InitDb(/* yb_enabled */ false);
}

Result<int> PgWrapper::Wait() {
  if (!pg_proc_) {
    return STATUS(IllegalState,
                  "PostgreSQL child process has not been started, cannot wait for it to exit");
  }
  return pg_proc_->Wait();
}

Status PgWrapper::InitDbForYSQL(const string& master_addresses, const string& tmp_dir_base) {
  LOG(INFO) << "Running initdb to initialize YSQL cluster with master addresses "
            << master_addresses;
  PgProcessConf conf;
  conf.master_addresses = master_addresses;
  conf.pg_port = 0;  // We should not use this port.
  std::mt19937 rng{std::random_device()()};
  conf.data_dir = Format("$0/tmp_pg_data_$1", tmp_dir_base, rng());
  auto se = ScopeExit([&conf] {
    auto is_dir = Env::Default()->IsDirectory(conf.data_dir);
    if (is_dir.ok()) {
      if (is_dir.get()) {
        Status del_status = Env::Default()->DeleteRecursively(conf.data_dir);
        if (!del_status.ok()) {
          LOG(WARNING) << "Failed to delete directory " << conf.data_dir;
        }
      }
    } else if (!is_dir.status().IsNotFound()) {
      LOG(WARNING) << "Failed to check directory existence for " << conf.data_dir << ": "
                   << is_dir.status();
    }
  });
  PgWrapper pg_wrapper(conf);
  auto start_time = std::chrono::steady_clock::now();
  Status initdb_status = pg_wrapper.InitDb(/* yb_enabled */ true);
  auto elapsed_time = std::chrono::steady_clock::now() - start_time;
  LOG(INFO)
      << "initdb took "
      << std::chrono::duration_cast<std::chrono::milliseconds>(elapsed_time).count() << " ms";
  if (!initdb_status.ok()) {
    LOG(ERROR) << "initdb failed: " << initdb_status;
  }
  return initdb_status;
}

string PgWrapper::GetPostgresExecutablePath() {
  return JoinPathSegments(GetPostgresInstallRoot(), "bin", "postgres");
}

string PgWrapper::GetInitDbExecutablePath() {
  return JoinPathSegments(GetPostgresInstallRoot(), "bin", "initdb");
}

Status PgWrapper::CheckExecutableValid(const std::string& executable_path) {
  if (VERIFY_RESULT(Env::Default()->IsExecutableFile(executable_path))) {
    return Status::OK();
  }
  return STATUS_FORMAT(NotFound, "Not an executable file: $0", executable_path);
}

void PgWrapper::SetCommonEnv(Subprocess* proc, bool yb_enabled) {
  // Used to resolve relative paths during YB init within PG code.
  // Needed because PG changes its current working dir to a data dir.
  char cwd[PATH_MAX];
  CHECK(getcwd(cwd, sizeof(cwd)) != nullptr);
  proc->SetEnv("YB_WORKING_DIR", cwd);
  // A temporary workaround for a failure to look up a user name by uid in an LDAP environment.
  proc->SetEnv("YB_PG_FALLBACK_SYSTEM_USER_NAME", "postgres");
  proc->SetEnv("YB_PG_ALLOW_RUNNING_AS_ANY_USER", "1");
  if (yb_enabled) {
    proc->SetEnv("YB_ENABLED_IN_POSTGRES", "1");
    proc->SetEnv("FLAGS_pggate_master_addresses", conf_.master_addresses);
    proc->SetEnv("FLAGS_pggate_tserver_shm_fd", std::to_string(conf_.tserver_shm_fd));
    // Postgres process can't compute default certs dir by itself
    // as it knows nothing about t-server's root data directory.
    // Solution is to specify it explicitly.
    proc->SetEnv("FLAGS_certs_dir", conf_.certs_dir);
    proc->SetEnv("FLAGS_certs_for_client_dir", conf_.certs_for_client_dir);

    proc->SetEnv("YB_PG_TRANSACTIONS_ENABLED", FLAGS_pg_transactions_enabled ? "1" : "0");

#ifdef ADDRESS_SANITIZER
    // Disable reporting signal-unsafe behavior for PostgreSQL because it does a lot of work in
    // signal handlers on shutdown.

    const char* asan_options = getenv("ASAN_OPTIONS");
    proc->SetEnv(
        "ASAN_OPTIONS",
        std::string(asan_options ? asan_options : "") + " report_signal_unsafe=0");
#endif

    // Pass non-default flags to the child process using FLAGS_... environment variables.
    static const std::vector<string> explicit_flags{"pggate_master_addresses",
                                                    "pggate_tserver_shm_fd",
                                                    "certs_dir",
                                                    "certs_for_client_dir"};
    std::vector<google::CommandLineFlagInfo> flag_infos;
    google::GetAllFlags(&flag_infos);
    for (const auto& flag_info : flag_infos) {
      // Skip the flags that we set explicitly using conf_ above.
      if (!flag_info.is_default &&
          std::find(explicit_flags.begin(),
                    explicit_flags.end(),
                    flag_info.name) == explicit_flags.end()) {
        proc->SetEnv("FLAGS_" + flag_info.name, flag_info.current_value);
      }
    }
  } else {
    proc->SetEnv("YB_PG_LOCAL_NODE_INITDB", "1");
  }
}

// ------------------------------------------------------------------------------------------------
// PgSupervisor: monitoring a PostgreSQL child process and restarting if needed
// ------------------------------------------------------------------------------------------------

PgSupervisor::PgSupervisor(PgProcessConf conf)
    : conf_(std::move(conf)) {
}

Status PgSupervisor::Start() {
  std::lock_guard<std::mutex> lock(mtx_);
  RETURN_NOT_OK(ExpectStateUnlocked(PgProcessState::kNotStarted));
  RETURN_NOT_OK(CleanupOldServerUnlocked());
  LOG(INFO) << "Starting PostgreSQL server";
  RETURN_NOT_OK(StartServerUnlocked());

  Status status = Thread::Create(
      "pg_supervisor", "pg_supervisor", &PgSupervisor::RunThread, this, &supervisor_thread_);
  if (!status.ok()) {
    supervisor_thread_.reset();
    return status;
  }

  state_ = PgProcessState::kRunning;

  return Status::OK();
}

CHECKED_STATUS PgSupervisor::CleanupOldServerUnlocked() {
  std::string postmaster_pid_filename = JoinPathSegments(conf_.data_dir, "postmaster.pid");
  if (Env::Default()->FileExists(postmaster_pid_filename)) {
    std::ifstream postmaster_pid_file;
    postmaster_pid_file.open(postmaster_pid_filename, std::ios_base::in);
    pid_t postgres_pid = 0;

    if (!postmaster_pid_file.eof()) {
      postmaster_pid_file >> postgres_pid;
    }

    if (!postmaster_pid_file.good() || postgres_pid == 0) {
      LOG(ERROR) << strings::Substitute("Error reading postgres process ID from file $0. $1 $2",
          postmaster_pid_filename, ErrnoToString(errno), errno);
    } else {
      LOG(WARNING) << "Killing older postgres process: " << postgres_pid;
      // If process does not exist, system may return "process does not exist" or
      // "operation not permitted" error. Ignore those errors.
      if (kill(postgres_pid, SIGKILL) != 0 && errno != ESRCH && errno != EPERM) {
        return STATUS(RuntimeError, "Unable to kill", Errno(errno));
      }
    }
    ignore_result(Env::Default()->DeleteFile(postmaster_pid_filename));
  }
  return Status::OK();
}

PgProcessState PgSupervisor::GetState() {
  std::lock_guard<std::mutex> lock(mtx_);
  return state_;
}

CHECKED_STATUS PgSupervisor::ExpectStateUnlocked(PgProcessState expected_state) {
  if (state_ != expected_state) {
    return STATUS_FORMAT(
        IllegalState, "Expected PostgreSQL server state to be $0, got $1", expected_state, state_);
  }
  return Status::OK();
}

CHECKED_STATUS PgSupervisor::StartServerUnlocked() {
  if (pg_wrapper_) {
    return STATUS(IllegalState, "Expecting pg_wrapper_ to not be set");
  }
  pg_wrapper_.emplace(conf_);
  auto start_status = pg_wrapper_->Start();
  if (!start_status.ok()) {
    pg_wrapper_.reset();
    return start_status;
  }
  return Status::OK();
}

void PgSupervisor::RunThread() {
  while (true) {
    Result<int> wait_result = pg_wrapper_->Wait();
    if (wait_result.ok()) {
      int ret_code = *wait_result;
      if (ret_code == 0) {
        LOG(INFO) << "PostgreSQL server exited normally";
      } else {
        LOG(WARNING) << "PostgreSQL server exited with code " << ret_code;
      }
      pg_wrapper_.reset();
    } else {
      // TODO: a better way to handle this error.
      LOG(WARNING) << "Failed when waiting for PostgreSQL server to exit: "
                   << wait_result.status() << ", waiting a bit";
      std::this_thread::sleep_for(1s);
      continue;
    }

    {
      std::lock_guard<std::mutex> lock(mtx_);
      if (state_ == PgProcessState::kStopping) {
        break;
      }
      LOG(INFO) << "Restarting PostgreSQL server";
      Status start_status = StartServerUnlocked();
      if (!start_status.ok()) {
        // TODO: a better way to handle this error.
        LOG(WARNING) << "Failed trying to start PostgreSQL server: "
                     << start_status << ", waiting a bit";
        std::this_thread::sleep_for(1s);
      }
    }
  }
}

void PgSupervisor::Stop() {
  {
    std::lock_guard<std::mutex> lock(mtx_);
    state_ = PgProcessState::kStopping;
    if (pg_wrapper_) {
      pg_wrapper_->Kill();
    }
  }
  supervisor_thread_->Join();
}

}  // namespace pgwrapper
}  // namespace yb