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
   9 #include <unistd.h>
  10 #include <ev++.h>
  11
  12 #include "dinit-log.h"
  13 #include "control-cmds.h"
  14 #include "service-listener.h"
  15 #include "cpbuffer.h"
  16
  17 // Control connection for dinit
  18
  19 // TODO: Use the input buffer as a circular buffer, instead of chomping data from
  20 // the front using a data move.
  21
  22 // forward-declaration of callback:
  23 static void control_conn_cb(struct ev_loop * loop, ev_io * w, int revents);
  24
  25 class ControlConn;
  26
  27 // Pointer to the control connection that is listening for rollback completion
  28 extern ControlConn * rollback_handler_conn;
  29
  30 extern int active_control_conns;
  31
  32 // "packet" format:
  33 // (1 byte) packet type
  34 // (N bytes) additional data (service name, etc)
  35 //   for LOADSERVICE/FINDSERVICE:
  36 //      (2 bytes) service name length
  37 //      (M bytes) service name (without nul terminator)
  38
  39 // Information packet:
  40 // (1 byte) packet type, >= 100
  41 // (1 byte) packet length (including all fields)
  42 //       N bytes: packet data (N = (length - 2))
  43
  44 class ServiceSet;
  45 class ServiceRecord;
  46
  47 class ControlConn : private ServiceListener
  48 {
  49     friend void control_conn_cb(struct ev_loop *, ev_io *, int);
  50
  51     struct ev_io iob;
  52     struct ev_loop *loop;
  53     ServiceSet *service_set;
  54
  55     bool bad_conn_close = false; // close when finished output?
  56     bool oom_close = false;      // send final 'out of memory' indicator
  57
  58     // The packet length before we need to re-check if the packet is complete.
  59     // processPacket() will not be called until the packet reaches this size.
  60     int chklen;
  61
  62     // Receive buffer
  63     CPBuffer rbuf;
  64
  65     template <typename T> using list = std::list<T>;
  66     template <typename T> using vector = std::vector<T>;
  67
  68     // A mapping between service records and their associated numerical identifier used
  69     // in communction
  70     using handle_t = uint32_t;
  71     std::unordered_multimap<ServiceRecord *, handle_t> serviceKeyMap;
  72     std::unordered_map<handle_t, ServiceRecord *> keyServiceMap;
  73
  74     // Buffer for outgoing packets. Each outgoing back is represented as a vector<char>.
  75     list<vector<char>> outbuf;
  76     // Current index within the first outgoing packet (all previous bytes have been sent).
  77     unsigned outpkt_index = 0;
  78
  79     // Queue a packet to be sent
  80     //  Returns:  true if the packet was successfully queued, false if otherwise
  81     //            (eg if out of memory); in the latter case the connection might
  82     //            no longer be valid (iff there are no outgoing packets queued).
  83     bool queuePacket(vector<char> &&v) noexcept;
  84     bool queuePacket(const char *pkt, unsigned size) noexcept;
  85
  86     // Process a packet. Can cause the ControlConn to be deleted iff there are no
  87     // outgoing packets queued.
  88     // Throws:
  89     //    std::bad_alloc - if an out-of-memory condition prevents processing
  90     void processPacket();
  91
  92     // Process a STARTSERVICE/STOPSERVICE packet. May throw std::bad_alloc.
  93     void processStartStop(int pktType);
  94
  95     // Process a FINDSERVICE/LOADSERVICE packet. May throw std::bad_alloc.
  96     void processFindLoad(int pktType);
  97
  98     // Process an UNPINSERVICE packet. May throw std::bad_alloc.
  99     void processUnpinService();
 100
 101     // Notify that data is ready to be read from the socket. Returns true in cases where the
 102     // connection was deleted with potentially pending outgoing packets.
 103     bool dataReady() noexcept;
 104
 105     void sendData() noexcept;
 106
 107     // Allocate a new handle for a service; may throw std::bad_alloc
 108     handle_t allocateServiceHandle(ServiceRecord *record);
 109
 110     ServiceRecord *findServiceForKey(uint32_t key)
 111     {
 112         try {
 113             return keyServiceMap.at(key);
 114         }
 115         catch (std::out_of_range &exc) {
 116             return nullptr;
 117         }
 118     }
 119
 120     // Close connection due to out-of-memory condition.
 121     void doOomClose()
 122     {
 123         bad_conn_close = true;
 124         oom_close = true;
 125         ev_io_set(&iob, iob.fd, EV_WRITE);
 126     }
 127
 128     // Process service event broadcast.
 129     void serviceEvent(ServiceRecord * service, ServiceEvent event) noexcept final override
 130     {
 131         // For each service handle corresponding to the event, send an information packet.
 132         auto range = serviceKeyMap.equal_range(service);
 133         auto & i = range.first;
 134         auto & end = range.second;
 135         try {
 136             while (i != end) {
 137                 uint32_t key = i->second;
 138                 std::vector<char> pkt;
 139                 constexpr int pktsize = 3 + sizeof(key);
 140                 pkt.reserve(pktsize);
 141                 pkt.push_back(DINIT_IP_SERVICEEVENT);
 142                 pkt.push_back(pktsize);
 143                 char * p = (char *) &key;
 144                 for (int j = 0; j < (int)sizeof(key); j++) {
 145                     pkt.push_back(*p++);
 146                 }
 147                 pkt.push_back(static_cast<char>(event));
 148                 queuePacket(std::move(pkt));
 149                 ++i;
 150             }
 151         }
 152         catch (std::bad_alloc &exc) {
 153             doOomClose();
 154         }
 155     }
 156
 157     public:
 158     ControlConn(struct ev_loop * loop, ServiceSet * service_set, int fd) : loop(loop), service_set(service_set), chklen(0)
 159     {
 160         ev_io_init(&iob, control_conn_cb, fd, EV_READ);
 161         iob.data = this;
 162         ev_io_start(loop, &iob);
 163
 164         active_control_conns++;
 165     }
 166
 167     bool rollbackComplete() noexcept;
 168
 169     virtual ~ControlConn() noexcept;
 170 };
 171
 172
 173 static void control_conn_cb(struct ev_loop * loop, ev_io * w, int revents)
 174 {
 175     ControlConn *conn = (ControlConn *) w->data;
 176     if (revents & EV_READ) {
 177         if (conn->dataReady()) {
 178             // ControlConn was deleted
 179             return;
 180         }
 181     }
 182     if (revents & EV_WRITE) {
 183         conn->sendData();
 184     }
 185 }
 186
 187 #endif