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 <glog/logging.h>
22
23#include <folly/Executor.h>
24#include <folly/synchronization/Baton.h>
25
26namespace folly {
27
28/// An Executor accepts units of work with add(), which should be
29/// threadsafe.
30class DefaultKeepAliveExecutor : public virtual Executor {
31 public:
32 virtual ~DefaultKeepAliveExecutor() {
33 DCHECK(!keepAlive_);
34 }
35
36 folly::Executor::KeepAlive<> weakRef() {
37 return WeakRef::create(controlBlock_, this);
38 }
39
40 protected:
41 void joinKeepAlive() {
42 DCHECK(keepAlive_);
43 keepAlive_.reset();
44 keepAliveReleaseBaton_.wait();
45 }
46
47 void joinAndResetKeepAlive() {
48 joinKeepAlive();
49 auto keepAliveCount =
50 controlBlock_->keepAliveCount_.exchange(1, std::memory_order_relaxed);
51 DCHECK_EQ(keepAliveCount, 0);
52 keepAliveReleaseBaton_.reset();
53 keepAlive_ = makeKeepAlive(this);
54 }
55
56 private:
57 struct ControlBlock {
58 std::atomic<ssize_t> keepAliveCount_{1};
59 };
60
61 class WeakRef : public Executor {
62 public:
63 static folly::Executor::KeepAlive<> create(
64 std::shared_ptr<ControlBlock> controlBlock,
65 Executor* executor) {
66 return makeKeepAlive(new WeakRef(std::move(controlBlock), executor));
67 }
68
69 void add(Func f) override {
70 if (auto executor = lock()) {
71 executor->add(std::move(f));
72 }
73 }
74
75 void addWithPriority(Func f, int8_t priority) override {
76 if (auto executor = lock()) {
77 executor->addWithPriority(std::move(f), priority);
78 }
79 }
80
81 virtual uint8_t getNumPriorities() const override {
82 return numPriorities_;
83 }
84
85 private:
86 WeakRef(std::shared_ptr<ControlBlock> controlBlock, Executor* executor)
87 : controlBlock_(std::move(controlBlock)),
88 executor_(executor),
89 numPriorities_(executor->getNumPriorities()) {}
90
91 bool keepAliveAcquire() override {
92 auto keepAliveCount =
93 keepAliveCount_.fetch_add(1, std::memory_order_relaxed);
94 // We should never increment from 0
95 DCHECK(keepAliveCount > 0);
96 return true;
97 }
98
99 void keepAliveRelease() override {
100 auto keepAliveCount =
101 keepAliveCount_.fetch_sub(1, std::memory_order_acq_rel);
102 DCHECK(keepAliveCount >= 1);
103
104 if (keepAliveCount == 1) {
105 delete this;
106 }
107 }
108
109 folly::Executor::KeepAlive<> lock() {
110 auto controlBlock =
111 controlBlock_->keepAliveCount_.load(std::memory_order_relaxed);
112 do {
113 if (controlBlock == 0) {
114 return {};
115 }
116 } while (!controlBlock_->keepAliveCount_.compare_exchange_weak(
117 controlBlock,
118 controlBlock + 1,
119 std::memory_order_release,
120 std::memory_order_relaxed));
121
122 return makeKeepAlive(executor_);
123 }
124
125 std::atomic<size_t> keepAliveCount_{1};
126
127 std::shared_ptr<ControlBlock> controlBlock_;
128 Executor* executor_;
129
130 uint8_t numPriorities_;
131 };
132
133 bool keepAliveAcquire() override {
134 auto keepAliveCount =
135 controlBlock_->keepAliveCount_.fetch_add(1, std::memory_order_relaxed);
136 // We should never increment from 0
137 DCHECK(keepAliveCount > 0);
138 return true;
139 }
140
141 void keepAliveRelease() override {
142 auto keepAliveCount =
143 controlBlock_->keepAliveCount_.fetch_sub(1, std::memory_order_acquire);
144 DCHECK(keepAliveCount >= 1);
145
146 if (keepAliveCount == 1) {
147 keepAliveReleaseBaton_.post(); // std::memory_order_release
148 }
149 }
150
151 std::shared_ptr<ControlBlock> controlBlock_{std::make_shared<ControlBlock>()};
152 Baton<> keepAliveReleaseBaton_;
153 KeepAlive<DefaultKeepAliveExecutor> keepAlive_{makeKeepAlive(this)};
154};
155
156} // namespace folly
157