1 /*
2  * Copyright (C) 2022 The Android Open Source Project
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 <atomic>
20 #include <chrono>
21 #include <functional>
22 #include <thread>
23 
24 #include <gtest/gtest_prod.h>
25 
26 #include "android-base/macros.h"
27 
28 namespace android {
29 
30 /*
31  * AtomicState manages updating or waiting on a state enum between multiple threads.
32  */
33 template <typename T>
34 class AtomicState {
35  public:
AtomicState(T state)36   explicit AtomicState(T state) : state_(state) {}
37   ~AtomicState() = default;
38 
39   /*
40    * Set the state to `to`.  Wakes up any waiters that are waiting on the new state.
41    */
set(T to)42   void set(T to) { state_.store(to); }
43 
44   /*
45    * If the state is `from`, change it to `to` and return true.  Otherwise don't change
46    * it and return false.  If the state is changed, wakes up any waiters that are waiting
47    * on the new state.
48    */
transition(T from,T to)49   bool transition(T from, T to) { return state_.compare_exchange_strong(from, to); }
50 
51   /*
52    * If the state is `from`, change it to `to` and return true.  Otherwise, call `or_func`,
53    * set the state to the value it returns and return false.  Wakes up any waiters that are
54    * waiting on the new state or on the state set by or_func.
55    */
transition_or(T from,T to,const std::function<T ()> & orFunc)56   bool transition_or(T from, T to, const std::function<T()>& orFunc) {
57     if (!state_.compare_exchange_strong(from, to)) {
58       state_.store(orFunc());
59       return false;
60     }
61     return true;
62   }
63 
64   /*
65    * Block until the state is either `state1` or `state2`, or the time limit is reached.
66    * Busy loops with short sleeps.  Returns true if the time limit was not reached, false
67    * if it was reached.
68    */
wait_for_either_of(T state1,T state2,std::chrono::nanoseconds timeout)69   bool wait_for_either_of(T state1, T state2, std::chrono::nanoseconds timeout) {
70     using namespace std::chrono_literals;
71     auto end = std::chrono::steady_clock::now() + timeout;
72     const std::chrono::nanoseconds sleep_time = 500us;
73     while (true) {
74       T state = state_.load();
75       if (state == state1 || state == state2) {
76         return true;
77       }
78       auto now = std::chrono::steady_clock::now();
79       if (now > end) {
80         return false;
81       }
82       std::this_thread::sleep_for(std::min(sleep_time, end - now));
83     }
84   }
85 
86  private:
87   std::atomic<T> state_;
88 
89   FRIEND_TEST(AtomicStateTest, transition);
90   FRIEND_TEST(AtomicStateTest, wait);
91 
92   DISALLOW_COPY_AND_ASSIGN(AtomicState);
93 };
94 
95 }  // namespace android
96