1 //
2 // Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2021
3 //
4 // Distributed under the Boost Software License, Version 1.0. (See accompanying
5 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
6 //
7 #pragma once
8 
9 #include "td/actor/actor.h"
10 
11 #include "td/utils/common.h"
12 #include "td/utils/optional.h"
13 
14 #include <functional>
15 
16 namespace td {
17 
18 template <class T>
19 class SchedulerLocalStorage {
20  public:
SchedulerLocalStorage()21   SchedulerLocalStorage() : data_(Scheduler::instance()->sched_count()) {
22   }
get()23   T &get() {
24     return data_[Scheduler::instance()->sched_id()];
25   }
26   template <class F>
for_each(F && f)27   void for_each(F &&f) {
28     for (auto &value : data_) {
29       f(value);
30     }
31   }
32   template <class F>
for_each(F && f)33   void for_each(F &&f) const {
34     for (const auto &value : data_) {
35       f(value);
36     }
37   }
38 
39  private:
40   std::vector<T> data_;
41 };
42 
43 template <class T>
44 class LazySchedulerLocalStorage {
45  public:
46   LazySchedulerLocalStorage() = default;
LazySchedulerLocalStorage(std::function<T ()> create_func)47   explicit LazySchedulerLocalStorage(std::function<T()> create_func) : create_func_(std::move(create_func)) {
48   }
set_create_func(std::function<T ()> create_func)49   void set_create_func(std::function<T()> create_func) {
50     CHECK(!create_func_);
51     create_func_ = create_func;
52   }
53 
set(T && t)54   void set(T &&t) {
55     auto &optional_value_ = sls_optional_value_.get();
56     CHECK(!optional_value_);
57     optional_value_ = std::move(t);
58   }
59 
get()60   T &get() {
61     auto &optional_value_ = sls_optional_value_.get();
62     if (!optional_value_) {
63       CHECK(create_func_);
64       optional_value_ = create_func_();
65     }
66     return *optional_value_;
67   }
clear_values()68   void clear_values() {
69     sls_optional_value_.for_each([&](auto &optional_value) {
70       if (optional_value) {
71         optional_value = optional<T>();
72       }
73     });
74   }
75 
76  private:
77   std::function<T()> create_func_;
78   SchedulerLocalStorage<optional<T>> sls_optional_value_;
79 };
80 
81 }  // namespace td
82