folly/io/async/AsyncUDPSocket.h

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