folly/io/async/AsyncPipe.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 #pragma once
  17
  18 #include <folly/io/async/AsyncTransport.h>
  19 #include <folly/io/async/EventHandler.h>
  20 #include <folly/io/async/DelayedDestruction.h>
  21 #include <folly/io/IOBufQueue.h>
  22
  23 #include <list>
  24 #include <system_error>
  25
  26 namespace folly {
  27
  28 class AsyncSocketException;
  29
  30 /**
  31  * Read from a pipe in an async manner.
  32  */
  33 class AsyncPipeReader : public EventHandler,
  34                         public AsyncReader,
  35                         public DelayedDestruction {
  36  public:
  37   typedef std::unique_ptr<AsyncPipeReader,
  38                           folly::DelayedDestruction::Destructor> UniquePtr;
  39
  40   template <typename... Args>
  41   static UniquePtr newReader(Args&&... args) {
  42     return UniquePtr(new AsyncPipeReader(std::forward<Args>(args)...));
  43   }
  44
  45   AsyncPipeReader(folly::EventBase* eventBase, int pipeFd)
  46     : EventHandler(eventBase, pipeFd),
  47     fd_(pipeFd) {}
  48
  49   /**
  50    * Set the read callback and automatically install/uninstall the handler
  51    * for events.
  52    */
  53   void setReadCB(AsyncReader::ReadCallback* callback) override {
  54     if (callback == readCallback_) {
  55       return;
  56     }
  57     readCallback_ = callback;
  58     if (readCallback_ && !isHandlerRegistered()) {
  59       registerHandler(EventHandler::READ | EventHandler::PERSIST);
  60     } else if (!readCallback_ && isHandlerRegistered()) {
  61       unregisterHandler();
  62     }
  63   }
  64
  65   /**
  66    * Get the read callback
  67    */
  68   AsyncReader::ReadCallback* getReadCallback() const override {
  69     return readCallback_;
  70   }
  71
  72   /**
  73    * Set a special hook to close the socket (otherwise, will call close())
  74    */
  75   void setCloseCallback(std::function<void(int)> closeCb) {
  76     closeCb_ = closeCb;
  77   }
  78
  79  private:
  80   ~AsyncPipeReader() override;
  81
  82   void handlerReady(uint16_t events) noexcept override;
  83   void failRead(const AsyncSocketException& ex);
  84   void close();
  85
  86   int fd_;
  87   AsyncReader::ReadCallback* readCallback_{nullptr};
  88   std::function<void(int)> closeCb_;
  89 };
  90
  91 /**
  92  * Write to a pipe in an async manner.
  93  */
  94 class AsyncPipeWriter : public EventHandler,
  95                         public AsyncWriter,
  96                         public DelayedDestruction {
  97  public:
  98   typedef std::unique_ptr<AsyncPipeWriter,
  99                           folly::DelayedDestruction::Destructor> UniquePtr;
 100
 101   template <typename... Args>
 102   static UniquePtr newWriter(Args&&... args) {
 103     return UniquePtr(new AsyncPipeWriter(std::forward<Args>(args)...));
 104   }
 105
 106   AsyncPipeWriter(folly::EventBase* eventBase, int pipeFd)
 107     : EventHandler(eventBase, pipeFd),
 108     fd_(pipeFd) {}
 109
 110   /**
 111    * Asynchronously write the given iobuf to this pipe, and invoke the callback
 112    * on success/error.
 113    */
 114   void write(std::unique_ptr<folly::IOBuf> iob,
 115              AsyncWriter::WriteCallback* wcb = nullptr);
 116
 117   /**
 118    * Set a special hook to close the socket (otherwise, will call close())
 119    */
 120   void setCloseCallback(std::function<void(int)> closeCb) {
 121     closeCb_ = closeCb;
 122   }
 123
 124   /**
 125    * Returns true if the pipe is closed
 126    */
 127   bool closed() const {
 128     return (fd_ < 0 || closeOnEmpty_);
 129   }
 130
 131   /**
 132    * Notify the pipe to close as soon as all pending writes complete
 133    */
 134   void closeOnEmpty();
 135
 136   /**
 137    * Close the pipe immediately, and fail all pending writes
 138    */
 139   void closeNow();
 140
 141   /**
 142    * Return true if there are currently writes pending (eg: the pipe is blocked
 143    * for writing)
 144    */
 145   bool hasPendingWrites() const {
 146     return !queue_.empty();
 147   }
 148
 149   // AsyncWriter methods
 150   void write(folly::AsyncWriter::WriteCallback* callback,
 151              const void* buf,
 152              size_t bytes,
 153              WriteFlags flags = WriteFlags::NONE) override {
 154     writeChain(callback, IOBuf::wrapBuffer(buf, bytes), flags);
 155   }
 156   void writev(folly::AsyncWriter::WriteCallback*,
 157               const iovec*,
 158               size_t,
 159               WriteFlags = WriteFlags::NONE) override {
 160     throw std::runtime_error("writev is not supported. Please use writeChain.");
 161   }
 162   void writeChain(folly::AsyncWriter::WriteCallback* callback,
 163                   std::unique_ptr<folly::IOBuf>&& buf,
 164                   WriteFlags flags = WriteFlags::NONE) override;
 165
 166  private:
 167   void handlerReady(uint16_t events) noexcept override;
 168   void handleWrite();
 169   void failAllWrites(const AsyncSocketException& ex);
 170
 171   int fd_;
 172   std::list<std::pair<folly::IOBufQueue, AsyncWriter::WriteCallback*>> queue_;
 173   bool closeOnEmpty_{false};
 174   std::function<void(int)> closeCb_;
 175
 176   ~AsyncPipeWriter() override {
 177     closeNow();
 178   }
 179 };
 180
 181 } // folly