git.m6w6.name Git - m6w6/libmemcached/blob - test/lib/Server.cpp

   1 #include "Server.hpp"
   2 #include "Retry.hpp"
   3 #include "ForkAndExec.hpp"
   4
   5 #include <sys/wait.h>
   6 #include <unistd.h>
   7
   8 Server::Server(string binary_, Server::argv_t args_)
   9     : binary{move(binary_)}
  10     , args{move(args_)}
  11 {}
  12
  13 Server::~Server() {
  14   stop();
  15   wait();
  16   if (pipe != -1) {
  17     close(pipe);
  18   }
  19   if (holds_alternative<string>(socket_or_port)) {
  20     unlink(get<string>(socket_or_port).c_str());
  21   }
  22 }
  23
  24 static inline string extractArg(const Server::arg_t &arg_cont, const string &func_arg) {
  25   if (holds_alternative<string>(arg_cont)) {
  26     return get<string>(arg_cont);
  27   } else {
  28     return get<Server::arg_func_t>(arg_cont)(func_arg);
  29   }
  30 }
  31
  32 static inline void pushArg(vector<char *> &arr, const string &arg) {
  33   auto len = arg.size();
  34   auto str = arg.data(), end = str + len + 1;
  35   auto ptr = new char[len + 1];
  36   copy(str, end, ptr);
  37   arr.push_back(ptr);
  38 }
  39
  40 optional<string> Server::handleArg(vector<char *> &arr, const string &arg, const arg_func_t &next_arg) {
  41   pushArg(arr, arg);
  42   if (arg == "-U" || arg == "--udp-port") {
  43     auto port = next_arg(arg);
  44     pushArg(arr, port);
  45     pushArg(arr, "-p");
  46     pushArg(arr, port);
  47     socket_or_port = stoi(port);
  48     return port;
  49   } else if (arg == "-p" || arg == "--port") {
  50     auto port = next_arg(arg);
  51     pushArg(arr, port);
  52     socket_or_port = stoi(port);
  53     return port;
  54   } else if (arg == "-s" || arg == "--unix-socket") {
  55     auto sock = next_arg(arg);
  56     pushArg(arr, sock);
  57     socket_or_port = sock;
  58     return sock;
  59   } else if (arg == "-S" || arg == "--enable-sasl") {
  60     sasl = true;
  61   }
  62   return {};
  63 }
  64
  65 [[nodiscard]]
  66 vector<char *> Server::createArgv()  {
  67   vector<char *> arr;
  68
  69   pushArg(arr, binary);
  70   //pushArg(arr, "-v");
  71
  72   for (auto it = args.cbegin(); it != args.cend(); ++it) {
  73     if (holds_alternative<arg_t>(*it)) {
  74       // a single argument
  75       auto arg = extractArg(get<arg_t>(*it), binary);
  76       handleArg(arr, arg, [&it](const string &arg_) {
  77         return extractArg(get<arg_t>(*++it), arg_);
  78       });
  79     } else {
  80       // an argument pair
  81       auto &[one, two] = get<arg_pair_t>(*it);
  82       auto arg_one = extractArg(one, binary);
  83       auto arg_two = extractArg(two, arg_one);
  84
  85       auto next = handleArg(arr, arg_one, [&arg_two](const string &) {
  86         return arg_two;
  87       });
  88
  89       if (!next.has_value()) {
  90         pushArg(arr, arg_two);
  91       }
  92     }
  93   }
  94
  95   arr.push_back(nullptr);
  96
  97   return arr;
  98 }
  99
 100 optional<Server::ChildProc> Server::start() {
 101   if (!pid) {
 102     auto argv = createArgv();
 103     ForkAndExec fork_and_exec{binary.c_str(), argv.data()};
 104
 105     pipe = fork_and_exec.createPipe();
 106     pid = fork_and_exec();
 107
 108     for (auto argp : argv) {
 109       delete [] argp;
 110     }
 111
 112     if (!pid) {
 113       return {};
 114     }
 115   }
 116   return ChildProc{pid, pipe};
 117 }
 118
 119 bool Server::isListening() {
 120   MemcachedPtr memc;
 121
 122   if (holds_alternative<string>(socket_or_port)) {
 123     if (memcached_server_add_unix_socket(*memc, get<string>(socket_or_port).c_str())) {
 124       return false;
 125     }
 126   } else {
 127     if (memcached_server_add(*memc, "localhost", get<int>(socket_or_port))) {
 128       return false;
 129     }
 130   }
 131
 132   if (sasl) {
 133     memcached_behavior_set(*memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, 1);
 134     memcached_set_sasl_auth_data(*memc, "memcached", "memcached");
 135   }
 136
 137   Malloced stat(memcached_stat(*memc, nullptr, nullptr));
 138   if (!*stat || !stat->pid || stat->pid == -1) {
 139     return false;
 140   }
 141   if (stat->pid != pid) {
 142     cerr << "Another server is listening on " << socket_or_port
 143          << " (expected pid " << pid << " found pid " << stat->pid << ")\n";
 144     return false;
 145   }
 146
 147   return true;
 148 }
 149
 150 bool Server::ensureListening() {
 151   return Retry{[this] {
 152     again:
 153     start();
 154     if (!isListening()) {
 155       if (tryWait()){
 156         goto again;
 157       }
 158     }
 159     return isListening();
 160   }}();
 161 }
 162
 163 bool Server::stop() {
 164   if (!pid) {
 165     return true;
 166   }
 167   if (signalled[SIGTERM]) {
 168     return signal(SIGKILL);
 169   } else {
 170     return signal(SIGTERM);
 171   }
 172 }
 173
 174 bool Server::signal(int signo) {
 175   if (!pid) {
 176     return false;
 177   }
 178   signalled[signo] += 1;
 179   return 0 <= kill(pid, signo);
 180 }
 181
 182 bool Server::check() {
 183   return signal(0);
 184 }
 185
 186 bool Server::wait(int flags) {
 187   if (pid && pid == waitpid(pid, &status, flags)) {
 188     if (drain().length() && output != "Signal handled: Terminated.\n") {
 189       cerr << "Output of " << *this << ":\n";
 190
 191       istringstream iss{output};
 192       string line;
 193
 194       while (getline(iss, line)) {
 195         cerr << "  " << line << "\n";
 196       }
 197
 198       if (output.back() != '\n') {
 199         cerr << endl;
 200       }
 201       output.clear();
 202     }
 203     pid = 0;
 204     if (pipe != -1) {
 205       close(pipe);
 206       pipe = -1;
 207     }
 208     return true;
 209   }
 210   return false;
 211 }
 212
 213 bool Server::tryWait() {
 214   return wait(WNOHANG);
 215 }
 216
 217 Server::Server(const Server &s) {
 218   binary = s.binary;
 219   args = s.args;
 220   socket_or_port = s.socket_or_port;
 221 }
 222
 223 Server &Server::operator=(const Server &s) {
 224   binary = s.binary;
 225   args = s.args;
 226   socket_or_port = s.socket_or_port;
 227   return *this;
 228 }
 229
 230 pid_t Server::getPid() const {
 231   return pid;
 232 }
 233
 234 const string &Server::getBinary() const {
 235   return binary;
 236 }
 237
 238 const Server::argv_t &Server::getArgs() const {
 239   return args;
 240 }
 241
 242 const socket_or_port_t &Server::getSocketOrPort() const {
 243   return socket_or_port;
 244 }
 245
 246 int Server::getPipe() const {
 247   return pipe;
 248 }
 249
 250 string &Server::drain() {
 251   if (pipe != -1) {
 252     again:
 253     char read_buf[1<<12];
 254     auto read_len = read(pipe, read_buf, sizeof(read_buf));
 255
 256     if (read_len > 0) {
 257       output.append(read_buf, read_len);
 258       goto again;
 259     }
 260     if (read_len == -1) {
 261       switch (errno) {
 262       case EINTR:
 263         goto again;
 264       default:
 265         perror("Server::drain read()");
 266         [[fallthrough]];
 267       case EAGAIN:
 268 #if EWOULDBLOCK != EAGAIN
 269       case EWOULDBLOCK:
 270 #endif
 271         break;
 272       }
 273     }
 274   }
 275   return output;
 276 }