Skip to content

Commit 2800c57

Browse files
sohamisohami
committed
DRILL-7191: RM blobs persistence in Zookeeper for Distributed RM.
Added stubs for QueryResourceManager exit and wait/cleanup thread Update MemoryCalculator to use DrillNode instead of DrillbitEndpoint Changes to support localbit resource registration to cluster state blob using DrillbitStatusListener Support ThrottledResourceManager via ResourceManagerBuilder Add some E2E tests and RMStateBlobs tests along with some bug fixes Fix TestRMConfigLoad tests to handle case where ZKQueues are explicitly enabled
1 parent 2b6a91a commit 2800c57

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

45 files changed

+1544
-575
lines changed

‎drill-yarn/src/main/java/org/apache/drill/yarn/zk/ZKClusterCoordinator.java

+30-34
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,6 @@
1717
*/
1818
package org.apache.drill.yarn.zk;
1919

20-
import static org.apache.drill.shaded.guava.com.google.common.collect.Collections2.transform;
21-
22-
import java.io.IOException;
23-
import java.util.Collection;
24-
import java.util.Collections;
25-
import java.util.HashSet;
26-
import java.util.Set;
27-
import java.util.concurrent.CountDownLatch;
28-
import java.util.concurrent.TimeUnit;
29-
30-
import org.apache.drill.shaded.guava.com.google.common.base.Throwables;
3120
import org.apache.commons.logging.Log;
3221
import org.apache.commons.logging.LogFactory;
3322
import org.apache.curator.RetryPolicy;
@@ -55,8 +44,16 @@
5544
import org.apache.drill.exec.coord.zk.ZkTransientStoreFactory;
5645
import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
5746
import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint.State;
47+
import org.apache.drill.shaded.guava.com.google.common.base.Throwables;
5848

59-
import org.apache.drill.shaded.guava.com.google.common.base.Function;
49+
import java.io.IOException;
50+
import java.util.Collection;
51+
import java.util.Collections;
52+
import java.util.HashMap;
53+
import java.util.Map;
54+
import java.util.concurrent.CountDownLatch;
55+
import java.util.concurrent.TimeUnit;
56+
import java.util.stream.Collectors;
6057

