Implement unregister on shutdown

arangod/Cluster/AgencyComm.h

@@ -255,6 +255,14 @@ struct AgencyWriteTransaction : public AgencyTransaction {
     operations.push_back(operation);
   }
   
+  //////////////////////////////////////////////////////////////////////////////
+  /// @brief shortcut to create a transaction with one operation
+  //////////////////////////////////////////////////////////////////////////////
+
+  explicit AgencyWriteTransaction(std::vector<AgencyOperation> const& _operations)
+  : operations(_operations) {
+  }
+  
   //////////////////////////////////////////////////////////////////////////////
   /// @brief shortcut to create a transaction with one operation and a
   /// precondition

arangod/Cluster/ClusterFeature.cpp

@@ -48,6 +48,7 @@ using namespace arangodb::options;
 ClusterFeature::ClusterFeature(application_features::ApplicationServer* server)
     : ApplicationFeature(server, "Cluster"),
       _username("root"),
+      _unregisterOnShutdown(false),
       _enableCluster(false),
       _heartbeatThread(nullptr),
       _heartbeatInterval(0),
@@ -466,7 +467,7 @@ void ClusterFeature::unprepare() {
     if (_heartbeatThread != nullptr) {
       _heartbeatThread->beginShutdown();
     }
-
+    
     // change into shutdown state
     ServerState::instance()->setState(ServerState::STATE_SHUTDOWN);
 
@@ -483,6 +484,9 @@ void ClusterFeature::unprepare() {
         }
       }
     }
+    if (_unregisterOnShutdown) {
+      ServerState::instance()->unregister();
+    }
   }
 
   ClusterComm::cleanup();
@@ -524,3 +528,7 @@ void ClusterFeature::unprepare() {
   // ClusterComm::cleanup();
   AgencyComm::cleanup();
 }
+
+void ClusterFeature::setUnregisterOnShutdown(bool unregisterOnShutdown) {
+  _unregisterOnShutdown = unregisterOnShutdown;
+}

arangod/Cluster/ClusterFeature.h

@@ -68,7 +68,10 @@ class ClusterFeature : public application_features::ApplicationFeature {
     return "/_api/agency/agency-callbacks";
   };
 
+  void setUnregisterOnShutdown(bool);
+
  private:
+  bool _unregisterOnShutdown;
   bool _enableCluster;
   HeartbeatThread* _heartbeatThread;
   uint64_t _heartbeatInterval;

arangod/Cluster/ServerState.cpp

@@ -209,6 +209,33 @@ ServerState::RoleEnum ServerState::getRole() {
   return loadRole();
 }
 
+bool ServerState::unregister() {
+  TRI_ASSERT(!getId().empty());
+  
+  std::string const& id = getId();
+
+  std::string localInfoEncoded = StringUtils::urlEncode(_localInfo);
+  AgencyOperation deleteLocalIdMap("Target/MapLocalToID/" + localInfoEncoded, AgencySimpleOperationType::DELETE_OP);
+  
+  std::vector<AgencyOperation> operations = {deleteLocalIdMap};
+
+  auto role = loadRole();
+  const std::string agencyKey = roleToAgencyKey(role);
+  TRI_ASSERT(isClusterRole(role));
+  if (role == ROLE_COORDINATOR || role == ROLE_PRIMARY) {
+    operations.push_back(AgencyOperation("Plan/" + agencyKey + "/" + id, AgencySimpleOperationType::DELETE_OP));
+    operations.push_back(AgencyOperation("Current/" + agencyKey + "/" + id, AgencySimpleOperationType::DELETE_OP));
+  }
+  
+  AgencyWriteTransaction unregisterTransaction(operations);
+  
+  AgencyComm comm;
+  AgencyCommResult result;
+  
+  result = comm.sendTransactionWithFailover(unregisterTransaction);
+  return result.successful();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 /// @brief try to register with a role
 ////////////////////////////////////////////////////////////////////////////////

arangod/Cluster/ServerState.h

@@ -134,6 +134,8 @@ class ServerState {
   RoleEnum getRole();
   
   bool registerWithRole(RoleEnum);
+  
+  bool unregister();
 
   /// @brief set the server role
   void setRole(RoleEnum);

arangod/RestHandler/RestShutdownHandler.cpp

@@ -24,6 +24,7 @@
 #include "RestShutdownHandler.h"
 
 #include "Rest/HttpRequest.h"
+#include "Cluster/ClusterFeature.h"
 
 #include <velocypack/Builder.h>
 #include <velocypack/velocypack-aliases.h>
@@ -47,6 +48,13 @@ HttpHandler::status_t RestShutdownHandler::execute() {
     return HttpHandler::status_t(HANDLER_DONE);
   }
 
+  bool found;
+  std::string const& remove = _request->value("remove_from_cluster", found);
+  if (found && remove == "1") {
+    ClusterFeature* clusterFeature = ApplicationServer::getFeature<ClusterFeature>("Cluster");
+    clusterFeature->setUnregisterOnShutdown(true);
+  }
+
   ApplicationServer::server->beginShutdown();
 
   try {