folly/io/ShutdownSocketSet.h

   1 /*
   2  * Copyright 2015 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(size_t maxFd = 1 << 18);
  41
  42   /**
  43    * Add an already open socket to the list of sockets managed by
  44    * ShutdownSocketSet. You MUST close the socket by calling
  45    * ShutdownSocketSet::close (which will, as a side effect, also handle EINTR
  46    * properly) and not by calling close() on the file descriptor.
  47    */
  48   void add(int fd);
  49
  50   /**
  51    * Remove a socket from the list of sockets managed by ShutdownSocketSet.
  52    * Note that remove() might block! (which we lamely implement by
  53    * sleep()-ing) in the extremely rare case that the fd is currently
  54    * being shutdown().
  55    */
  56   void remove(int fd);
  57
  58   /**
  59    * Close a socket managed by ShutdownSocketSet. Returns the same return code
  60    * as ::close() (and sets errno accordingly).
  61    */
  62   int close(int fd);
  63
  64   /**
  65    * Shut down a socket. If abortive is true, we perform an abortive
  66    * shutdown (send RST rather than FIN). Note that we might still end up
  67    * sending FIN due to the rather interesting implementation.
  68    *
  69    * This is async-signal-safe and ignores errors.  Obviously, subsequent
  70    * read() and write() operations to the socket will fail. During normal
  71    * operation, just call ::shutdown() on the socket.
  72    */
  73   void shutdown(int fd, bool abortive=false);
  74
  75   /**
  76    * Shut down all sockets managed by ShutdownSocketSet. This is
  77    * async-signal-safe and ignores errors.
  78    */
  79   void shutdownAll(bool abortive=false);
  80
  81  private:
  82   void doShutdown(int fd, bool abortive);
  83
  84   // State transitions:
  85   // add():
  86   //   FREE -> IN_USE
  87   //
  88   // close():
  89   //   IN_USE -> (::close()) -> FREE
  90   //   SHUT_DOWN -> (::close()) -> FREE
  91   //   IN_SHUTDOWN -> MUST_CLOSE
  92   //   (If the socket is currently being shut down, we must defer the
  93   //    ::close() until the shutdown completes)
  94   //
  95   // shutdown():
  96   //   IN_USE -> IN_SHUTDOWN
  97   //   (::shutdown())
  98   //   IN_SHUTDOWN -> SHUT_DOWN
  99   //   MUST_CLOSE -> (::close()) -> FREE
 100   enum State : uint8_t {
 101     FREE = 0,
 102     IN_USE,
 103     IN_SHUTDOWN,
 104     SHUT_DOWN,
 105     MUST_CLOSE
 106   };
 107
 108   struct Free {
 109     template <class T>
 110     void operator()(T* ptr) const {
 111       ::free(ptr);
 112     }
 113   };
 114
 115   const size_t maxFd_;
 116   std::unique_ptr<std::atomic<uint8_t>[], Free> data_;
 117   folly::File nullFile_;
 118 };
 119
 120 }  // namespaces