folly/io/async/AsyncUDPSocket.h

   1 /*
   2  * Copyright 2014-present 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 <memory>
  20
  21 #include <folly/ScopeGuard.h>
  22 #include <folly/SocketAddress.h>
  23 #include <folly/io/IOBuf.h>
  24 #include <folly/io/async/AsyncSocketBase.h>
  25 #include <folly/io/async/AsyncSocketException.h>
  26 #include <folly/io/async/EventBase.h>
  27 #include <folly/io/async/EventHandler.h>
  28
  29 namespace folly {
  30
  31 /**
  32  * UDP socket
  33  */
  34 class AsyncUDPSocket : public EventHandler {
  35  public:
  36   enum class FDOwnership {
  37     OWNS,
  38     SHARED
  39   };
  40
  41   class ReadCallback {
  42    public:
  43     /**
  44      * Invoked when the socket becomes readable and we want buffer
  45      * to write to.
  46      *
  47      * NOTE: From socket we will end up reading at most `len` bytes
  48      *       and if there were more bytes in datagram, we will end up
  49      *       dropping them.
  50      */
  51      virtual void getReadBuffer(void** buf, size_t* len) noexcept = 0;
  52
  53     /**
  54      * Invoked when a new datagraom is available on the socket. `len`
  55      * is the number of bytes read and `truncated` is true if we had
  56      * to drop few bytes because of running out of buffer space.
  57      */
  58     virtual void onDataAvailable(const folly::SocketAddress& client,
  59                                  size_t len,
  60                                  bool truncated) noexcept = 0;
  61
  62     /**
  63      * Invoked when there is an error reading from the socket.
  64      *
  65      * NOTE: Since UDP is connectionless, you can still read from the socket.
  66      *       But you have to re-register readCallback yourself after
  67      *       onReadError.
  68      */
  69     virtual void onReadError(const AsyncSocketException& ex)
  70         noexcept = 0;
  71
  72     /**
  73      * Invoked when socket is closed and a read callback is registered.
  74      */
  75     virtual void onReadClosed() noexcept = 0;
  76
  77     virtual ~ReadCallback() = default;
  78   };
  79
  80   /**
  81    * Create a new UDP socket that will run in the
  82    * given eventbase
  83    */
  84   explicit AsyncUDPSocket(EventBase* evb);
  85   ~AsyncUDPSocket() override;
  86
  87   /**
  88    * Returns the address server is listening on
  89    */
  90   virtual const folly::SocketAddress& address() const {
  91     CHECK_NE(-1, fd_) << "Server not yet bound to an address";
  92     return localAddress_;
  93   }
  94
  95   /**
  96    * Bind the socket to the following address. If port is not
  97    * set in the `address` an ephemeral port is chosen and you can
  98    * use `address()` method above to get it after this method successfully
  99    * returns.
 100    */
 101   virtual void bind(const folly::SocketAddress& address);
 102
 103   /**
 104    * Use an already bound file descriptor. You can either transfer ownership
 105    * of this FD by using ownership = FDOwnership::OWNS or share it using
 106    * FDOwnership::SHARED. In case FD is shared, it will not be `close`d in
 107    * destructor.
 108    */
 109   virtual void setFD(int fd, FDOwnership ownership);
 110
 111   /**
 112    * Send the data in buffer to destination. Returns the return code from
 113    * ::sendmsg.
 114    */
 115   virtual ssize_t write(const folly::SocketAddress& address,
 116                         const std::unique_ptr<folly::IOBuf>& buf);
 117
 118   /**
 119    * Send data in iovec to destination. Returns the return code from sendmsg.
 120    */
 121   virtual ssize_t writev(const folly::SocketAddress& address,
 122                          const struct iovec* vec, size_t veclen);
 123
 124   /**
 125    * Start reading datagrams
 126    */
 127   virtual void resumeRead(ReadCallback* cob);
 128
 129   /**
 130    * Pause reading datagrams
 131    */
 132   virtual void pauseRead();
 133
 134   /**
 135    * Stop listening on the socket.
 136    */
 137   virtual void close();
 138
 139   /**
 140    * Get internal FD used by this socket
 141    */
 142   virtual int getFD() const {
 143     CHECK_NE(-1, fd_) << "Need to bind before getting FD out";
 144     return fd_;
 145   }
 146
 147   /**
 148    * Set reuse port mode to call bind() on the same address multiple times
 149    */
 150   virtual void setReusePort(bool reusePort) {
 151     reusePort_ = reusePort;
 152   }
 153
 154   /**
 155    * Set SO_REUSEADDR flag on the socket. Default is ON.
 156    */
 157   virtual void setReuseAddr(bool reuseAddr) {
 158     reuseAddr_ = reuseAddr;
 159   }
 160
 161   EventBase* getEventBase() const {
 162     return eventBase_;
 163   }
 164
 165  protected:
 166   virtual ssize_t sendmsg(int socket, const struct msghdr* message, int flags) {
 167     return ::sendmsg(socket, message, flags);
 168   }
 169
 170   // Non-null only when we are reading
 171   ReadCallback* readCallback_;
 172
 173  private:
 174   AsyncUDPSocket(const AsyncUDPSocket&) = delete;
 175   AsyncUDPSocket& operator=(const AsyncUDPSocket&) = delete;
 176
 177   // EventHandler
 178   void handlerReady(uint16_t events) noexcept override;
 179
 180   void handleRead() noexcept;
 181   bool updateRegistration() noexcept;
 182
 183   EventBase* eventBase_;
 184   folly::SocketAddress localAddress_;
 185
 186   int fd_;
 187   FDOwnership ownership_;
 188
 189   // Temp space to receive client address
 190   folly::SocketAddress clientAddress_;
 191
 192   bool reuseAddr_{true};
 193   bool reusePort_{false};
 194 };
 195
 196 } // namespace folly