folly/io/ShutdownSocketSet.h

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #include <atomic>
  20 #include <cstdlib>
  21 #include <memory>
  22
  23 #include <boost/noncopyable.hpp>
  24
  25 #include <folly/File.h>
  26
  27 namespace folly {
  28
  29 /**
  30  * Set of sockets that allows immediate, take-no-prisoners abort.
  31  */
  32 class ShutdownSocketSet : private boost::noncopyable {
  33  public:
  34   /**
  35    * Create a socket set that can handle file descriptors < maxFd.
  36    * The default value (256Ki) is high enough for just about all
  37    * applications, even if you increased the number of file descriptors
  38    * on your system.
  39    */
  40   explicit ShutdownSocketSet(int maxFd = 1 << 18);
  41
  42   // Singleton instance used by all thrift servers.
  43   // May return nullptr on startup/shutdown.
  44   static std::shared_ptr<ShutdownSocketSet> getInstance();
  45
  46   /**
  47    * Add an already open socket to the list of sockets managed by
  48    * ShutdownSocketSet. You MUST close the socket by calling
  49    * ShutdownSocketSet::close (which will, as a side effect, also handle EINTR
  50    * properly) and not by calling close() on the file descriptor.
  51    */
  52   void add(int fd);
  53
  54   /**
  55    * Remove a socket from the list of sockets managed by ShutdownSocketSet.
  56    * Note that remove() might block! (which we lamely implement by
  57    * sleep()-ing) in the extremely rare case that the fd is currently
  58    * being shutdown().
  59    */
  60   void remove(int fd);
  61
  62   /**
  63    * Close a socket managed by ShutdownSocketSet. Returns the same return code
  64    * as ::close() (and sets errno accordingly).
  65    */
  66   int close(int fd);
  67
  68   /**
  69    * Shut down a socket. If abortive is true, we perform an abortive
  70    * shutdown (send RST rather than FIN). Note that we might still end up
  71    * sending FIN due to the rather interesting implementation.
  72    *
  73    * This is async-signal-safe and ignores errors.  Obviously, subsequent
  74    * read() and write() operations to the socket will fail. During normal
  75    * operation, just call ::shutdown() on the socket.
  76    */
  77   void shutdown(int fd, bool abortive=false);
  78
  79   /**
  80    * Immediate shutdown of all connections. This is a hard-hitting hammer;
  81    * all reads and writes will return errors and no new connections will
  82    * be accepted.
  83    *
  84    * To be used only in dire situations. We're using it from the failure
  85    * signal handler to close all connections quickly, even though the server
  86    * might take multiple seconds to finish crashing.
  87    *
  88    * The optional bool parameter indicates whether to set the active
  89    * connections in to not linger.  The effect of that includes RST packets
  90    * being immediately sent to clients which will result
  91    * in errors (and not normal EOF) on the client side.  This also causes
  92    * the local (ip, tcp port number) tuple to be reusable immediately, instead
  93    * of having to wait the standard amount of time.  For full details see
  94    * the `shutdown` method of `ShutdownSocketSet` (incl. notes about the
  95    * `abortive` parameter).
  96    *
  97    * This is async-signal-safe and ignores errors.
  98    */
  99   void shutdownAll(bool abortive=false);
 100
 101  private:
 102   void doShutdown(int fd, bool abortive);
 103
 104   // State transitions:
 105   // add():
 106   //   FREE -> IN_USE
 107   //
 108   // close():
 109   //   IN_USE -> (::close()) -> FREE
 110   //   SHUT_DOWN -> (::close()) -> FREE
 111   //   IN_SHUTDOWN -> MUST_CLOSE
 112   //   (If the socket is currently being shut down, we must defer the
 113   //    ::close() until the shutdown completes)
 114   //
 115   // shutdown():
 116   //   IN_USE -> IN_SHUTDOWN
 117   //   (::shutdown())
 118   //   IN_SHUTDOWN -> SHUT_DOWN
 119   //   MUST_CLOSE -> (::close()) -> FREE
 120   enum State : uint8_t {
 121     FREE = 0,
 122     IN_USE,
 123     IN_SHUTDOWN,
 124     SHUT_DOWN,
 125     MUST_CLOSE
 126   };
 127
 128   struct Free {
 129     template <class T>
 130     void operator()(T* ptr) const {
 131       ::free(ptr);
 132     }
 133   };
 134
 135   const int maxFd_;
 136   std::unique_ptr<std::atomic<uint8_t>[], Free> data_;
 137   folly::File nullFile_;
 138 };
 139
 140 } // namespace folly