folly/wangle/channel/OutputBufferingHandler.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 <folly/wangle/channel/Handler.h>
  20 #include <folly/io/async/EventBase.h>
  21 #include <folly/io/async/EventBaseManager.h>
  22 #include <folly/io/IOBuf.h>
  23 #include <folly/io/IOBufQueue.h>
  24
  25 namespace folly { namespace wangle {
  26
  27 /*
  28  * OutputBufferingHandler buffers writes in order to minimize syscalls. The
  29  * transport will be written to once per event loop instead of on every write.
  30  */
  31 class OutputBufferingHandler : public BytesToBytesHandler,
  32                                protected EventBase::LoopCallback {
  33  public:
  34   Future<void> write(Context* ctx, std::unique_ptr<IOBuf> buf) override {
  35     CHECK(buf);
  36     if (!queueSends_) {
  37       return ctx->fireWrite(std::move(buf));
  38     } else {
  39       ctx_ = ctx;
  40       // Delay sends to optimize for fewer syscalls
  41       if (!sends_) {
  42         DCHECK(!isLoopCallbackScheduled());
  43         // Buffer all the sends, and call writev once per event loop.
  44         sends_ = std::move(buf);
  45         ctx->getTransport()->getEventBase()->runInLoop(this);
  46       } else {
  47         DCHECK(isLoopCallbackScheduled());
  48         sends_->prependChain(std::move(buf));
  49       }
  50       Promise<void> p;
  51       auto f = p.getFuture();
  52       promises_.push_back(std::move(p));
  53       return f;
  54     }
  55   }
  56
  57   void runLoopCallback() noexcept override {
  58     MoveWrapper<std::vector<Promise<void>>> promises(std::move(promises_));
  59     ctx_->fireWrite(std::move(sends_)).then([promises](Try<void> t) mutable {
  60       for (auto& p : *promises) {
  61         p.setTry(t);
  62       }
  63     });
  64   }
  65
  66   Future<void> close(Context* ctx) override {
  67     if (isLoopCallbackScheduled()) {
  68       cancelLoopCallback();
  69     }
  70
  71     // If there are sends queued, cancel them
  72     for (auto& promise : promises_) {
  73       promise.setException(
  74         folly::make_exception_wrapper<std::runtime_error>(
  75           "close() called while sends still pending"));
  76     }
  77     sends_.reset();
  78     promises_.clear();
  79     return ctx->fireClose();
  80   }
  81
  82   std::vector<Promise<void>> promises_;
  83   std::unique_ptr<IOBuf> sends_{nullptr};
  84   bool queueSends_{true};
  85   Context* ctx_;
  86 };
  87
  88 }}