查看本类的 API文档回源码主页即时通讯网 - 即时通讯开发者社区!
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             ids[count] = ids[insertionPosition];
184             ids[insertionPosition] = (short) id;
185             count++;
186         }
187 
188         /**
189          * Return how much more usable ids are left.
190          *
191          * @return the number of ids that are left for usage.
192          */
193         int usableIds() {
194             return count;
195         }
196 
197         /**
198          * Return the maximum number of ids that are supported.
199          *
200          * @return the maximum number of ids.
201          */
202         int maxUsableIds() {
203             return ids.length;
204         }
205     }
206 }