1/*
2 * Copyright 2018-present 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 <memory>
20#include <thread>
21
22#include <folly/ScopeGuard.h>
23#include <folly/executors/thread_factory/ThreadFactory.h>
24
25namespace folly {
26
27class InitThreadFactory : public ThreadFactory {
28 public:
29 explicit InitThreadFactory(
30 std::shared_ptr<ThreadFactory> threadFactory,
31 Func&& threadInitializer,
32 Func&& threadFinializer = [] {})
33 : threadFactory_(std::move(threadFactory)),
34 threadInitFini_(std::make_shared<ThreadInitFini>(
35 std::move(threadInitializer),
36 std::move(threadFinializer))) {}
37
38 std::thread newThread(Func&& func) override {
39 return threadFactory_->newThread(
40 [func = std::move(func), threadInitFini = threadInitFini_]() mutable {
41 threadInitFini->initializer();
42 SCOPE_EXIT {
43 threadInitFini->finalizer();
44 };
45 func();
46 });
47 }
48
49 private:
50 std::shared_ptr<ThreadFactory> threadFactory_;
51 struct ThreadInitFini {
52 ThreadInitFini(Func&& init, Func&& fini)
53 : initializer(std::move(init)), finalizer(std::move(fini)) {}
54
55 Func initializer;
56 Func finalizer;
57 };
58 std::shared_ptr<ThreadInitFini> threadInitFini_;
59};
60
61} // namespace folly
62