Handler::getContext() when possible

Summary:
Only allow this if the handler is only ever attached to a single pipeline once. i.e. only ever associated with one Context

Test Plan: unit, thrift unit

Reviewed By: davejwatson@fb.com

Subscribers: fugalh, alandau, bmatheny, folly-diffs@, jsedgwick, yfeldblum, chalfant

FB internal diff: D2024007

Tasks: 6836580

Signature: t1:2024007:1430157264:efcf70ca3531c10eec5d458c9e9d6cda60c507c3

folly/wangle/channel/AsyncSocketHandler.h

@@ -25,6 +25,7 @@
 
 namespace folly { namespace wangle {
 
+// This handler may only be used in a single Pipeline
 class AsyncSocketHandler
   : public folly::wangle::BytesToBytesHandler,
     public AsyncSocket::ReadCallback {
@@ -65,8 +66,6 @@ class AsyncSocketHandler
   }
 
   void attachPipeline(Context* ctx) override {
-    CHECK(!ctx_);
-    ctx_ = ctx;
     attachReadCallback();
   }
 
@@ -105,7 +104,7 @@ class AsyncSocketHandler
   }
 
   void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
-    const auto readBufferSettings = ctx_->getReadBufferSettings();
+    const auto readBufferSettings = getContext()->getReadBufferSettings();
     const auto ret = bufQueue_.preallocate(
         readBufferSettings.first,
         readBufferSettings.second);
@@ -115,16 +114,17 @@ class AsyncSocketHandler
 
   void readDataAvailable(size_t len) noexcept override {
     bufQueue_.postallocate(len);
-    ctx_->fireRead(bufQueue_);
+    getContext()->fireRead(bufQueue_);
   }
 
   void readEOF() noexcept override {
-    ctx_->fireReadEOF();
+    getContext()->fireReadEOF();
   }
 
   void readErr(const AsyncSocketException& ex)
     noexcept override {
-    ctx_->fireReadException(make_exception_wrapper<AsyncSocketException>(ex));
+    getContext()->fireReadException(
+        make_exception_wrapper<AsyncSocketException>(ex));
   }
 
  private:
@@ -146,7 +146,6 @@ class AsyncSocketHandler
     folly::Promise<void> promise_;
   };
 
-  Context* ctx_{nullptr};
   folly::IOBufQueue bufQueue_{folly::IOBufQueue::cacheChainLength()};
   std::shared_ptr<AsyncSocket> socket_{nullptr};
 };

folly/wangle/channel/Handler.h