6158
/**
6259
* Manages cluster coordination utilizing zookeeper.
@@ -87,8 +84,7 @@ public class ZKClusterCoordinator extends ClusterCoordinator {
8784

8885
private CuratorFramework curator;
8986
private ServiceDiscovery<DrillbitEndpoint> discovery;
90-
private volatile Collection<DrillbitEndpoint> endpoints = Collections
91-
.emptyList();
87+
private volatile Map<String, DrillbitEndpoint> endpointsMap = Collections.emptyMap();
9288
private final String serviceName;
9389
private final CountDownLatch initialConnection = new CountDownLatch(1);
9490
private final TransientStoreFactory factory;
@@ -214,7 +210,7 @@ public void unregister(RegistrationHandle handle) {
214210

215211
@Override
216212
public Collection<DrillbitEndpoint> getAvailableEndpoints() {
217-
return this.endpoints;
213+
return this.endpointsMap.values();
218214
}
219215

220216
@Override
@@ -233,35 +229,33 @@ public <V> TransientStore<V> getOrCreateTransientStore(
233229

234230
private synchronized void updateEndpoints() {
235231
try {
236-
Collection<DrillbitEndpoint> newDrillbitSet = transform(
237-
discovery.queryForInstances(serviceName),
238-
new Function<ServiceInstance<DrillbitEndpoint>, DrillbitEndpoint>() {
239-
@Override
240-
public DrillbitEndpoint apply(
241-
ServiceInstance<DrillbitEndpoint> input) {
242-
return input.getPayload();
243-
}
244-
});
232+
// All active bits in the Zookeeper
233+
final Map<String, DrillbitEndpoint> UUIDtoEndpoints = discovery.queryForInstances(serviceName).stream()
234+
.collect(Collectors.toConcurrentMap(ServiceInstance::getId, ServiceInstance::getPayload));
245235

246236
// set of newly dead bits : original bits - new set of active bits.
247-
Set<DrillbitEndpoint> unregisteredBits = new HashSet<>(endpoints);
248-
unregisteredBits.removeAll(newDrillbitSet);
237+
Map<String, DrillbitEndpoint> unregisteredBits = new HashMap<>(endpointsMap);
238+
for (Map.Entry<String, DrillbitEndpoint> newEndpoint : UUIDtoEndpoints.entrySet()) {
239+
unregisteredBits.remove(newEndpoint.getKey());
240+
}
249241

250242
// Set of newly live bits : new set of active bits - original bits.
251-
Set<DrillbitEndpoint> registeredBits = new HashSet<>(newDrillbitSet);
252-
registeredBits.removeAll(endpoints);
243+
Map<String, DrillbitEndpoint> registeredBits = new HashMap<>(UUIDtoEndpoints);
244+
for (Map.Entry<String, DrillbitEndpoint> newEndpoint : endpointsMap.entrySet()) {
245+
registeredBits.remove(newEndpoint.getKey());
246+
}
253247

254-
endpoints = newDrillbitSet;
248+
endpointsMap = UUIDtoEndpoints;
255249

256250
if (logger.isDebugEnabled()) {
257251
StringBuilder builder = new StringBuilder();
258252
builder.append("Active drillbit set changed. Now includes ");
259-
builder.append(newDrillbitSet.size());
253+
builder.append(UUIDtoEndpoints.size());
260254
builder.append(" total bits.");
261-
if (!newDrillbitSet.isEmpty()) {
255+
if (!UUIDtoEndpoints.isEmpty()) {
262256
builder.append(" New active drillbits: \n");
263257
}
264-
for (DrillbitEndpoint bit : newDrillbitSet) {
258+
for (DrillbitEndpoint bit : UUIDtoEndpoints.values()) {
265259
builder.append('\t');
266260
builder.append(bit.getAddress());
267261
builder.append(':');
@@ -277,11 +271,13 @@ public DrillbitEndpoint apply(
277271

278272
// Notify the drillbit listener for newly unregistered bits.
279273
if (!(unregisteredBits.isEmpty())) {
280-
drillbitUnregistered(unregisteredBits);
274+
drillbitUnregistered(unregisteredBits.entrySet().stream().collect(
275+
Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey)));
281276
}
282277
// Notify the drillbit listener for newly registered bits.
283278
if (!(registeredBits.isEmpty())) {
284-
drillbitRegistered(registeredBits);
279+
drillbitRegistered(registeredBits.entrySet().stream().collect(
280+
Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey)));
285281
}
286282

287283
} catch (Exception e) {

‎drill-yarn/src/main/java/org/apache/drill/yarn/zk/ZKRegistry.java

+13-12
Original file line numberDiff line numberDiff line change
@@ -17,24 +17,24 @@
1717
*/
1818
package org.apache.drill.yarn.zk;
1919

20-
import java.util.ArrayList;
21-
import java.util.HashMap;
22-
import java.util.List;
23-
import java.util.Map;
24-
import java.util.Set;
25-
26-
import org.apache.drill.shaded.guava.com.google.common.annotations.VisibleForTesting;
2720
import org.apache.commons.logging.Log;
2821
import org.apache.commons.logging.LogFactory;
2922
import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
3023
import org.apache.drill.exec.work.foreman.DrillbitStatusListener;
24+
import org.apache.drill.shaded.guava.com.google.common.annotations.VisibleForTesting;
3125
import org.apache.drill.yarn.appMaster.AMWrapperException;
3226
import org.apache.drill.yarn.appMaster.EventContext;
3327
import org.apache.drill.yarn.appMaster.Pollable;
3428
import org.apache.drill.yarn.appMaster.RegistryHandler;
3529
import org.apache.drill.yarn.appMaster.Task;
3630
import org.apache.drill.yarn.appMaster.TaskLifecycleListener;
3731

32+
import java.util.ArrayList;
33+
import java.util.HashMap;
34+
import java.util.List;
35+
import java.util.Map;
36+
import java.util.Set;
37+
3838
/**
3939
* AM-specific implementation of a Drillbit registry backed by ZooKeeper.
4040
* Listens to ZK events for registering a Drillbit and deregistering. Alerts the
@@ -297,11 +297,12 @@ public AckEvent(Task task, DrillbitEndpoint endpoint) {
297297
* Callback from ZK to indicate that one or more drillbits have become
298298
* registered. We handle registrations in a critical section, then alert the
299299
* cluster controller outside the critical section.
300+
* @param registeredDrillbitsUUID
300301
*/
301302

