View Javadoc

1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements.  See the NOTICE file distributed with
4    * this work for additional information regarding copyright ownership.
5    * The ASF licenses this file to you under the Apache License, Version 2.0
6    * (the "License"); you may not use this file except in compliance with
7    * the License.  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  package org.apache.hadoop.hbase.quotas;
18  
19  import static org.junit.Assert.fail;
20  
21  import java.io.IOException;
22  import java.util.ArrayList;
23  import java.util.List;
24  import java.util.Objects;
25  import java.util.Random;
26  import java.util.Set;
27  import java.util.Map.Entry;
28  import java.util.concurrent.atomic.AtomicLong;
29  
30  import org.apache.commons.logging.Log;
31  import org.apache.commons.logging.LogFactory;
32  import org.apache.hadoop.hbase.HBaseTestingUtility;
33  import org.apache.hadoop.hbase.HColumnDescriptor;
34  import org.apache.hadoop.hbase.HTableDescriptor;
35  import org.apache.hadoop.hbase.NamespaceDescriptor;
36  import org.apache.hadoop.hbase.TableName;
37  import org.apache.hadoop.hbase.Waiter.Predicate;
38  import org.apache.hadoop.hbase.classification.InterfaceAudience;
39  import org.apache.hadoop.hbase.client.Admin;
40  import org.apache.hadoop.hbase.client.Connection;
41  import org.apache.hadoop.hbase.client.Put;
42  import org.apache.hadoop.hbase.client.Table;
43  import org.apache.hadoop.hbase.util.Bytes;
44  import org.junit.rules.TestName;
45  
46  import com.google.common.collect.HashMultimap;
47  import com.google.common.collect.Iterables;
48  import com.google.common.collect.Multimap;
49  
50  @InterfaceAudience.Private
51  public class SpaceQuotaHelperForTests {
52    private static final Log LOG = LogFactory.getLog(SpaceQuotaHelperForTests.class);
53  
54    public static final int SIZE_PER_VALUE = 256;
55    public static final String F1 = "f1";
56    public static final long ONE_KILOBYTE = 1024L;
57    public static final long ONE_MEGABYTE = ONE_KILOBYTE * ONE_KILOBYTE;
58  
59    private final HBaseTestingUtility testUtil;
60    private final TestName testName;
61    private final AtomicLong counter;
62  
63    public SpaceQuotaHelperForTests(
64        HBaseTestingUtility testUtil, TestName testName, AtomicLong counter) {
65      this.testUtil = Objects.requireNonNull(testUtil);
66      this.testName = Objects.requireNonNull(testName);
67      this.counter = Objects.requireNonNull(counter);
68    }
69  
70    //
71    // Helpers
72    //
73  
74    /**
75     * Returns the number of quotas defined in the HBase quota table.
76     */
77    long listNumDefinedQuotas(Connection conn) throws IOException {
78      QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
79      try {
80        return Iterables.size(scanner);
81      } finally {
82        if (scanner != null) {
83          scanner.close();
84        }
85      }
86    }
87  
88    /**
89     * Removes all quotas defined in the HBase quota table.
90     */
91    void removeAllQuotas(Connection conn) throws IOException {
92      QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
93      try {
94        for (QuotaSettings quotaSettings : scanner) {
95          final String namespace = quotaSettings.getNamespace();
96          final TableName tableName = quotaSettings.getTableName();
97          if (namespace != null) {
98            LOG.debug("Deleting quota for namespace: " + namespace);
99            QuotaUtil.deleteNamespaceQuota(conn, namespace);
100         } else {
101           assert tableName != null;
102           LOG.debug("Deleting quota for table: "+ tableName);
103           QuotaUtil.deleteTableQuota(conn, tableName);
104         }
105       }
106     } finally {
107       if (scanner != null) {
108         scanner.close();
109       }
110     }
111   }
112 
113   QuotaSettings getTableSpaceQuota(Connection conn, TableName tn) throws IOException {
114     try (QuotaRetriever scanner = QuotaRetriever.open(
115         conn.getConfiguration(), new QuotaFilter().setTableFilter(tn.getNameAsString()))) {
116       for (QuotaSettings setting : scanner) {
117         if (setting.getTableName().equals(tn) && setting.getQuotaType() == QuotaType.SPACE) {
118           return setting;
119         }
120       }
121       return null;
122     }
123   }
124 
125   /**
126    * Waits 30seconds for the HBase quota table to exist.
127    */
128   void waitForQuotaTable(Connection conn) throws IOException {
129     waitForQuotaTable(conn, 30_000);
130   }
131 
132   /**
133    * Waits {@code timeout} milliseconds for the HBase quota table to exist.
134    */
135   void waitForQuotaTable(final Connection conn, long timeout) throws IOException {
136     testUtil.waitFor(timeout, 1000, new Predicate<IOException>() {
137       @Override
138       public boolean evaluate() throws IOException {
139         return conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME);
140       }
141     });
142   }
143 
144   void writeData(TableName tn, long sizeInBytes) throws IOException {
145     writeData(testUtil.getConnection(), tn, sizeInBytes);
146   }
147 
148   void writeData(Connection conn, TableName tn, long sizeInBytes) throws IOException {
149     final Table table = conn.getTable(tn);
150     try {
151       List<Put> updates = new ArrayList<>();
152       long bytesToWrite = sizeInBytes;
153       long rowKeyId = 0L;
154       final StringBuilder sb = new StringBuilder();
155       final Random r = new Random();
156       while (bytesToWrite > 0L) {
157         sb.setLength(0);
158         sb.append(Long.toString(rowKeyId));
159         // Use the reverse counter as the rowKey to get even spread across all regions
160         Put p = new Put(Bytes.toBytes(sb.reverse().toString()));
161         byte[] value = new byte[SIZE_PER_VALUE];
162         r.nextBytes(value);
163         p.addColumn(Bytes.toBytes(F1), Bytes.toBytes("q1"), value);
164         updates.add(p);
165 
166         // Batch 50K worth of updates
167         if (updates.size() > 50) {
168           table.put(updates);
169           updates.clear();
170         }
171 
172         // Just count the value size, ignore the size of rowkey + column
173         bytesToWrite -= SIZE_PER_VALUE;
174         rowKeyId++;
175       }
176 
177       // Write the final batch
178       if (!updates.isEmpty()) {
179         table.put(updates);
180       }
181 
182       LOG.debug("Data was written to HBase");
183       // Push the data to disk.
184       testUtil.getHBaseAdmin().flush(tn);
185       LOG.debug("Data flushed to disk");
186     } finally {
187       table.close();
188     }
189   }
190 
191   Multimap<TableName, QuotaSettings> createTablesWithSpaceQuotas() throws Exception {
192     final Admin admin = testUtil.getHBaseAdmin();
193     final Multimap<TableName, QuotaSettings> tablesWithQuotas = HashMultimap.create();
194 
195     final TableName tn1 = createTable();
196     final TableName tn2 = createTable();
197 
198     NamespaceDescriptor nd = NamespaceDescriptor.create("ns" + counter.getAndIncrement()).build();
199     admin.createNamespace(nd);
200     final TableName tn3 = createTableInNamespace(nd);
201     final TableName tn4 = createTableInNamespace(nd);
202     final TableName tn5 = createTableInNamespace(nd);
203 
204     final long sizeLimit1 = 1024L * 1024L * 1024L * 1024L * 5L; // 5TB
205     final SpaceViolationPolicy violationPolicy1 = SpaceViolationPolicy.NO_WRITES;
206     QuotaSettings qs1 = QuotaSettingsFactory.limitTableSpace(tn1, sizeLimit1, violationPolicy1);
207     tablesWithQuotas.put(tn1, qs1);
208     admin.setQuota(qs1);
209 
210     final long sizeLimit2 = 1024L * 1024L * 1024L * 200L; // 200GB
211     final SpaceViolationPolicy violationPolicy2 = SpaceViolationPolicy.NO_WRITES_COMPACTIONS;
212     QuotaSettings qs2 = QuotaSettingsFactory.limitTableSpace(tn2, sizeLimit2, violationPolicy2);
213     tablesWithQuotas.put(tn2, qs2);
214     admin.setQuota(qs2);
215 
216     final long sizeLimit3 = 1024L * 1024L * 1024L * 1024L * 100L; // 100TB
217     final SpaceViolationPolicy violationPolicy3 = SpaceViolationPolicy.NO_INSERTS;
218     QuotaSettings qs3 = QuotaSettingsFactory.limitNamespaceSpace(
219         nd.getName(), sizeLimit3, violationPolicy3);
220     tablesWithQuotas.put(tn3, qs3);
221     tablesWithQuotas.put(tn4, qs3);
222     tablesWithQuotas.put(tn5, qs3);
223     admin.setQuota(qs3);
224 
225     final long sizeLimit4 = 1024L * 1024L * 1024L * 5L; // 5GB
226     final SpaceViolationPolicy violationPolicy4 = SpaceViolationPolicy.NO_INSERTS;
227     QuotaSettings qs4 = QuotaSettingsFactory.limitTableSpace(tn5, sizeLimit4, violationPolicy4);
228     // Override the ns quota for tn5, import edge-case to catch table quota taking
229     // precedence over ns quota.
230     tablesWithQuotas.put(tn5, qs4);
231     admin.setQuota(qs4);
232 
233     return tablesWithQuotas;
234   }
235 
236   TableName createTable() throws Exception {
237     return createTableWithRegions(1);
238   }
239 
240   TableName createTableWithRegions(int numRegions) throws Exception {
241     return createTableWithRegions(NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
242   }
243 
244   TableName createTableWithRegions(Admin admin, int numRegions) throws Exception {
245     return createTableWithRegions(
246         testUtil.getHBaseAdmin(), NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
247   }
248 
249   TableName createTableWithRegions(String namespace, int numRegions) throws Exception {
250     return createTableWithRegions(testUtil.getHBaseAdmin(), namespace, numRegions);
251   }
252 
253   TableName createTableWithRegions(Admin admin, String namespace, int numRegions) throws Exception {
254     final TableName tn = TableName.valueOf(
255         namespace, testName.getMethodName() + counter.getAndIncrement());
256 
257     // Delete the old table
258     if (admin.tableExists(tn)) {
259       admin.disableTable(tn);
260       admin.deleteTable(tn);
261     }
262 
263     // Create the table
264     HTableDescriptor tableDesc = new HTableDescriptor(tn);
265     tableDesc.addFamily(new HColumnDescriptor(F1));
266     if (numRegions == 1) {
267       admin.createTable(tableDesc);
268     } else {
269       admin.createTable(tableDesc, Bytes.toBytes("0"), Bytes.toBytes("9"), numRegions);
270     }
271     return tn;
272   }
273 
274   TableName createTableInNamespace(NamespaceDescriptor nd) throws Exception {
275     final Admin admin = testUtil.getHBaseAdmin();
276     final TableName tn = TableName.valueOf(nd.getName(),
277         testName.getMethodName() + counter.getAndIncrement());
278 
279     // Delete the old table
280     if (admin.tableExists(tn)) {
281       admin.disableTable(tn);
282       admin.deleteTable(tn);
283     }
284 
285     // Create the table
286     HTableDescriptor tableDesc = new HTableDescriptor(tn);
287     tableDesc.addFamily(new HColumnDescriptor(F1));
288 
289     admin.createTable(tableDesc);
290     return tn;
291   }
292 
293   void partitionTablesByQuotaTarget(Multimap<TableName,QuotaSettings> quotas,
294       Set<TableName> tablesWithTableQuota, Set<TableName> tablesWithNamespaceQuota) {
295     // Partition the tables with quotas by table and ns quota
296     for (Entry<TableName, QuotaSettings> entry : quotas.entries()) {
297       SpaceLimitSettings settings = (SpaceLimitSettings) entry.getValue();
298       TableName tn = entry.getKey();
299       if (settings.getTableName() != null) {
300         tablesWithTableQuota.add(tn);
301       }
302       if (settings.getNamespace() != null) {
303         tablesWithNamespaceQuota.add(tn);
304       }
305 
306       if (settings.getTableName() == null && settings.getNamespace() == null) {
307         fail("Unexpected table name with null tableName and namespace: " + tn);
308       }
309     }
310   }
311 }