![]() |
RTBKit
0.9
Open-source framework to create real-time ad bidding systems.
|
00001 /* zookeeper_temporary_server.h -*- C++ -*- 00002 Eric Robert, 25 February 2013 00003 Copyright (c) 2012 Datacratic Inc. All rights reserved. 00004 00005 A temporary server for testing zookeeper services. 00006 */ 00007 00008 #pragma once 00009 00010 #include "jml/utils/environment.h" 00011 #include "jml/utils/file_functions.h" 00012 #include "jml/arch/timers.h" 00013 #include <fstream> 00014 #include <sys/stat.h> 00015 #include <sys/types.h> 00016 #include <sys/wait.h> 00017 #include <sys/un.h> 00018 #include <sys/socket.h> 00019 #include <sys/prctl.h> 00020 #include <signal.h> 00021 00022 namespace ZooKeeper { 00023 00024 struct TemporaryServer : boost::noncopyable 00025 { 00026 TemporaryServer(std::string uniquePath = "") 00027 : server(-1) 00028 { 00029 if (uniquePath == "") { 00030 ML::Env_Option<std::string> tmpDir("TMP", "./tmp"); 00031 uniquePath = ML::format("%s/zookeeper-temporary-server-%d", tmpDir.get(), getpid()); 00032 } 00033 00034 static int uid; 00035 port = 4096 + (getpid() % 16768) + uid++; 00036 00037 std::cerr << "starting zookeeper temporary server under " << uniquePath << std::endl; 00038 this->uniquePath = uniquePath; 00039 } 00040 00041 ~TemporaryServer() { 00042 shutdown(); 00043 } 00044 00045 void start() { 00046 if (uniquePath == "" || uniquePath[0] == '/' || uniquePath == "." || uniquePath == "..") { 00047 throw ML::Exception("unacceptable path"); 00048 } 00049 00050 createDirectory(); 00051 createConfig(); 00052 createServer(); 00053 } 00054 00055 void shutdown() { 00056 if (server == -1) 00057 return; 00058 00059 // Stop boost test framework from interpreting this as a problem... 00060 signal(SIGCHLD, SIG_DFL); 00061 00062 int res = kill(server, SIGTERM); 00063 if (res == -1) { 00064 throw ML::Exception(errno, "cannot kill zookeeper"); 00065 } 00066 00067 int status = 0; 00068 res = waitpid(server, &status, 0); 00069 if (res == -1) { 00070 throw ML::Exception(errno, "failed to wait for zookeeper to shutdown"); 00071 } 00072 00073 server = -1; 00074 00075 if (uniquePath != "") { 00076 std::cerr << "removing " << uniquePath << std::endl; 00077 int res = system(("rm -rf " + uniquePath).c_str()); 00078 if (res) { 00079 throw ML::Exception(errno, "failed to remove zookeeper path"); 00080 } 00081 } 00082 } 00083 00084 int getPort() const { 00085 return port; 00086 } 00087 00088 private: 00089 void createDirectory() { 00090 struct stat stats; 00091 int res = stat(uniquePath.c_str(), &stats); 00092 if (res != -1 || (errno != EEXIST && errno != ENOENT)) { 00093 throw ML::Exception(errno, "path " + uniquePath + " already exists"); 00094 } 00095 00096 res = system(ML::format("mkdir -p %s", uniquePath).c_str()); 00097 if (res == -1) { 00098 throw ML::Exception(errno, "couldn't create directory"); 00099 } 00100 00101 res = system(ML::format("mkdir -p %s/data", uniquePath).c_str()); 00102 if (res == -1) { 00103 throw ML::Exception(errno, "couldn't create data directory"); 00104 } 00105 } 00106 00107 void createConfig() { 00108 std::string filename = uniquePath + "/zoo.cfg"; 00109 std::ofstream file(filename); 00110 if (!file) { 00111 throw ML::Exception("couldn't create zoo.cfg"); 00112 } 00113 00114 std::cerr << "zookeeper is using port " << port << std::endl; 00115 00116 file << "tickTime=2000" << std::endl; 00117 file << "dataDir=" << uniquePath << "/data" << std::endl; 00118 file << "clientPort=" << port << std::endl; 00119 file << "dataLogDir=" << uniquePath << std::endl; 00120 file << "initLimit=10" << std::endl; 00121 file << "syncLimit=5" << std::endl; 00122 file << "maxClientCnxns=4096" << std::endl; 00123 } 00124 00125 void createServer() { 00126 int pid = fork(); 00127 if (pid == -1) { 00128 throw ML::Exception(errno, "fork"); 00129 } 00130 00131 if (pid == 0) { 00132 signal(SIGTERM, SIG_DFL); 00133 signal(SIGKILL, SIG_DFL); 00134 00135 int res = prctl(PR_SET_PDEATHSIG, SIGHUP); 00136 if(res == -1) { 00137 throw ML::Exception(errno, "prctl failed"); 00138 } 00139 00140 std::string home = getenv("HOME"); 00141 std::string path = home + "/local/bin/zookeeper/bin/zkServer.sh"; 00142 std::string file = uniquePath + "/zoo.cfg"; 00143 std::string logs = "ZOO_LOG_DIR=" + uniquePath; 00144 00145 char const * args[] = { path.c_str(), "start-foreground", file.c_str(), (char *) 0 }; 00146 char const * envp[] = { logs.c_str(), (char *) 0 }; 00147 00148 std::cerr << "running zookeeper from " << file << std::endl; 00149 00150 res = execvpe(path.c_str(), (char **) args, (char **) envp); 00151 if (res == -1) { 00152 throw ML::Exception(errno, "zookeeper failed to start"); 00153 } 00154 00155 throw ML::Exception(errno, "zookeeper server didn't start"); 00156 } 00157 else { 00158 server = pid; 00159 } 00160 00161 ML::sleep(1); 00162 } 00163 00164 int server; 00165 int port; 00166 std::string uniquePath; 00167 }; 00168 00169 } // namespace ZooKeeper 00170