1 /*
2  * Copyright (c) Facebook, Inc. and its affiliates.
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/Chrono.h>
22 #include <folly/futures/Future.h>
23 #include <folly/io/async/EventBase.h>
24 #include <folly/io/async/HHWheelTimer.h>
25 
26 namespace folly {
27 // Our Callback object for HHWheelTimer
28 template <class TBase>
29 struct WTCallback : public std::enable_shared_from_this<WTCallback<TBase>>,
30                     public TBase::Callback {
31   struct PrivateConstructorTag {};
32 
33  public:
WTCallbackWTCallback34   WTCallback(PrivateConstructorTag, EventBase* base) : base_(base) {}
35 
36   // Only allow creation by this factory, to ensure heap allocation.
createWTCallback37   static std::shared_ptr<WTCallback> create(EventBase* base) {
38     // optimization opportunity: memory pool
39     auto cob = std::make_shared<WTCallback>(PrivateConstructorTag{}, base);
40     // Capture shared_ptr of cob in lambda so that Core inside Promise will
41     // hold a ref count to it. The ref count will be released when Core goes
42     // away which happens when both Promise and Future go away
43     cob->promise_.setInterruptHandler(
44         [cob](exception_wrapper ew) { cob->interruptHandler(std::move(ew)); });
45     return cob;
46   }
47 
getSemiFutureWTCallback48   SemiFuture<Unit> getSemiFuture() { return promise_.getSemiFuture(); }
49 
stealPromiseWTCallback50   FOLLY_NODISCARD Promise<Unit> stealPromise() {
51     // Don't need promise anymore. Break the circular reference as promise_
52     // is holding a ref count to us via Core. Core won't go away until both
53     // Promise and Future go away.
54     return std::move(promise_);
55   }
56 
57  protected:
58   folly::Synchronized<EventBase*> base_;
59   Promise<Unit> promise_;
60 
timeoutExpiredWTCallback61   void timeoutExpired() noexcept override {
62     base_ = nullptr;
63     // Don't need Promise anymore, break the circular reference
64     auto promise = stealPromise();
65     if (!promise.isFulfilled()) {
66       promise.setValue();
67     }
68   }
69 
callbackCanceledWTCallback70   void callbackCanceled() noexcept override {
71     base_ = nullptr;
72     // Don't need Promise anymore, break the circular reference
73     auto promise = stealPromise();
74     if (!promise.isFulfilled()) {
75       promise.setException(FutureNoTimekeeper{});
76     }
77   }
78 
interruptHandlerWTCallback79   void interruptHandler(exception_wrapper ew) {
80     auto rBase = base_.rlock();
81     if (!*rBase) {
82       return;
83     }
84     // Capture shared_ptr of self in lambda, if we don't do this, object
85     // may go away before the lambda is executed from event base thread.
86     // This is not racing with timeoutExpired anymore because this is called
87     // through Future, which means Core is still alive and keeping a ref count
88     // on us, so what timeouExpired is doing won't make the object go away
89     (*rBase)->runInEventBaseThread([me = std::enable_shared_from_this<
90                                         WTCallback<TBase>>::shared_from_this(),
91                                     ew = std::move(ew)]() mutable {
92       me->cancelTimeout();
93       // Don't need Promise anymore, break the circular reference
94       auto promise = me->stealPromise();
95       if (!promise.isFulfilled()) {
96         promise.setException(std::move(ew));
97       }
98     });
99   }
100 };
101 
102 } // namespace folly
103