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(const T& initial) noexcept {
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) noexcept {
103        value.store(desired, std::memory_order_relaxed);
104    }
105
106    /**
107     * Add 'arg' to the current value. If the new value would underflow (i.e. if
108     * arg was negative and current less than arg) then calls underflow() on the
109     * selected UnderflowPolicy.
110     *
111     * Note: Not marked 'noexcept' as underflow() could throw.
112     */
113    T fetch_add(SignedT arg) {
114        T current = load();
115        T desired;
116        do {
117            if (arg < 0) {
118                desired = current - T(std::abs(arg));
119                if (SignedT(current) + arg < 0) {
120                    UnderflowPolicy<T>::underflow(desired, current, arg);
121                }
122            } else {
123                desired = current + T(arg);
124                if (desired > std::numeric_limits<SignedT>::max()) {
125                    UnderflowPolicy<T>::underflow(desired, current, arg);
126                }
127            }
128            // Attempt to set the atomic value to desired. If the atomic value
129            // is not the same as current then it has changed during
130            // operation. compare_exchange_weak will reload the new value
131            // into current if it fails, and we will retry.
132        } while (!value.compare_exchange_weak(
133                current, desired, std::memory_order_relaxed));
134
135        return current;
136    }
137
138    /**
139     * Subtract 'arg' from the current value. If the new value would underflow
140     * then calls underflow() on the selected UnderflowPolicy.
141     *
142     * Note: Not marked 'noexcept' as underflow() could throw.
143     */
144    T fetch_sub(SignedT arg) {
145        T current = load();
146        T desired;
147        do {
148            if (arg < 0) {
149                desired = current + T(std::abs(arg));
150            } else {
151                desired = current - T(arg);
152            }
153
154            if (desired > std::numeric_limits<SignedT>::max()) {
155                UnderflowPolicy<T>::underflow(desired, current, arg);
156            }
157            // Attempt to set the atomic value to desired. If the atomic value
158            // is not the same as current then it has changed during
159            // operation. compare_exchange_weak will reload the new value
160            // into current if it fails, and we will retry.
161        } while (!value.compare_exchange_weak(
162                current, desired, std::memory_order_relaxed));
163
164        return current;
165    }
166
167    T exchange(T arg) noexcept {
168        return value.exchange(arg, std::memory_order_relaxed);
169    }
170
171    NonNegativeCounter& operator=(const NonNegativeCounter& rhs) noexcept {
172        store(rhs.load());
173        return *this;
174    }
175
176    NonNegativeCounter& operator+=(const T rhs) {
177        fetch_add(rhs);
178        return *this;
179    }
180
181    NonNegativeCounter& operator+=(const NonNegativeCounter& rhs) {
182        fetch_add(rhs.load());
183        return *this;
184    }
185
186    NonNegativeCounter& operator-=(const T rhs) {
187        fetch_sub(rhs);
188        return *this;
189    }
190
191    NonNegativeCounter& operator-=(const NonNegativeCounter& rhs) {
192        fetch_sub(rhs.load());
193        return *this;
194    }
195
196    T operator++() {
197        return fetch_add(1) + 1;
198    }
199
200    T operator++(int) {
201        return fetch_add(1);
202    }
203
204    T operator--() {
205        T previous = fetch_sub(1);
206        if (previous == 0) {
207            // If we are doing a clamp underflow we can pass in previous,
208            // it's already 0 and we are returning 0. If we are going to
209            // throw, we want to print previous.
210            UnderflowPolicy<T>::underflow(previous, previous, 1);
211            return 0;
212        }
213        return previous - 1;
214    }
215
216    T operator--(int) {
217        return fetch_sub(1);
218    }
219
220    NonNegativeCounter& operator=(T val) noexcept {
221        store(val);
222        return *this;
223    }
224
225private:
226    std::atomic<T> value;
227};
228
229} // namespace cb
230