302303
@Override
303-
public void drillbitRegistered(Set<DrillbitEndpoint> registeredDrillbits) {
304-
List<AckEvent> updates = registerDrillbits(registeredDrillbits);
304+
public void drillbitRegistered(Map<DrillbitEndpoint, String> registeredDrillbitsUUID) {
305+
List<AckEvent> updates = registerDrillbits(registeredDrillbitsUUID.keySet());
305306
for (AckEvent event : updates) {
306307
if (event.task == null) {
307308
registryHandler.reserveHost(event.endpoint.getAddress());
@@ -363,12 +364,12 @@ private AckEvent drillbitRegistered(DrillbitEndpoint dbe) {
363364
* Callback from ZK to indicate that one or more drillbits have become
364365
* deregistered from ZK. We handle the deregistrations in a critical section,
365366
* but updates to the cluster controller outside of a critical section.
367+
* @param unregisteredDrillbitsUUID
366368
*/
367369

368370
@Override
369-
public void drillbitUnregistered(
370-
Set<DrillbitEndpoint> unregisteredDrillbits) {
371-
List<AckEvent> updates = unregisterDrillbits(unregisteredDrillbits);
371+
public void drillbitUnregistered(Map<DrillbitEndpoint, String> unregisteredDrillbitsUUID) {
372+
List<AckEvent> updates = unregisterDrillbits(unregisteredDrillbitsUUID.keySet());
372373
for (AckEvent event : updates) {
373374
registryHandler.completionAck(event.task, ENDPOINT_PROPERTY);
374375
}

‎drill-yarn/src/test/java/org/apache/drill/yarn/zk/TestZkRegistry.java

+13-14
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,6 @@
1717
*/
1818
package org.apache.drill.yarn.zk;
1919

20-
import static org.junit.Assert.assertEquals;
21-
import static org.junit.Assert.assertNotNull;
22-
import static org.junit.Assert.assertNull;
23-
import static org.junit.Assert.assertTrue;
24-
25-
import java.util.List;
26-
import java.util.Map;
27-
import java.util.Set;
28-
2920
import org.apache.curator.framework.CuratorFramework;
3021
import org.apache.curator.framework.CuratorFrameworkFactory;
3122
import org.apache.curator.retry.RetryNTimes;
@@ -41,6 +32,15 @@
4132
import org.apache.drill.yarn.zk.ZKRegistry.DrillbitTracker;
4233
import org.junit.Test;
4334

35+
import java.util.List;
36+
import java.util.Map;
37+
import java.util.Set;
38+
39+
import static org.junit.Assert.assertEquals;
40+
import static org.junit.Assert.assertNotNull;
41+
import static org.junit.Assert.assertNull;
42+
import static org.junit.Assert.assertTrue;
43+
4444
/**
4545
* Tests for the AM version of the cluster coordinator. The AM version has no
4646
* dependencies on the DoY config system or other systems, making it easy to
@@ -108,14 +108,13 @@ private class TestDrillbitStatusListener implements DrillbitStatusListener {
108108
protected Set<DrillbitEndpoint> removed;
109109

110110
@Override
111-
public void drillbitUnregistered(
112-
Set<DrillbitEndpoint> unregisteredDrillbits) {
113-
removed = unregisteredDrillbits;
111+
public void drillbitUnregistered(Map<DrillbitEndpoint, String> unregisteredDrillbitsUUID) {
112+
removed = unregisteredDrillbitsUUID.keySet();
114113
}
115114

116115
@Override
117-
public void drillbitRegistered(Set<DrillbitEndpoint> registeredDrillbits) {
118-
added = registeredDrillbits;
116+
public void drillbitRegistered(Map<DrillbitEndpoint, String> registeredDrillbitsUUID) {
117+
added = registeredDrillbitsUUID.keySet();
119118
}
120119

121120
public void clear() {

‎exec/java-exec/src/main/java/org/apache/drill/exec/coord/ClusterCoordinator.java

+5-6
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525

2626
import java.util.Collection;
2727
import java.util.Map;
28-
import java.util.Set;
2928
import java.util.concurrent.ConcurrentHashMap;
3029

3130
/**
@@ -107,17 +106,17 @@ public interface RegistrationHandle {
107106

108107
/**
109108
* Actions to take when there are a set of new de-active drillbits.
110-
* @param unregisteredBits
109+
* @param unregisteredBitsUUID
111110
*/
112-
protected void drillbitUnregistered(Set<DrillbitEndpoint> unregisteredBits) {
111+
protected void drillbitUnregistered(Map<DrillbitEndpoint, String> unregisteredBitsUUID) {
113112
for (DrillbitStatusListener listener : listeners.keySet()) {
114-
listener.drillbitUnregistered(unregisteredBits);
113+
listener.drillbitUnregistered(unregisteredBitsUUID);
115114
}
116115
}
117116

118-
protected void drillbitRegistered(Set<DrillbitEndpoint> registeredBits) {
117+
protected void drillbitRegistered(Map<DrillbitEndpoint, String> registeredBitsUUID) {
119118
for (DrillbitStatusListener listener : listeners.keySet()) {
120-
listener.drillbitRegistered(registeredBits);
119+
listener.drillbitRegistered(registeredBitsUUID);
121120
}
122121
}
123122

‎exec/java-exec/src/main/java/org/apache/drill/exec/coord/zk/ZKClusterCoordinator.java

+23-22
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import org.apache.curator.x.discovery.ServiceInstance;
3232
import org.apache.curator.x.discovery.details.ServiceCacheListener;
3333
import org.apache.drill.common.AutoCloseables;
34+
import org.apache.drill.common.DrillNode;
3435
import org.apache.drill.common.config.DrillConfig;
3536
import org.apache.drill.exec.ExecConstants;
3637
import org.apache.drill.exec.coord.ClusterCoordinator;
@@ -46,11 +47,8 @@
4647

4748
import java.io.IOException;
4849
import java.util.Collection;
49-
import java.util.Collections;
5050
import java.util.HashMap;
51-
import java.util.HashSet;
5251
import java.util.Map;
53-
import java.util.Set;
5452
import java.util.concurrent.ConcurrentHashMap;
5553
import java.util.concurrent.CountDownLatch;
5654
import java.util.concurrent.TimeUnit;
@@ -66,15 +64,13 @@ public class ZKClusterCoordinator extends ClusterCoordinator {
6664

6765
private CuratorFramework curator;
6866
private ServiceDiscovery<DrillbitEndpoint> discovery;
69-
private volatile Collection<DrillbitEndpoint> endpoints = Collections.emptyList();
7067
private final String serviceName;
7168
private final CountDownLatch initialConnection = new CountDownLatch(1);
7269
private final TransientStoreFactory factory;
7370
private ServiceCache<DrillbitEndpoint> serviceCache;
74-
private DrillbitEndpoint endpoint;
7571

7672
// endpointsMap maps String UUID to Drillbit endpoints
77-
private ConcurrentHashMap<String, DrillbitEndpoint> endpointsMap = new ConcurrentHashMap<>();
73+
private Map<String, DrillbitEndpoint> endpointsMap = new ConcurrentHashMap<>();
7874
private static final Pattern ZK_COMPLEX_STRING = Pattern.compile("(^.*?)/(.*)/([^/]*)$");
7975

8076
public ZKClusterCoordinator(DrillConfig config, String connect) {
@@ -220,7 +216,8 @@ public void unregister(RegistrationHandle handle) {
220216
*/
221217
public RegistrationHandle update(RegistrationHandle handle, State state) {
222218
ZKRegistrationHandle h = (ZKRegistrationHandle) handle;
223-
try {
219+
final DrillbitEndpoint endpoint;
220+
try {
224221
endpoint = h.endpoint.toBuilder().setState(state).build();
225222
ServiceInstance<DrillbitEndpoint> serviceInstance = ServiceInstance.<DrillbitEndpoint>builder()
226223
.name(serviceName)
@@ -231,6 +228,7 @@ public RegistrationHandle update(RegistrationHandle handle, State state) {
231228
Throwables.throwIfUnchecked(e);
232229
throw new RuntimeException(e);
233230
}
231+
handle.setEndPoint(endpoint);
234232
return handle;
235233
}
236234

@@ -282,46 +280,49 @@ public <V> TransientStore<V> getOrCreateTransientStore(final TransientStoreConfi
282280
private synchronized void updateEndpoints() {
283281
try {
284282
// All active bits in the Zookeeper
285-
final Map<String, DrillbitEndpoint> activeEndpointsUUID = discovery.queryForInstances(serviceName).stream()
283+
final Map<String, DrillbitEndpoint> UUIDtoEndpoints = discovery.queryForInstances(serviceName).stream()
286284
.collect(Collectors.toMap(ServiceInstance::getId, ServiceInstance::getPayload));
287285

288-
final Map<DrillbitEndpoint, String> UUIDtoEndpoints = activeEndpointsUUID.entrySet().stream()
289-
.collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
286+
final Map<DrillNode, String> activeEndpointsUUID = UUIDtoEndpoints.entrySet().stream()
287+
.collect(Collectors.toMap(x -> DrillNode.create(x.getValue()), Map.Entry::getKey));
290288

291289
// set of newly dead bits : original bits - new set of active bits.
292-
Set<DrillbitEndpoint> unregisteredBits = new HashSet<>();
290+
final Map<DrillbitEndpoint, String> unregisteredBits = new HashMap<>();
293291
// Set of newly live bits : new set of active bits - original bits.
294-
Set<DrillbitEndpoint> registeredBits = new HashSet<>();
292+
final Map<DrillbitEndpoint, String> registeredBits = new HashMap<>();
295293

296294

297295
// Updates the endpoints map if there is a change in state of the endpoint or with the addition
298296
// of new drillbit endpoints. Registered endpoints is set to newly live drillbit endpoints.
299-
for (Map.Entry<String, DrillbitEndpoint> endpointToUUID : activeEndpointsUUID.entrySet()) {
300-
endpointsMap.put(endpointToUUID.getKey(), endpointToUUID.getValue());
297+
for (Map.Entry<String, DrillbitEndpoint> endpoint : UUIDtoEndpoints.entrySet()) {
298+
// check if this bit is newly added bit
299+
if (!endpointsMap.containsKey(endpoint.getKey())) {
300+
registeredBits.put(endpoint.getValue(), endpoint.getKey());
301+
}
302+
endpointsMap.put(endpoint.getKey(), endpoint.getValue());
301303
}
302304

303305
// Remove all the endpoints that are newly dead
304306
for ( String bitUUID: endpointsMap.keySet()) {
305-
if (!activeEndpointsUUID.containsKey(bitUUID)) {
307+
if (!UUIDtoEndpoints.containsKey(bitUUID)) {
306308
final DrillbitEndpoint unregisteredBit = endpointsMap.get(bitUUID);
307-
unregisteredBits.add(unregisteredBit);
308-
309-
if (UUIDtoEndpoints.containsKey(unregisteredBit)) {
309+
unregisteredBits.put(unregisteredBit, bitUUID);
310+
final DrillNode unregisteredNode = DrillNode.create(unregisteredBit);
311+
if (activeEndpointsUUID.containsKey(unregisteredNode)) {
310312
logger.info("Drillbit registered again with different UUID. [Details: Address: {}, UserPort: {}," +
311313
" PreviousUUID: {}, CurrentUUID: {}", unregisteredBit.getAddress(), unregisteredBit.getUserPort(),
312-
bitUUID, UUIDtoEndpoints.get(unregisteredBit));
314+
bitUUID, activeEndpointsUUID.get(unregisteredNode));
313315
}
314316
endpointsMap.remove(bitUUID);
315317
}
316318
}
317-
endpoints = endpointsMap.values();
318319
if (logger.isDebugEnabled()) {
319320
StringBuilder builder = new StringBuilder();
320321
builder.append("Active drillbit set changed. Now includes ");
321-
builder.append(activeEndpointsUUID.size());
322+
builder.append(UUIDtoEndpoints.size());
322323
builder.append(" total bits. New active drillbits:\n");
323324
builder.append("Address | User Port | Control Port | Data Port | Version | State\n");
324-
for (DrillbitEndpoint bit: activeEndpointsUUID.values()) {
325+
for (DrillbitEndpoint bit: UUIDtoEndpoints.values()) {
325326
builder.append(bit.getAddress()).append(" | ");
326327
builder.append(bit.getUserPort()).append(" | ");
327328
builder.append(bit.getControlPort()).append(" | ");

0 commit comments

Comments
 (0)