1 /*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
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 * A copy of the License is located at
7 *
8 * http://aws.amazon.com/apache2.0
9 *
10 * or in the "license" file accompanying this file. This file is distributed
11 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12 * express or implied. See the License for the specific language governing
13 * permissions and limitations under the License.
14 */
15
16 package software.amazon.awssdk.core.internal.capacity;
17
18 import java.util.Optional;
19 import java.util.concurrent.atomic.AtomicInteger;
20 import software.amazon.awssdk.annotations.SdkInternalApi;
21 import software.amazon.awssdk.core.retry.conditions.TokenBucketRetryCondition.Capacity;
22 import software.amazon.awssdk.utils.Validate;
23
24 /**
25 * A lock-free implementation of a token bucket. Tokens can be acquired from the bucket as long as there is sufficient capacity
26 * in the bucket.
27 */
28 @SdkInternalApi
29 public class TokenBucket {
30 private final int maxCapacity;
31 private final AtomicInteger capacity;
32
33 /**
34 * Create a bucket containing the specified number of tokens.
35 */
36 public TokenBucket(int maxCapacity) {
37 this.maxCapacity = maxCapacity;
38 this.capacity = new AtomicInteger(maxCapacity);
39 }
40
41 /**
42 * Try to acquire a certain number of tokens from this bucket. If there aren't sufficient tokens in this bucket,
43 * {@link Optional#empty()} is returned.
44 */
45 public Optional<Capacity> tryAcquire(int amountToAcquire) {
46 Validate.isTrue(amountToAcquire >= 0, "Amount must not be negative.");
47
48 if (amountToAcquire == 0) {
49 return Optional.of(Capacity.builder()
50 .capacityAcquired(0)
51 .capacityRemaining(capacity.get())
52 .build());
53 }
54
55 int currentCapacity;
56 int newCapacity;
57 do {
58 currentCapacity = capacity.get();
59 newCapacity = currentCapacity - amountToAcquire;
60
61 if (newCapacity < 0) {
62 return Optional.empty();
63 }
64 } while (!capacity.compareAndSet(currentCapacity, newCapacity));
65
66 return Optional.of(Capacity.builder()
67 .capacityAcquired(amountToAcquire)
68 .capacityRemaining(newCapacity)
69 .build());
70 }
71
72 /**
73 * Release a certain number of tokens back to this bucket. If this number of tokens would exceed the maximum number of tokens
74 * configured for the bucket, the bucket is instead set to the maximum value and the additional tokens are discarded.
75 */
76 public void release(int amountToRelease) {
77 Validate.isTrue(amountToRelease >= 0, "Amount must not be negative.");
78
79 if (amountToRelease == 0) {
80 return;
81 }
82
83 int currentCapacity;
84 int newCapacity;
85 do {
86 currentCapacity = capacity.get();
87
88 if (currentCapacity == maxCapacity) {
89 return;
90 }
91
92 newCapacity = Math.min(currentCapacity + amountToRelease, maxCapacity);
93 } while (!capacity.compareAndSet(currentCapacity, newCapacity));
94 }
95
96 /**
97 * Retrieve a snapshot of the current number of tokens in the bucket. Because this number is constantly changing, it's
98 * recommended to refer to the {@link Capacity#capacityRemaining()} returned by the {@link #tryAcquire(int)} method whenever
99 * possible.
100 */
101 public int currentCapacity() {
102 return capacity.get();
103 }
104
105 /**
106 * Retrieve the maximum capacity of the bucket configured when the bucket was created.
107 */
108 public int maxCapacity() {
109 return maxCapacity;
110 }
111
112 @Override
113 public boolean equals(Object o) {
114 if (this == o) {
115 return true;
116 }
117 if (o == null || getClass() != o.getClass()) {
118 return false;
119 }
120
121 TokenBucket that = (TokenBucket) o;
122
123 if (maxCapacity != that.maxCapacity) {
124 return false;
125 }
126 return capacity.get() == that.capacity.get();
127 }
128
129 @Override
130 public int hashCode() {
131 int result = maxCapacity;
132 result = 31 * result + capacity.get();
133 return result;
134 }
135 }
136