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