1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20 package org.apache.hadoop.hbase.regionserver;
21
22 import static org.junit.Assert.assertEquals;
23 import static org.junit.Assert.assertTrue;
24
25 import java.io.IOException;
26
27 import org.apache.hadoop.hbase.security.Superusers;
28 import org.apache.hadoop.hbase.security.User;
29 import org.apache.hadoop.hbase.util.ByteStringer;
30 import org.apache.hadoop.conf.Configuration;
31 import org.apache.hadoop.hbase.CoordinatedStateManagerFactory;
32 import org.apache.hadoop.hbase.HBaseConfiguration;
33 import org.apache.hadoop.hbase.HConstants;
34 import org.apache.hadoop.hbase.HRegionInfo;
35 import org.apache.hadoop.hbase.testclassification.MediumTests;
36 import org.apache.hadoop.hbase.CoordinatedStateManager;
37 import org.apache.hadoop.hbase.ipc.PriorityFunction;
38 import org.apache.hadoop.hbase.protobuf.generated.ClientProtos.Get;
39 import org.apache.hadoop.hbase.protobuf.generated.ClientProtos.GetRequest;
40 import org.apache.hadoop.hbase.protobuf.generated.ClientProtos.ScanRequest;
41 import org.apache.hadoop.hbase.protobuf.generated.HBaseProtos.RegionSpecifier;
42 import org.apache.hadoop.hbase.protobuf.generated.HBaseProtos.RegionSpecifier.RegionSpecifierType;
43 import org.apache.hadoop.hbase.protobuf.generated.RPCProtos.RequestHeader;
44 import org.junit.Before;
45 import org.junit.Test;
46 import org.junit.experimental.categories.Category;
47 import org.mockito.Mockito;
48
49 import com.google.protobuf.ByteString;
50
51
52
53
54 @Category(MediumTests.class)
55 public class TestPriorityRpc {
56 private HRegionServer regionServer = null;
57 private PriorityFunction priority = null;
58
59 @Before
60 public void setup() {
61 Configuration conf = HBaseConfiguration.create();
62 conf.setBoolean("hbase.testing.nocluster", true);
63 CoordinatedStateManager cp = CoordinatedStateManagerFactory.getCoordinatedStateManager(conf);
64 regionServer = HRegionServer.constructRegionServer(HRegionServer.class, conf, cp);
65 priority = regionServer.rpcServices.getPriority();
66 }
67
68 @Test
69 public void testQosFunctionForMeta() throws IOException {
70 priority = regionServer.rpcServices.getPriority();
71 RequestHeader.Builder headerBuilder = RequestHeader.newBuilder();
72
73
74
75 headerBuilder.setMethodName("foo");
76
77 GetRequest.Builder getRequestBuilder = GetRequest.newBuilder();
78 RegionSpecifier.Builder regionSpecifierBuilder = RegionSpecifier.newBuilder();
79 regionSpecifierBuilder.setType(RegionSpecifierType.REGION_NAME);
80 ByteString name = ByteStringer.wrap(HRegionInfo.FIRST_META_REGIONINFO.getRegionName());
81 regionSpecifierBuilder.setValue(name);
82 RegionSpecifier regionSpecifier = regionSpecifierBuilder.build();
83 getRequestBuilder.setRegion(regionSpecifier);
84 Get.Builder getBuilder = Get.newBuilder();
85 getBuilder.setRow(ByteStringer.wrap("somerow".getBytes()));
86 getRequestBuilder.setGet(getBuilder.build());
87 GetRequest getRequest = getRequestBuilder.build();
88 RequestHeader header = headerBuilder.build();
89 HRegion mockRegion = Mockito.mock(HRegion.class);
90 HRegionServer mockRS = Mockito.mock(HRegionServer.class);
91 RSRpcServices mockRpc = Mockito.mock(RSRpcServices.class);
92 Mockito.when(mockRS.getRSRpcServices()).thenReturn(mockRpc);
93 HRegionInfo mockRegionInfo = Mockito.mock(HRegionInfo.class);
94 Mockito.when(mockRpc.getRegion((RegionSpecifier) Mockito.any())).thenReturn(mockRegion);
95 Mockito.when(mockRegion.getRegionInfo()).thenReturn(mockRegionInfo);
96 Mockito.when(mockRegionInfo.isSystemTable()).thenReturn(true);
97
98 ((AnnotationReadingPriorityFunction)priority).setRegionServer(mockRS);
99 assertEquals(HConstants.SYSTEMTABLE_QOS, priority.getPriority(header, getRequest,
100 User.createUserForTesting(regionServer.conf, "someuser", new String[]{"somegroup"})));
101 }
102
103 @Test
104 public void testQosFunctionWithoutKnownArgument() throws IOException {
105
106
107
108
109 RequestHeader.Builder headerBuilder = RequestHeader.newBuilder();
110 headerBuilder.setMethodName("foo");
111 RequestHeader header = headerBuilder.build();
112 PriorityFunction qosFunc = regionServer.rpcServices.getPriority();
113 assertEquals(HConstants.NORMAL_QOS, qosFunc.getPriority(header, null,
114 User.createUserForTesting(regionServer.conf, "someuser", new String[]{"somegroup"})));
115 }
116
117 @Test
118 public void testQosFunctionForRequestCalledBySuperUser() throws Exception {
119 RequestHeader.Builder headerBuilder = RequestHeader.newBuilder();
120 headerBuilder.setMethodName("foo");
121 RequestHeader header = headerBuilder.build();
122 PriorityFunction qosFunc = regionServer.rpcServices.getPriority();
123
124
125 regionServer.conf.set(Superusers.SUPERUSER_CONF_KEY, "samplesuperuser");
126 Superusers.initialize(regionServer.conf);
127 assertEquals(HConstants.ADMIN_QOS, qosFunc.getPriority(header, null,
128 User.createUserForTesting(regionServer.conf, "samplesuperuser",
129 new String[]{"somegroup"})));
130
131
132 regionServer.conf.set(Superusers.SUPERUSER_CONF_KEY, "@samplesupergroup");
133 Superusers.initialize(regionServer.conf);
134 assertEquals(HConstants.ADMIN_QOS, qosFunc.getPriority(header, null,
135 User.createUserForTesting(regionServer.conf, "regularuser",
136 new String[]{"samplesupergroup"})));
137 }
138
139 @Test
140 public void testQosFunctionForScanMethod() throws IOException {
141 RequestHeader.Builder headerBuilder = RequestHeader.newBuilder();
142 headerBuilder.setMethodName("Scan");
143 RequestHeader header = headerBuilder.build();
144
145
146 ScanRequest.Builder scanBuilder = ScanRequest.newBuilder();
147 ScanRequest scanRequest = scanBuilder.build();
148 HRegion mockRegion = Mockito.mock(HRegion.class);
149 HRegionServer mockRS = Mockito.mock(HRegionServer.class);
150 RSRpcServices mockRpc = Mockito.mock(RSRpcServices.class);
151 Mockito.when(mockRS.getRSRpcServices()).thenReturn(mockRpc);
152 HRegionInfo mockRegionInfo = Mockito.mock(HRegionInfo.class);
153 Mockito.when(mockRpc.getRegion((RegionSpecifier)Mockito.any())).thenReturn(mockRegion);
154 Mockito.when(mockRegion.getRegionInfo()).thenReturn(mockRegionInfo);
155 Mockito.when(mockRegionInfo.isSystemTable()).thenReturn(false);
156
157 ((AnnotationReadingPriorityFunction)priority).setRegionServer(mockRS);
158 int qos = priority.getPriority(header, scanRequest,
159 User.createUserForTesting(regionServer.conf, "someuser", new String[]{"somegroup"}));
160 assertTrue ("" + qos, qos == HConstants.NORMAL_QOS);
161
162
163 scanBuilder = ScanRequest.newBuilder();
164 scanBuilder.setScannerId(12345);
165 scanRequest = scanBuilder.build();
166
167 RegionScanner mockRegionScanner = Mockito.mock(RegionScanner.class);
168 Mockito.when(mockRpc.getScanner(12345)).thenReturn(mockRegionScanner);
169 Mockito.when(mockRegionScanner.getRegionInfo()).thenReturn(mockRegionInfo);
170 Mockito.when(mockRpc.getRegion((RegionSpecifier)Mockito.any())).thenReturn(mockRegion);
171 Mockito.when(mockRegion.getRegionInfo()).thenReturn(mockRegionInfo);
172 Mockito.when(mockRegionInfo.isSystemTable()).thenReturn(true);
173
174
175 ((AnnotationReadingPriorityFunction)priority).setRegionServer(mockRS);
176
177 assertEquals(HConstants.SYSTEMTABLE_QOS, priority.getPriority(header, scanRequest,
178 User.createUserForTesting(regionServer.conf, "someuser", new String[]{"somegroup"})));
179
180
181 Mockito.when(mockRegionInfo.isSystemTable()).thenReturn(false);
182 assertEquals(HConstants.NORMAL_QOS, priority.getPriority(header, scanRequest,
183 User.createUserForTesting(regionServer.conf, "someuser", new String[]{"somegroup"})));
184 }
185 }