1 /* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2 /*
3  *     Copyright 2017 Couchbase, Inc.
4  *
5  *   Licensed under the Apache License, Version 2.0 (the "License");
6  *   you may not use this file except in compliance with the License.
7  *   You may obtain a copy of the License at
8  *
9  *       http://www.apache.org/licenses/LICENSE-2.0
10  *
11  *   Unless required by applicable law or agreed to in writing, software
12  *   distributed under the License is distributed on an "AS IS" BASIS,
13  *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *   See the License for the specific language governing permissions and
15  *   limitations under the License.
16  */
17 #pragma once
18 
19 #include <atomic>
20 #include <limits>
21 #include <stdexcept>
22 #include <string>
23 #include <type_traits>
24 
25 namespace cb {
26 
27 /// Policy class for handling underflow by clamping the value at zero.
28 template <class T>
29 struct ClampAtZeroUnderflowPolicy {
30     using SignedT = typename std::make_signed<T>::type;
31 
underflowcb::ClampAtZeroUnderflowPolicy32     void underflow(T& desired, T current, SignedT arg) {
33         desired = 0;
34     }
35 };
36 
37 /// Policy class for handling underflow by throwing an exception. Prints the
38 /// previous value stored in the counter and the argument (the value that we
39 // were attempting to subtract)
40 template <class T>
41 struct ThrowExceptionUnderflowPolicy {
42     using SignedT = typename std::make_signed<T>::type;
43 
underflowcb::ThrowExceptionUnderflowPolicy44     void underflow(T& desired, T current, SignedT arg) {
45         using std::to_string;
46         throw std::underflow_error("ThrowExceptionUnderflowPolicy current:" +
47                                    to_string(current) + " arg:" +
48                                    to_string(arg));
49     }
50 };
51 
52 // Default NonNegativeCounter OrdereReversedPolicy (if user doesn't explicitly
53 // specify otherwise) - use ClampAtZeroUnderflowPolicy for Release builds, and
54 // ThrowExceptionPolicy for Pre-Release builds.
55 template <class T>
56 #if CB_DEVELOPMENT_ASSERTS
57 using DefaultUnderflowPolicy = ThrowExceptionUnderflowPolicy<T>;
58 #else
59 using DefaultUnderflowPolicy = ClampAtZeroUnderflowPolicy<T>;
60 #endif
61 
62 /**
63  * The NonNegativeCounter class wraps std::atomic<> and prevents it
64  * underflowing over overflowing. By default will clamp the value at 0 on
65  * underflow, but behaviour can be customized by specifying a different
66  * UnderflowPolicy class. Same for overflow.
67  *
68  * Even though this counter can only be templated on unsigned types, it has the
69  * maximum value of the corresponding signed type. This is because we need to
70  * allow and check the addition of negative values.
71  */
72 template <typename T,
73           template <class> class UnderflowPolicy = DefaultUnderflowPolicy>
74 class NonNegativeCounter : public UnderflowPolicy<T> {
75     static_assert(
76             std::is_unsigned<T>::value,
77             "NonNegativeCounter should only be templated over unsigned types");
78 
79     using SignedT = typename std::make_signed<T>::type;
80 
81 public:
82     NonNegativeCounter() noexcept {
83         store(0);
84     }
85 
NonNegativeCounter(T initial)86     NonNegativeCounter(T initial) {
87         store(initial);
88     }
89 
90     NonNegativeCounter(const NonNegativeCounter& other) noexcept {
91         store(other.load());
92     }
93 
94     operator T() const noexcept {
95         return load();
96     }
97 
98     T load() const noexcept {
99         return value.load(std::memory_order_relaxed);
100     }
101 
store(T desired)102     void store(T desired) {
103         if (desired > T(std::numeric_limits<SignedT>::max())) {
104             UnderflowPolicy<T>::underflow(desired, load(), desired);
105         }
106         value.store(desired, std::memory_order_relaxed);
107     }
108 
109     /**
110      * Add 'arg' to the current value. If the new value would underflow (i.e. if
111      * arg was negative and current less than arg) then calls underflow() on the
112      * selected UnderflowPolicy.
113      *
114      * Note: Not marked 'noexcept' as underflow() could throw.
115      */
fetch_add(SignedT arg)116     T fetch_add(SignedT arg) {
117         T current = load();
118         T desired;
119         do {
120             if (arg < 0) {
121                 desired = current - T(std::abs(arg));
122                 if (SignedT(current) + arg < 0) {
123                     UnderflowPolicy<T>::underflow(desired, current, arg);
124                 }
125             } else {
126                 desired = current + T(arg);
127                 if (desired > T(std::numeric_limits<SignedT>::max())) {
128                     UnderflowPolicy<T>::underflow(desired, current, arg);
129                 }
130             }
131             // Attempt to set the atomic value to desired. If the atomic value
132             // is not the same as current then it has changed during
133             // operation. compare_exchange_weak will reload the new value
134             // into current if it fails, and we will retry.
135         } while (!value.compare_exchange_weak(
136                 current, desired, std::memory_order_relaxed));
137 
138         return current;
139     }
140 
141     /**
142      * Subtract 'arg' from the current value. If the new value would underflow
143      * then calls underflow() on the selected UnderflowPolicy.
144      *
145      * Note: Not marked 'noexcept' as underflow() could throw.
146      */
fetch_sub(SignedT arg)147     T fetch_sub(SignedT arg) {
148         T current = load();
149         T desired;
150         do {
151             if (arg < 0) {
152                 desired = current + T(std::abs(arg));
153             } else {
154                 desired = current - T(arg);
155             }
156 
157             if (desired > T(std::numeric_limits<SignedT>::max())) {
158                 UnderflowPolicy<T>::underflow(desired, current, arg);
159             }
160             // Attempt to set the atomic value to desired. If the atomic value
161             // is not the same as current then it has changed during
162             // operation. compare_exchange_weak will reload the new value
163             // into current if it fails, and we will retry.
164         } while (!value.compare_exchange_weak(
165                 current, desired, std::memory_order_relaxed));
166 
167         return current;
168     }
169 
170     T exchange(T arg) noexcept {
171         return value.exchange(arg, std::memory_order_relaxed);
172     }
173 
174     NonNegativeCounter& operator=(const NonNegativeCounter& rhs) noexcept {
175         value.store(rhs.load(), std::memory_order_relaxed);
176         return *this;
177     }
178 
operator +=(const T rhs)179     NonNegativeCounter& operator+=(const T rhs) {
180         fetch_add(rhs);
181         return *this;
182     }
183 
operator +=(const NonNegativeCounter& rhs)184     NonNegativeCounter& operator+=(const NonNegativeCounter& rhs) {
185         fetch_add(rhs.load());
186         return *this;
187     }
188 
operator -=(const T rhs)189     NonNegativeCounter& operator-=(const T rhs) {
190         fetch_sub(rhs);
191         return *this;
192     }
193 
operator -=(const NonNegativeCounter& rhs)194     NonNegativeCounter& operator-=(const NonNegativeCounter& rhs) {
195         fetch_sub(rhs.load());
196         return *this;
197     }
198 
operator ++()199     T operator++() {
200         return fetch_add(1) + 1;
201     }
202 
operator ++(int)203     T operator++(int) {
204         return fetch_add(1);
205     }
206 
operator --()207     T operator--() {
208         T previous = fetch_sub(1);
209         if (previous == 0) {
210             // If we are doing a clamp underflow we can pass in previous,
211             // it's already 0 and we are returning 0. If we are going to
212             // throw, we want to print previous.
213             UnderflowPolicy<T>::underflow(previous, previous, 1);
214             return 0;
215         }
216         return previous - 1;
217     }
218 
operator --(int)219     T operator--(int) {
220         return fetch_sub(1);
221     }
222 
operator =(T val)223     NonNegativeCounter& operator=(T val) {
224         store(val);
225         return *this;
226     }
227 
228 private:
229     std::atomic<T> value;
230 };
231 
232 } // namespace cb
233