@@ -23,8 +23,33 @@
 
 namespace folly { namespace wangle {
 
+template <class Context>
+class HandlerBase {
+ public:
+  virtual ~HandlerBase() {}
+
+  virtual void attachPipeline(Context* ctx) {}
+  virtual void attachTransport(Context* ctx) {}
+
+  virtual void detachPipeline(Context* ctx) {}
+  virtual void detachTransport(Context* ctx) {}
+
+  Context* getContext() {
+    if (attachCount_ != 1) {
+      return nullptr;
+    }
+    CHECK(ctx_);
+    return ctx_;
+  }
+
+ private:
+  friend detail::HandlerContextBase<Context>;
+  uint64_t attachCount_{0};
+  Context* ctx_{nullptr};
+};
+
 template <class Rin, class Rout = Rin, class Win = Rout, class Wout = Rin>
-class Handler {
+class Handler : public HandlerBase<HandlerContext<Rout, Wout>> {
  public:
   typedef Rin rin;
   typedef Rout rout;
@@ -46,12 +71,6 @@ class Handler {
     return ctx->fireClose();
   }
 
-  virtual void attachPipeline(Context* ctx) {}
-  virtual void attachTransport(Context* ctx) {}
-
-  virtual void detachPipeline(Context* ctx) {}
-  virtual void detachTransport(Context* ctx) {}
-
   /*
   // Other sorts of things we might want, all shamelessly stolen from Netty
   // inbound

folly/wangle/channel/HandlerContext.h

@@ -22,8 +22,26 @@
 
 namespace folly { namespace wangle {
 
+namespace detail {
+
+template <class HandlerContext>
+class HandlerContextBase {
+ protected:
+  template <class H>
+  void attachContext(H* handler, HandlerContext* ctx) {
+    if (++handler->attachCount_ == 1) {
+      handler->ctx_ = ctx;
+    } else {
+      handler->ctx_ = nullptr;
+    }
+  }
+};
+
+} // detail
+
 template <class In, class Out>
-class HandlerContext {
+class HandlerContext
+  : public detail::HandlerContextBase<HandlerContext<In, Out>> {
  public:
   virtual ~HandlerContext() {}
 
@@ -125,6 +143,7 @@ class ContextImpl : public HandlerContext<typename H::rout,
   // PipelineContext overrides
   void attachPipeline() override {
     if (!attached_) {
+      this->attachContext(handler_.get(), this);
       handler_->attachPipeline(this);
       attached_ = true;
     }

folly/wangle/channel/OutputBufferingHandler.h

@@ -27,6 +27,8 @@ namespace folly { namespace wangle {
 /*
  * OutputBufferingHandler buffers writes in order to minimize syscalls. The
  * transport will be written to once per event loop instead of on every write.
+ *
+ * This handler may only be used in a single Pipeline.
  */
 class OutputBufferingHandler : public BytesToBytesHandler,
                                protected EventBase::LoopCallback {
@@ -36,7 +38,6 @@ class OutputBufferingHandler : public BytesToBytesHandler,
     if (!queueSends_) {
       return ctx->fireWrite(std::move(buf));
     } else {
-      ctx_ = ctx;
       // Delay sends to optimize for fewer syscalls
       if (!sends_) {
         DCHECK(!isLoopCallbackScheduled());
@@ -56,11 +57,12 @@ class OutputBufferingHandler : public BytesToBytesHandler,
 
   void runLoopCallback() noexcept override {
     MoveWrapper<std::vector<Promise<void>>> promises(std::move(promises_));
-    ctx_->fireWrite(std::move(sends_)).then([promises](Try<void> t) mutable {
-      for (auto& p : *promises) {
-        p.setTry(t);
-      }
-    });
+    getContext()->fireWrite(std::move(sends_))
+      .then([promises](Try<void> t) mutable {
+        for (auto& p : *promises) {
+          p.setTry(t);
+        }
+      });
   }
 
   Future<void> close(Context* ctx) override {
@@ -82,7 +84,6 @@ class OutputBufferingHandler : public BytesToBytesHandler,
   std::vector<Promise<void>> promises_;
   std::unique_ptr<IOBuf> sends_{nullptr};
   bool queueSends_{true};
-  Context* ctx_;
 };
 
 }}

folly/wangle/channel/test/PipelineTest.cpp

@@ -225,11 +225,20 @@ TEST(PipelineTest, DynamicAttachDetachOrder) {
   }
 }
 
+TEST(PipelineTest, GetContext) {
+  IntHandler handler;
+  EXPECT_CALL(handler, attachPipeline(_));
+  StaticPipeline<int, int, IntHandler> pipeline(&handler);
+  EXPECT_TRUE(handler.getContext());
+  EXPECT_CALL(handler, detachPipeline(_));
+}
+
 TEST(PipelineTest, HandlerInMultiplePipelines) {
   IntHandler handler;
   EXPECT_CALL(handler, attachPipeline(_)).Times(2);
   StaticPipeline<int, int, IntHandler> pipeline1(&handler);
   StaticPipeline<int, int, IntHandler> pipeline2(&handler);
+  EXPECT_FALSE(handler.getContext());
   EXPECT_CALL(handler, detachPipeline(_)).Times(2);
 }
 
@@ -237,6 +246,7 @@ TEST(PipelineTest, HandlerInPipelineTwice) {
   IntHandler handler;
   EXPECT_CALL(handler, attachPipeline(_)).Times(2);
   StaticPipeline<int, int, IntHandler, IntHandler> pipeline(&handler, &handler);
+  EXPECT_FALSE(handler.getContext());
   EXPECT_CALL(handler, detachPipeline(_)).Times(2);
 }