1 /*
2 * Copyright 2024 The Netty Project
3 *
4 * The Netty Project licenses this file to you under the Apache License,
5 * version 2.0 (the "License"); you may not use this file except in compliance
6 * with the License. You may obtain a copy of the License at:
7 *
8 * https://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, WITHOUT
12 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13 * License for the specific language governing permissions and limitations
14 * under the License.
15 */
16 package io.netty.resolver.dns;
17
18 import io.netty.util.internal.MathUtil;
19 import io.netty.util.internal.PlatformDependent;
20
21 import java.util.Random;
22
23 /**
24 * Special data-structure that will allow to retrieve the next query id to use, while still guarantee some sort
25 * of randomness.
26 * The query id will be between 0 (inclusive) and 65535 (inclusive) as defined by the RFC.
27 */
28 final class DnsQueryIdSpace {
29 private static final int MAX_ID = 65535;
30 private static final int BUCKETS = 4;
31 // Each bucket is 16kb of size.
32 private static final int BUCKET_SIZE = (MAX_ID + 1) / BUCKETS;
33
34 // If there are other buckets left that have at least 500 usable ids we will drop an unused bucket.
35 private static final int BUCKET_DROP_THRESHOLD = 500;
36 private final DnsQueryIdRange[] idBuckets = new DnsQueryIdRange[BUCKETS];
37
38 DnsQueryIdSpace() {
39 assert idBuckets.length == MathUtil.findNextPositivePowerOfTwo(idBuckets.length);
40 // We start with 1 bucket.
41 idBuckets[0] = newBucket(0);
42 }
43
44 private static DnsQueryIdRange newBucket(int idBucketsIdx) {
45 return new DnsQueryIdRange(BUCKET_SIZE, idBucketsIdx * BUCKET_SIZE);
46 }
47
48 /**
49 * Returns the next ID to use for a query or {@code -1} if there is none left to use.
50 *
51 * @return next id to use.
52 */
53 int nextId() {
54 int freeIdx = -1;
55 for (int bucketIdx = 0; bucketIdx < idBuckets.length; bucketIdx++) {
56 DnsQueryIdRange bucket = idBuckets[bucketIdx];
57 if (bucket != null) {
58 int id = bucket.nextId();
59 if (id != -1) {
60 return id;
61 }
62 } else if (freeIdx == -1 ||
63 // Let's make it somehow random which free slot is used.
64 PlatformDependent.threadLocalRandom().nextBoolean()) {
65 // We have a slot that we can use to create a new bucket if we need to.
66 freeIdx = bucketIdx;
67 }
68 }
69 if (freeIdx == -1) {
70 // No ids left and no slot left to create a new bucket.
71 return -1;
72 }
73
74 // We still have some slots free to store a new bucket. Let's do this now and use it to generate the next id.
75 DnsQueryIdRange bucket = newBucket(freeIdx);
76 idBuckets[freeIdx] = bucket;
77 int id = bucket.nextId();
78 assert id >= 0;
79 return id;
80 }
81
82 /**
83 * Push back the id, so it can be used again for the next query.
84 *
85 * @param id the id.
86 */
87 void pushId(int id) {
88 int bucketIdx = id / BUCKET_SIZE;
89 if (bucketIdx >= idBuckets.length) {
90 throw new IllegalArgumentException("id too large: " + id);
91 }
92 DnsQueryIdRange bucket = idBuckets[bucketIdx];
93 assert bucket != null;
94 bucket.pushId(id);
95
96 if (bucket.usableIds() == bucket.maxUsableIds()) {
97 // All ids are usable in this bucket. Let's check if there are other buckets left that have still
98 // some space left and if so drop this bucket.
99 for (int idx = 0; idx < idBuckets.length; idx++) {
100 if (idx != bucketIdx) {
101 DnsQueryIdRange otherBucket = idBuckets[idx];
102 if (otherBucket != null && otherBucket.usableIds() > BUCKET_DROP_THRESHOLD) {
103 // Drop bucket on the floor to reduce memory usage, there is another bucket left we can
104 // use that still has enough ids to use.
105 idBuckets[bucketIdx] = null;
106 return;
107 }
108 }
109 }
110 }
111 }
112
113 /**
114 * Return how much more usable ids are left.
115 *
116 * @return the number of ids that are left for usage.
117 */
118 int usableIds() {
119 int usableIds = 0;
120 for (DnsQueryIdRange bucket: idBuckets) {
121 // If there is nothing stored in the index yet we can assume the whole bucket is usable
122 usableIds += bucket == null ? BUCKET_SIZE : bucket.usableIds();
123 }
124 return usableIds;
125 }
126
127 /**
128 * Return the maximum number of ids that are supported.
129 *
130 * @return the maximum number of ids.
131 */
132 int maxUsableIds() {
133 return BUCKET_SIZE * idBuckets.length;
134 }
135
136 /**
137 * Provides a query if from a range of possible ids.
138 */
139 private static final class DnsQueryIdRange {
140
141 // Holds all possible ids which are stored as unsigned shorts
142 private final short[] ids;
143 private final int startId;
144 private int count;
145
146 DnsQueryIdRange(int bucketSize, int startId) {
147 this.ids = new short[bucketSize];
148 this.startId = startId;
149 for (int v = startId; v < bucketSize + startId; v++) {
150 pushId(v);
151 }
152 }
153
154 /**
155 * Returns the next ID to use for a query or {@code -1} if there is none left to use.
156 *
157 * @return next id to use.
158 */
159 int nextId() {
160 assert count >= 0;
161 if (count == 0) {
162 return -1;
163 }
164 short id = ids[count - 1];
165 count--;
166
167 return id & 0xFFFF;
168 }
169
170 /**
171 * Push back the id, so it can be used again for the next query.
172 *
173 * @param id the id.
174 */
175 void pushId(int id) {
176 if (count == ids.length) {
177 throw new IllegalStateException("overflow");
178 }
179 assert id <= startId + ids.length && id >= startId;
180 // pick a slot for our index, and whatever was in that slot before will get moved to the tail.
181 Random random = PlatformDependent.threadLocalRandom();
182 int insertionPosition = random.nextInt(count + 1);
183 short moveId = ids[insertionPosition];
184 short insertId = (short) id;
185
186 // Assert that the ids are different or its the same index.
187 assert moveId != insertId || insertionPosition == count;
188
189 ids[count] = moveId;
190 ids[insertionPosition] = insertId;
191 count++;
192 }
193
194 /**
195 * Return how much more usable ids are left.
196 *
197 * @return the number of ids that are left for usage.
198 */
199 int usableIds() {
200 return count;
201 }
202
203 /**
204 * Return the maximum number of ids that are supported.
205 *
206 * @return the maximum number of ids.
207 */
208 int maxUsableIds() {
209 return ids.length;
210 }
211 }
212 }