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
25namespace cb {
26
27/// Policy class for handling underflow by clamping the value at zero.
28template <class T>
29struct ClampAtZeroUnderflowPolicy {
30    using SignedT = typename std::make_signed<T>::type;
31
32    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)
40template <class T>
41struct ThrowExceptionUnderflowPolicy {
42    using SignedT = typename std::make_signed<T>::type;
43
44    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.
55template <class T>
56#if CB_DEVELOPMENT_ASSERTS
57using DefaultUnderflowPolicy = ThrowExceptionUnderflowPolicy<T>;
58#else
59using 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 */
72template <typename T,
73          template <class> class UnderflowPolicy = DefaultUnderflowPolicy>
74class 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
81public:
82    NonNegativeCounter() noexcept {
83        store(0);
84    }
85
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
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     */
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     */
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
179    NonNegativeCounter& operator+=(const T rhs) {
180        fetch_add(rhs);
181        return *this;
182    }
183
184    NonNegativeCounter& operator+=(const NonNegativeCounter& rhs) {
185        fetch_add(rhs.load());
186        return *this;
187    }
188
189    NonNegativeCounter& operator-=(const T rhs) {
190        fetch_sub(rhs);
191        return *this;
192    }
193
194    NonNegativeCounter& operator-=(const NonNegativeCounter& rhs) {
195        fetch_sub(rhs.load());
196        return *this;
197    }
198
199    T operator++() {
200        return fetch_add(1) + 1;
201    }
202
203    T operator++(int) {
204        return fetch_add(1);
205    }
206
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
219    T operator--(int) {
220        return fetch_sub(1);
221    }
222
223    NonNegativeCounter& operator=(T val) {
224        store(val);
225        return *this;
226    }
227
228private:
229    std::atomic<T> value;
230};
231
232} // namespace cb
233