src/control.h

   1 #ifndef DINIT_CONTROL_H
   2 #define DINIT_CONTROL_H
   3
   4 #include <list>
   5 #include <vector>
   6 #include <unordered_map>
   7 #include <limits>
   8 #include <cstddef>
   9
  10 #include <unistd.h>
  11
  12 #include "dasynq.h"
  13
  14 #include "dinit-log.h"
  15 #include "control-cmds.h"
  16 #include "service-listener.h"
  17 #include "cpbuffer.h"
  18
  19 // Control connection for dinit
  20
  21 using namespace dasynq;
  22 using EventLoop_t = EventLoop<NullMutex>;
  23
  24 class ControlConn;
  25 class ControlConnWatcher;
  26
  27 // forward-declaration of callback:
  28 static Rearm control_conn_cb(EventLoop_t *loop, ControlConnWatcher *watcher, int revents);
  29
  30 // Pointer to the control connection that is listening for rollback completion
  31 extern ControlConn * rollback_handler_conn;
  32
  33 extern int active_control_conns;
  34
  35 // "packet" format:
  36 // (1 byte) packet type
  37 // (N bytes) additional data (service name, etc)
  38 //   for LOADSERVICE/FINDSERVICE:
  39 //      (2 bytes) service name length
  40 //      (M bytes) service name (without nul terminator)
  41
  42 // Information packet:
  43 // (1 byte) packet type, >= 100
  44 // (1 byte) packet length (including all fields)
  45 //       N bytes: packet data (N = (length - 2))
  46
  47 class ServiceSet;
  48 class ServiceRecord;
  49
  50 class ControlConnWatcher : public EventLoop_t::BidiFdWatcher
  51 {
  52     inline Rearm receiveEvent(EventLoop_t &loop, int fd, int flags) noexcept;
  53
  54     Rearm readReady(EventLoop_t &loop, int fd) noexcept override
  55     {
  56         return receiveEvent(loop, fd, IN_EVENTS);
  57     }
  58
  59     Rearm writeReady(EventLoop_t &loop, int fd) noexcept override
  60     {
  61         return receiveEvent(loop, fd, OUT_EVENTS);
  62     }
  63
  64     public:
  65     int fd; // TODO this is already stored, find a better way to access it.
  66     EventLoop_t * eventLoop;
  67
  68     void setWatches(int flags)
  69     {
  70         EventLoop_t::BidiFdWatcher::setWatches(*eventLoop, flags);
  71     }
  72
  73     void registerWith(EventLoop_t &loop, int fd, int flags)
  74     {
  75         this->fd = fd;
  76         this->eventLoop = &loop;
  77         BidiFdWatcher<EventLoop_t>::addWatch(loop, fd, flags);
  78     }
  79 };
  80
  81 inline Rearm ControlConnWatcher::receiveEvent(EventLoop_t &loop, int fd, int flags) noexcept
  82 {
  83     return control_conn_cb(&loop, this, flags);
  84 }
  85
  86
  87 class ControlConn : private ServiceListener
  88 {
  89     friend Rearm control_conn_cb(EventLoop_t *loop, ControlConnWatcher *watcher, int revents);
  90
  91     ControlConnWatcher iob;
  92     EventLoop_t *loop;
  93     ServiceSet *service_set;
  94
  95     bool bad_conn_close = false; // close when finished output?
  96     bool oom_close = false;      // send final 'out of memory' indicator
  97
  98     // The packet length before we need to re-check if the packet is complete.
  99     // processPacket() will not be called until the packet reaches this size.
 100     int chklen;
 101
 102     // Receive buffer
 103     CPBuffer<1024> rbuf;
 104
 105     template <typename T> using list = std::list<T>;
 106     template <typename T> using vector = std::vector<T>;
 107
 108     // A mapping between service records and their associated numerical identifier used
 109     // in communction
 110     using handle_t = uint32_t;
 111     std::unordered_multimap<ServiceRecord *, handle_t> serviceKeyMap;
 112     std::unordered_map<handle_t, ServiceRecord *> keyServiceMap;
 113
 114     // Buffer for outgoing packets. Each outgoing back is represented as a vector<char>.
 115     list<vector<char>> outbuf;
 116     // Current index within the first outgoing packet (all previous bytes have been sent).
 117     unsigned outpkt_index = 0;
 118
 119     // Queue a packet to be sent
 120     //  Returns:  false if the packet could not be queued and a suitable error packet
 121     //              could not be sent/queued (the connection should be closed);
 122     //            true (with bad_conn_close == false) if the packet was successfully
 123     //              queued;
 124     //            true (with bad_conn_close == true) if the packet was not successfully
 125     //              queued (but a suitable error packate has been queued).
 126     // The in/out watch enabled state will also be set appropriately.
 127     bool queuePacket(vector<char> &&v) noexcept;
 128     bool queuePacket(const char *pkt, unsigned size) noexcept;
 129
 130     // Process a packet.
 131     //  Returns:  true (with bad_conn_close == false) if successful
 132     //            true (with bad_conn_close == true) if an error packet was queued
 133     //            false if an error occurred but no error packet could be queued
 134     //                (connection should be closed).
 135     // Throws:
 136     //    std::bad_alloc - if an out-of-memory condition prevents processing
 137     bool processPacket();
 138
 139     // Process a STARTSERVICE/STOPSERVICE packet. May throw std::bad_alloc.
 140     bool processStartStop(int pktType);
 141
 142     // Process a FINDSERVICE/LOADSERVICE packet. May throw std::bad_alloc.
 143     bool processFindLoad(int pktType);
 144
 145     // Process an UNPINSERVICE packet. May throw std::bad_alloc.
 146     bool processUnpinService();
 147
 148     // Notify that data is ready to be read from the socket. Returns true if the connection should
 149     // be closed.
 150     bool dataReady() noexcept;
 151
 152     bool sendData() noexcept;
 153
 154     // Allocate a new handle for a service; may throw std::bad_alloc
 155     handle_t allocateServiceHandle(ServiceRecord *record);
 156
 157     ServiceRecord *findServiceForKey(uint32_t key)
 158     {
 159         try {
 160             return keyServiceMap.at(key);
 161         }
 162         catch (std::out_of_range &exc) {
 163             return nullptr;
 164         }
 165     }
 166
 167     // Close connection due to out-of-memory condition.
 168     void doOomClose()
 169     {
 170         bad_conn_close = true;
 171         oom_close = true;
 172         iob.setWatches(OUT_EVENTS);
 173     }
 174
 175     // Process service event broadcast.
 176     // Note that this can potentially be called during packet processing (upon issuing
 177     // service start/stop orders etc).
 178     void serviceEvent(ServiceRecord * service, ServiceEvent event) noexcept final override
 179     {
 180         // For each service handle corresponding to the event, send an information packet.
 181         auto range = serviceKeyMap.equal_range(service);
 182         auto & i = range.first;
 183         auto & end = range.second;
 184         try {
 185             while (i != end) {
 186                 uint32_t key = i->second;
 187                 std::vector<char> pkt;
 188                 constexpr int pktsize = 3 + sizeof(key);
 189                 pkt.reserve(pktsize);
 190                 pkt.push_back(DINIT_IP_SERVICEEVENT);
 191                 pkt.push_back(pktsize);
 192                 char * p = (char *) &key;
 193                 for (int j = 0; j < (int)sizeof(key); j++) {
 194                     pkt.push_back(*p++);
 195                 }
 196                 pkt.push_back(static_cast<char>(event));
 197                 queuePacket(std::move(pkt));
 198                 ++i;
 199             }
 200         }
 201         catch (std::bad_alloc &exc) {
 202             doOomClose();
 203         }
 204     }
 205
 206     public:
 207     ControlConn(EventLoop_t * loop, ServiceSet * service_set, int fd) : loop(loop), service_set(service_set), chklen(0)
 208     {
 209         iob.registerWith(*loop, fd, IN_EVENTS);
 210         active_control_conns++;
 211     }
 212
 213     bool rollbackComplete() noexcept;
 214
 215     virtual ~ControlConn() noexcept;
 216 };
 217
 218
 219 static Rearm control_conn_cb(EventLoop_t * loop, ControlConnWatcher * watcher, int revents)
 220 {
 221     char * cc_addr = (reinterpret_cast<char *>(watcher)) - offsetof(ControlConn, iob);
 222     ControlConn *conn = reinterpret_cast<ControlConn *>(cc_addr);
 223     if (revents & IN_EVENTS) {
 224         if (conn->dataReady()) {
 225             delete conn;
 226             return Rearm::REMOVED;
 227         }
 228     }
 229     if (revents & OUT_EVENTS) {
 230         if (conn->sendData()) {
 231             delete conn;
 232             return Rearm::REMOVED;
 233         }
 234     }
 235
 236     return Rearm::NOOP;
 237 }
 238
 239 #endif