Add keepAlive() mechanism
[folly.git] / folly / io / async / VirtualEventBase.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
17 #pragma once
18
19 #include <folly/Baton.h>
20 #include <folly/Executor.h>
21 #include <folly/io/async/EventBase.h>
22
23 namespace folly {
24
25 /**
26  * VirtualEventBase implements a light-weight view onto existing EventBase.
27  *
28  * Multiple VirtualEventBases can be backed by a single EventBase. Similarly
29  * to EventBase, VirtualEventBase implements loopKeepAlive() functionality,
30  * which allows callbacks holding KeepAlive token to keep EventBase looping
31  * until they are complete.
32  *
33  * VirtualEventBase destructor blocks until all its KeepAliveTokens are released
34  * and all tasks scheduled through it are complete. EventBase destructor also
35  * blocks until all VirtualEventBases backed by it are released.
36  */
37 class VirtualEventBase : public folly::Executor, public folly::TimeoutManager {
38  public:
39   explicit VirtualEventBase(EventBase& evb);
40
41   VirtualEventBase(const VirtualEventBase&) = delete;
42   VirtualEventBase& operator=(const VirtualEventBase&) = delete;
43
44   ~VirtualEventBase();
45
46   EventBase& getEventBase() {
47     return evb_;
48   }
49
50   /**
51    * Adds the given callback to a queue of things run before destruction
52    * of current VirtualEventBase.
53    *
54    * This allows users of VirtualEventBase that run in it, but don't control it,
55    * to be notified before VirtualEventBase gets destructed.
56    *
57    * Note: this will be called from the loop of the EventBase, backing this
58    * VirtualEventBase
59    */
60   void runOnDestruction(EventBase::LoopCallback* callback);
61
62   /**
63    * @see EventBase::runInLoop
64    */
65   template <typename F>
66   void runInLoop(F&& f, bool thisIteration = false) {
67     evb_.runInLoop(std::forward<F>(f), thisIteration);
68   }
69
70   /**
71    * VirtualEventBase destructor blocks until all tasks scheduled through its
72    * runInEventBaseThread are complete.
73    *
74    * @see EventBase::runInEventBaseThread
75    */
76   template <typename F>
77   void runInEventBaseThread(F&& f) {
78     // KeepAlive token has to be released in the EventBase thread. If
79     // runInEventBaseThread() fails, we can't extract the KeepAlive token
80     // from the callback to properly release it.
81     CHECK(evb_.runInEventBaseThread([
82       keepAliveToken = getKeepAliveToken(),
83       f = std::forward<F>(f)
84     ]() mutable { f(); }));
85   }
86
87   HHWheelTimer& timer() {
88     return evb_.timer();
89   }
90
91   void attachTimeoutManager(
92       AsyncTimeout* obj,
93       TimeoutManager::InternalEnum internal) override {
94     evb_.attachTimeoutManager(obj, internal);
95   }
96
97   void detachTimeoutManager(AsyncTimeout* obj) override {
98     evb_.detachTimeoutManager(obj);
99   }
100
101   bool scheduleTimeout(AsyncTimeout* obj, TimeoutManager::timeout_type timeout)
102       override {
103     return evb_.scheduleTimeout(obj, timeout);
104   }
105
106   void cancelTimeout(AsyncTimeout* obj) override {
107     evb_.cancelTimeout(obj);
108   }
109
110   void bumpHandlingTime() override {
111     evb_.bumpHandlingTime();
112   }
113
114   bool isInTimeoutManagerThread() override {
115     return evb_.isInTimeoutManagerThread();
116   }
117
118   /**
119    * @see runInEventBaseThread
120    */
121   void add(folly::Func f) override {
122     runInEventBaseThread(std::move(f));
123   }
124
125   /**
126    * Returns you a handle which prevents VirtualEventBase from being destroyed.
127    * KeepAlive handle can be released from EventBase loop only.
128    */
129   KeepAlive getKeepAliveToken() override {
130     if (evb_.inRunningEventBaseThread()) {
131       ++loopKeepAliveCount_;
132     } else {
133       ++loopKeepAliveCountAtomic_;
134     }
135     return makeKeepAlive();
136   }
137
138   bool inRunningEventBaseThread() const {
139     return evb_.inRunningEventBaseThread();
140   }
141
142  protected:
143   void keepAliveRelease() override {
144     DCHECK(getEventBase().inRunningEventBaseThread());
145     if (loopKeepAliveCountAtomic_.load()) {
146       loopKeepAliveCount_ += loopKeepAliveCountAtomic_.exchange(0);
147     }
148     DCHECK(loopKeepAliveCount_ > 0);
149     if (--loopKeepAliveCount_ == 0) {
150       loopKeepAliveBaton_.post();
151     }
152   }
153
154  private:
155   using LoopCallbackList = EventBase::LoopCallback::List;
156
157   EventBase& evb_;
158
159   ssize_t loopKeepAliveCount_{0};
160   std::atomic<ssize_t> loopKeepAliveCountAtomic_{0};
161   folly::Baton<> loopKeepAliveBaton_;
162   KeepAlive loopKeepAlive_;
163
164   KeepAlive evbLoopKeepAlive_;
165
166   folly::Synchronized<LoopCallbackList> onDestructionCallbacks_;
167 };
168 }