Split functions in main.cc to multiple headers (#1787)

CMakeLists.txt

@@ -196,7 +196,7 @@ endif()
 
 # kvrocks objects target
 file(GLOB_RECURSE KVROCKS_SRCS src/*.cc)
-list(FILTER KVROCKS_SRCS EXCLUDE REGEX src/main.cc)
+list(FILTER KVROCKS_SRCS EXCLUDE REGEX src/cli/main.cc)
 
 add_library(kvrocks_objs OBJECT ${KVROCKS_SRCS})
 
@@ -239,7 +239,7 @@ if(ENABLE_IPO)
 endif()
 
 # kvrocks main target
-add_executable(kvrocks src/main.cc)
+add_executable(kvrocks src/cli/main.cc)
 target_link_libraries(kvrocks PRIVATE kvrocks_objs ${EXTERNAL_LIBS})
 
 # kvrocks2redis sync tool

src/cli/daemon_util.h

@@ -0,0 +1,127 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ *
+ */
+
+#pragma once
+
+#include <config/config.h>
+#include <glog/logging.h>
+#include <signal.h>
+#include <sys/socket.h>
+#include <sys/stat.h>
+#include <sys/un.h>
+
+#include <cstdlib>
+
+#include "unique_fd.h"
+
+inline bool SupervisedUpstart() {
+  const char *upstart_job = getenv("UPSTART_JOB");
+  if (!upstart_job) {
+    LOG(WARNING) << "upstart supervision requested, but UPSTART_JOB not found";
+    return false;
+  }
+  LOG(INFO) << "supervised by upstart, will stop to signal readiness";
+  raise(SIGSTOP);
+  unsetenv("UPSTART_JOB");
+  return true;
+}
+
+inline bool SupervisedSystemd() {
+  const char *notify_socket = getenv("NOTIFY_SOCKET");
+  if (!notify_socket) {
+    LOG(WARNING) << "systemd supervision requested, but NOTIFY_SOCKET not found";
+    return false;
+  }
+
+  auto fd = UniqueFD(socket(AF_UNIX, SOCK_DGRAM, 0));
+  if (!fd) {
+    LOG(WARNING) << "Cannot connect to systemd socket " << notify_socket;
+    return false;
+  }
+
+  sockaddr_un su;
+  memset(&su, 0, sizeof(su));
+  su.sun_family = AF_UNIX;
+  strncpy(su.sun_path, notify_socket, sizeof(su.sun_path) - 1);
+  su.sun_path[sizeof(su.sun_path) - 1] = '\0';
+  if (notify_socket[0] == '@') su.sun_path[0] = '\0';
+
+  iovec iov;
+  memset(&iov, 0, sizeof(iov));
+  std::string ready = "READY=1";
+  iov.iov_base = &ready[0];
+  iov.iov_len = ready.size();
+
+  msghdr hdr;
+  memset(&hdr, 0, sizeof(hdr));
+  hdr.msg_name = &su;
+  hdr.msg_namelen = offsetof(struct sockaddr_un, sun_path) + strlen(notify_socket);
+  hdr.msg_iov = &iov;
+  hdr.msg_iovlen = 1;
+
+  int sendto_flags = 0;
+  unsetenv("NOTIFY_SOCKET");
+#ifdef HAVE_MSG_NOSIGNAL
+  sendto_flags |= MSG_NOSIGNAL;
+#endif
+  if (sendmsg(*fd, &hdr, sendto_flags) < 0) {
+    LOG(WARNING) << "Cannot send notification to systemd";
+    return false;
+  }
+  return true;
+}
+
+inline bool IsSupervisedMode(SupervisedMode mode) {
+  if (mode == kSupervisedAutoDetect) {
+    const char *upstart_job = getenv("UPSTART_JOB");
+    const char *notify_socket = getenv("NOTIFY_SOCKET");
+    if (upstart_job) {
+      mode = kSupervisedUpStart;
+    } else if (notify_socket) {
+      mode = kSupervisedSystemd;
+    }
+  }
+  if (mode == kSupervisedUpStart) {
+    return SupervisedUpstart();
+  } else if (mode == kSupervisedSystemd) {
+    return SupervisedSystemd();
+  }
+  return false;
+}
+
+inline void Daemonize() {
+  pid_t pid = fork();
+  if (pid < 0) {
+    LOG(ERROR) << "Failed to fork the process, err: " << strerror(errno);
+    exit(1);
+  }
+
+  if (pid > 0) exit(EXIT_SUCCESS);  // parent process
+  // change the file mode
+  umask(0);
+  if (setsid() < 0) {
+    LOG(ERROR) << "Failed to setsid, err: %s" << strerror(errno);
+    exit(1);
+  }
+
+  close(STDIN_FILENO);
+  close(STDOUT_FILENO);
+  close(STDERR_FILENO);
+}

src/cli/main.cc

@@ -0,0 +1,186 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ *
+ */
+
+#ifdef __linux__
+#define _XOPEN_SOURCE 700  // NOLINT
+#else
+#define _XOPEN_SOURCE
+#endif
+
+#include <event2/thread.h>
+#include <glog/logging.h>
+
+#include <iomanip>
+#include <ostream>
+
+#include "config.h"
+#include "daemon_util.h"
+#include "io_util.h"
+#include "pid_util.h"
+#include "scope_exit.h"
+#include "server/server.h"
+#include "signal_util.h"
+#include "storage/storage.h"
+#include "string_util.h"
+#include "time_util.h"
+#include "unique_fd.h"
+#include "vendor/crc64.h"
+#include "version.h"
+#include "version_util.h"
+
+Server *srv = nullptr;
+
+extern "C" void SignalHandler(int sig) {
+  if (srv && !srv->IsStopped()) {
+    LOG(INFO) << "Bye Bye";
+    srv->Stop();
+  }
+}
+
+struct NewOpt {
+  friend auto &operator<<(std::ostream &os, NewOpt) { return os << std::string(4, ' ') << std::setw(32); }
+} new_opt;
+
+static void PrintUsage(const char *program) {
+  std::cout << program << " implements the Redis protocol based on rocksdb" << std::endl
+            << "Usage:" << std::endl
+            << std::left << new_opt << "-c, --config <filename>"
+            << "set config file to <filename>, or `-` for stdin" << std::endl
+            << new_opt << "-v, --version"
+            << "print version information" << std::endl
+            << new_opt << "-h, --help"
+            << "print this help message" << std::endl
+            << new_opt << "--<config-key> <config-value>"
+            << "overwrite specific config option <config-key> to <config-value>" << std::endl;
+}
+
+static CLIOptions ParseCommandLineOptions(int argc, char **argv) {
+  using namespace std::string_view_literals;
+  CLIOptions opts;
+
+  for (int i = 1; i < argc; ++i) {
+    if ((argv[i] == "-c"sv || argv[i] == "--config"sv) && i + 1 < argc) {
+      opts.conf_file = argv[++i];
+    } else if (argv[i] == "-v"sv || argv[i] == "--version"sv) {
+      std::cout << "kvrocks " << PrintVersion << std::endl;
+      std::exit(0);
+    } else if (argv[i] == "-h"sv || argv[i] == "--help"sv) {
+      PrintUsage(*argv);
+      std::exit(0);
+    } else if (std::string_view(argv[i], 2) == "--" && std::string_view(argv[i]).size() > 2 && i + 1 < argc) {
+      auto key = std::string_view(argv[i] + 2);
+      opts.cli_options.emplace_back(key, argv[++i]);
+    } else {
+      PrintUsage(*argv);
+      std::exit(1);
+    }
+  }
+
+  return opts;
+}
+
+static void InitGoogleLog(const Config *config) {
+  FLAGS_minloglevel = config->log_level;
+  FLAGS_max_log_size = 100;
+  FLAGS_logbufsecs = 0;
+
+  if (util::EqualICase(config->log_dir, "stdout")) {
+    for (int level = google::INFO; level <= google::FATAL; level++) {
+      google::SetLogDestination(level, "");
+    }
+    FLAGS_stderrthreshold = google::ERROR;
+    FLAGS_logtostdout = true;
+  } else {
+    FLAGS_log_dir = config->log_dir + "/";
+    if (config->log_retention_days != -1) {
+      google::EnableLogCleaner(config->log_retention_days);
+    }
+  }
+}
+
+int main(int argc, char *argv[]) {
+  srand(static_cast<unsigned>(util::GetTimeStamp()));
+
+  google::InitGoogleLogging("kvrocks");
+  auto glog_exit = MakeScopeExit(google::ShutdownGoogleLogging);
+
+  evthread_use_pthreads();
+  auto event_exit = MakeScopeExit(libevent_global_shutdown);
+
+  signal(SIGPIPE, SIG_IGN);
+  SetupSigSegvAction(SignalHandler);
+
+  auto opts = ParseCommandLineOptions(argc, argv);
+
+  Config config;
+  Status s = config.Load(opts);
+  if (!s.IsOK()) {
+    std::cout << "Failed to load config. Error: " << s.Msg() << std::endl;
+    return 1;
+  }
+
+  crc64_init();
+  InitGoogleLog(&config);
+  LOG(INFO) << "kvrocks " << PrintVersion;
+  // Tricky: We don't expect that different instances running on the same port,
+  // but the server use REUSE_PORT to support the multi listeners. So we connect
+  // the listen port to check if the port has already listened or not.
+  if (!config.binds.empty()) {
+    uint32_t ports[] = {config.port, config.tls_port, 0};
+    for (uint32_t *port = ports; *port; ++port) {
+      if (util::IsPortInUse(*port)) {
+        LOG(ERROR) << "Could not create server TCP since the specified port[" << *port << "] is already in use";
+        return 1;
+      }
+    }
+  }
+  bool is_supervised = IsSupervisedMode(static_cast<SupervisedMode>(config.supervised_mode));
+  if (config.daemonize && !is_supervised) Daemonize();
+  s = CreatePidFile(config.pidfile);
+  if (!s.IsOK()) {
+    LOG(ERROR) << "Failed to create pidfile: " << s.Msg();
+    return 1;
+  }
+  auto pidfile_exit = MakeScopeExit([&config] { RemovePidFile(config.pidfile); });
+
+#ifdef ENABLE_OPENSSL
+  // initialize OpenSSL
+  if (config.tls_port || config.tls_replication) {
+    InitSSL();
+  }
+#endif
+
+  engine::Storage storage(&config);
+  s = storage.Open();
+  if (!s.IsOK()) {
+    LOG(ERROR) << "Failed to open: " << s.Msg();
+    return 1;
+  }
+  Server server(&storage, &config);
+  srv = &server;
+  s = srv->Start();
+  if (!s.IsOK()) {
+    LOG(ERROR) << "Failed to start server: " << s.Msg();
+    return 1;
+  }
+  srv->Join();
+
+  return 0;
+}

src/cli/pid_util.h

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ *
+ */
+
+#pragma once
+
+#include <fcntl.h>
+
+#include "io_util.h"
+#include "status.h"
+#include "unique_fd.h"
+
+inline Status CreatePidFile(const std::string &path) {
+  auto fd = UniqueFD(open(path.data(), O_RDWR | O_CREAT, 0660));
+  if (!fd) {
+    return Status::FromErrno();
+  }
+
+  std::string pid_str = std::to_string(getpid());
+  auto s = util::Write(*fd, pid_str);
+  if (!s.IsOK()) {
+    return s.Prefixed("failed to write to PID-file");
+  }
+
+  return Status::OK();
+}
+
+inline void RemovePidFile(const std::string &path) { std::remove(path.data()); }