Some work on exitmessage

louxiu · Feb 21, 2014 · d7a6dfa · d7a6dfa
1 parent 222c3e4
commit d7a6dfa
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 8 deletions.
diff --git a/src/RClient.cpp b/src/RClient.cpp
@@ -122,7 +122,7 @@ struct Option opts[] = {
 
     { None, 0, 0, 0, "" },
     { None, 0, 0, 0, "Rdm:" },
-    { QuitRdm, "quit-rdm", 'q', no_argument, "Tell server to shut down." },
+    { QuitRdm, "quit-rdm", 'q', optional_argument, "Tell server to shut down. If arg is broadcast:10 tell all rdm on the farm to quit with exit code 10." },
     { ConnectTimeout, "connect-timeout", 0, required_argument, "Timeout for connecting to rdm in ms (default " STR(DEFAULT_CONNECT_TIMEOUT)  ")." },
 
     { None, 0, 0, 0, "" },
@@ -789,9 +789,32 @@ bool RClient::parse(int &argc, char **argv)
         case XmlDiagnostics:
             addLog(RTags::CompilationErrorXml);
             break;
-        case QuitRdm:
-            addQuery(QueryMessage::Shutdown);
-            break;
+        case QuitRdm: {
+            const char *arg;
+            if (optarg) {
+                arg = optarg;
+            } else if (optind < argc && argv[optind][0] != '-') {
+                arg = argv[optind++];
+            }
+            if (arg) {
+                if (strncmp(arg, "broadcast:", 10)) {
+                    fprintf(stderr, "Invalid argument to -q\n");
+                    return 1;
+                }
+                bool ok;
+                const int exit = String(arg + 10).toLongLong(&ok);
+                if (!ok) {
+                    fprintf(stderr, "Invalid argument to -q\n");
+                    return 1;
+                }
+                String query;
+                Serializer serializer(query);
+                serializer << exit;
+                addQuery(QueryMessage::Shutdown, query);
+            } else {
+                addQuery(QueryMessage::Shutdown);
+            }
+            break; }
         case DeleteProject:
             addQuery(QueryMessage::DeleteProject, optarg);
             break;

diff --git a/src/RTagsMessage.h b/src/RTagsMessage.h
@@ -35,7 +35,8 @@ class RTagsMessage : public Message
         JobAnnouncementId,
         ProxyJobAnnouncementId,
         ClientMessageId,
-        ClientConnectedId
+        ClientConnectedId,
+        ExitMessageId
     };
 
     RTagsMessage(uint8_t id) : Message(id) {}

diff --git a/src/Server.cpp b/src/Server.cpp
@@ -23,6 +23,7 @@
 #include "VisitFileResponseMessage.h"
 #include "JobAnnouncementMessage.h"
 #include "ProxyJobAnnouncementMessage.h"
+#include "ExitMessage.h"
 #include "ClientMessage.h"
 #include "ClientConnectedMessage.h"
 #include "PreprocessJob.h"
@@ -112,11 +113,12 @@ static const bool debugMulti = getenv("RDM_DEBUG_MULTI");
 Server *Server::sInstance = 0;
 Server::Server()
     : mVerbose(false), mCurrentFileId(0), mThreadPool(0), mServerConnection(0), mHostName(Rct::hostName()),
-      mCompletionThread(0), mFirstRemote(0), mLastRemote(0), mAnnounced(false), mWorkPending(false)
+      mCompletionThread(0), mFirstRemote(0), mLastRemote(0), mAnnounced(false), mWorkPending(false), mExitCode(0)
 {
     Messages::registerMessage<JobRequestMessage>();
     Messages::registerMessage<JobResponseMessage>();
     Messages::registerMessage<ClientConnectedMessage>();
+    Messages::registerMessage<ExitMessage>();
 
     assert(!sInstance);
     sInstance = this;
@@ -383,6 +385,9 @@ void Server::onNewMessage(Message *message, Connection *connection)
         error() << m->raw();
         handleLogOutputMessage(static_cast<const LogOutputMessage&>(*m), connection);
         break;
+    case ExitMessage::MessageId:
+        handleExitMessage(static_cast<const ExitMessage&>(*m));
+        break;
     case VisitFileMessage::MessageId:
         handleVisitFileMessage(static_cast<const VisitFileMessage&>(*m), connection);
         break;
@@ -464,6 +469,17 @@ void Server::handleCompileMessage(CompileMessage &message, Connection *conn)
     index(message.arguments(), message.workingDirectory(), message.projects());
 }
 
+void Server::handleExitMessage(const ExitMessage &message)
+{
+    mExitCode = message.exitCode();
+    for (auto client : mClients) {
+        client->send(message);
+    }
+
+    EventLoop::eventLoop()->registerTimer(std::bind(&EventLoop::quit, EventLoop::eventLoop()),
+                                          1000, Timer::SingleShot);
+}
+
 void Server::handleLogOutputMessage(const LogOutputMessage &message, Connection *conn)
 {
     new LogObject(conn, message.level());
@@ -1330,7 +1346,15 @@ void Server::shutdown(const QueryMessage &query, Connection *conn)
         if (it.second)
             it.second->unload();
     }
-    EventLoop::eventLoop()->quit();
+    if (!query.query().isEmpty()) {
+        int exitCode;
+        Deserializer deserializer(query.query());
+        deserializer >> exitCode;
+        ExitMessage msg(exitCode);
+        handleExitMessage(msg);
+    } else {
+        EventLoop::eventLoop()->quit();
+    }
     conn->write("Shutting down");
     conn->finish();
 }

diff --git a/src/Server.h b/src/Server.h
@@ -48,6 +48,7 @@ class JobResponseMessage;
 class JobAnnouncementMessage;
 class ProxyJobAnnouncementMessage;
 class ClientConnectedMessage;
+class ExitMessage;
 class ClientMessage;
 class CompletionThread;
 class PreprocessJob;
@@ -114,6 +115,7 @@ class Server
     void stopServers();
     int mongooseStatistics(struct mg_connection *conn);
     void dumpJobs(Connection *conn);
+    int exitCode() const { return mExitCode; }
 private:
     bool selectProject(const Match &match, Connection *conn, unsigned int queryFlags);
     bool updateProject(const List<String> &projects, unsigned int queryFlags);
@@ -129,6 +131,7 @@ class Server
     void onNewMessage(Message *message, Connection *conn);
     void onConnectionDisconnected(Connection *o);
     void clearProjects();
+    void handleExitMessage(const ExitMessage &message);
     void handleCompileMessage(CompileMessage &message, Connection *conn);
     void handleIndexerMessage(const IndexerMessage &message, Connection *conn);
     void handleQueryMessage(const QueryMessage &message, Connection *conn);
@@ -233,6 +236,7 @@ class Server
     bool mAnnounced;
     Hash<std::string, Remote *> mRemotes;
     bool mWorkPending;
+    int mExitCode;
 };
 
 #endif
diff --git a/src/rdm.cpp b/src/rdm.cpp
@@ -596,5 +596,5 @@ int main(int argc, char** argv)
 
     const unsigned int ret = loop->exec();
     cleanupLogging();
-    return ret == EventLoop::Success ? 0 : 1;
+    return ret == EventLoop::Success ? server->exitCode() : 1;
 }