private void updateNodes(MemoryPoolAssignmentsRequest assignments) {
Set<Node> activeNodes = nodeManager.getActiveNodes();
ImmutableSet<String> activeNodeIds =
activeNodes.stream().map(Node::getNodeIdentifier).collect(toImmutableSet());
// Remove nodes that don't exist anymore
// Make a copy to materialize the set difference
Set<String> deadNodes = ImmutableSet.copyOf(difference(nodes.keySet(), activeNodeIds));
nodes.keySet().removeAll(deadNodes);
// Add new nodes
for (Node node : activeNodes) {
if (!nodes.containsKey(node.getNodeIdentifier())) {
nodes.put(
node.getNodeIdentifier(),
new RemoteNodeMemory(
httpClient,
memoryInfoCodec,
assignmentsRequestJsonCodec,
locationFactory.createMemoryInfoLocation(node)));
}
}
// Schedule refresh
for (RemoteNodeMemory node : nodes.values()) {
node.asyncRefresh(assignments);
}
}
@Override
public ConnectorSplitSource getSplits(
ConnectorTransactionHandle transactionHandle,
ConnectorSession session,
ConnectorTableLayoutHandle layoutHandle) {
AtopTableLayoutHandle handle =
checkType(layoutHandle, AtopTableLayoutHandle.class, "layoutHandle");
AtopTableHandle table = handle.getTableHandle();
List<ConnectorSplit> splits = new ArrayList<>();
DateTime end = DateTime.now().withZone(timeZone);
for (Node node : nodeManager.getActiveDatasourceNodes(connectorId.getId())) {
DateTime start = end.minusDays(maxHistoryDays - 1).withTimeAtStartOfDay();
while (start.isBefore(end)) {
DateTime splitEnd = start.withTime(23, 59, 59, 999);
Domain splitDomain =
Domain.create(
ValueSet.ofRanges(
Range.range(TIMESTAMP, start.getMillis(), true, splitEnd.getMillis(), true)),
false);
if (handle.getStartTimeConstraint().overlaps(splitDomain)
&& handle.getEndTimeConstraint().overlaps(splitDomain)) {
splits.add(new AtopSplit(table.getTable(), node.getHostAndPort(), start));
}
start = start.plusDays(1).withTimeAtStartOfDay();
}
}
return new FixedSplitSource(connectorId.getId(), splits);
}
private void addRows(Builder table, Set<Node> nodes, NodeState state) {
for (Node node : nodes) {
table.addRow(
node.getNodeIdentifier(),
node.getHttpUri().toString(),
getNodeVersion(node),
isCoordinator(node),
state.toString().toLowerCase());
}
}
@Override
public URI createTaskLocation(Node node, TaskId taskId) {
Preconditions.checkNotNull(node, "node is null");
Preconditions.checkNotNull(taskId, "taskId is null");
return uriBuilderFrom(node.getHttpUri())
.appendPath("/v1/task")
.appendPath(taskId.toString())
.build();
}
@Test
public void testPredicatePushdown() throws Exception {
for (Node node : nodes) {
String nodeIdentifier = node.getNodeIdentifier();
TupleDomain<ColumnHandle> nodeTupleDomain =
TupleDomain.fromFixedValues(
ImmutableMap.of(columnHandle, NullableValue.of(VARCHAR, utf8Slice(nodeIdentifier))));
ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, nodeTupleDomain);
ConnectorSplitSource splitSource =
splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout);
List<ConnectorSplit> allSplits = getAllSplits(splitSource);
assertEquals(allSplits.size(), 1);
assertEquals(allSplits.get(0).getAddresses().size(), 1);
assertEquals(allSplits.get(0).getAddresses().get(0).getHostText(), nodeIdentifier);
}
}
@Override
public ConnectorSplitSource getSplits(
ConnectorSession session, ConnectorTableLayoutHandle layout) {
TpchTableHandle tableHandle =
checkType(layout, TpchTableLayoutHandle.class, "layout").getTable();
Set<Node> nodes = nodeManager.getActiveDatasourceNodes(connectorId);
checkState(!nodes.isEmpty(), "No TPCH nodes available");
int totalParts = nodes.size() * splitsPerNode;
int partNumber = 0;
// Split the data using split and skew by the number of nodes available.
ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
for (Node node : nodes) {
for (int i = 0; i < splitsPerNode; i++) {
splits.add(
new TpchSplit(
tableHandle, partNumber, totalParts, ImmutableList.of(node.getHostAndPort())));
partNumber++;
}
}
return new FixedSplitSource(connectorId, splits.build());
}
@Override
public RemoteTask createRemoteTask(
Session session,
TaskId taskId,
com.facebook.presto.spi.Node node,
PlanFragment fragment,
Multimap<PlanNodeId, Split> initialSplits,
OutputBuffers outputBuffers) {
return new HttpRemoteTask(
session,
taskId,
node.getNodeIdentifier(),
locationFactory.createTaskLocation(node, taskId),
fragment,
initialSplits,
outputBuffers,
httpClient,
executor,
maxConsecutiveErrorCount,
minErrorDuration,
taskInfoCodec,
taskUpdateRequestCodec);
}
public class TestJmxSplitManager {
private final Node localNode = new TestingNode("host1");
private final Set<Node> nodes =
ImmutableSet.of(localNode, new TestingNode("host2"), new TestingNode("host3"));
private final JmxColumnHandle columnHandle = new JmxColumnHandle("test", "node", VARCHAR);
private final JmxTableHandle tableHandle =
new JmxTableHandle("test", "objectName", ImmutableList.of(columnHandle));
private final JmxSplitManager splitManager =
new JmxSplitManager("test", new TestingNodeManager());
private final JmxMetadata metadata = new JmxMetadata("test", getPlatformMBeanServer());
private final JmxRecordSetProvider recordSetProvider =
new JmxRecordSetProvider(getPlatformMBeanServer(), localNode.getNodeIdentifier());
@Test
public void testPredicatePushdown() throws Exception {
for (Node node : nodes) {
String nodeIdentifier = node.getNodeIdentifier();
TupleDomain<ColumnHandle> nodeTupleDomain =
TupleDomain.fromFixedValues(
ImmutableMap.of(columnHandle, NullableValue.of(VARCHAR, utf8Slice(nodeIdentifier))));
ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, nodeTupleDomain);
ConnectorSplitSource splitSource =
splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout);
List<ConnectorSplit> allSplits = getAllSplits(splitSource);
assertEquals(allSplits.size(), 1);
assertEquals(allSplits.get(0).getAddresses().size(), 1);
assertEquals(allSplits.get(0).getAddresses().get(0).getHostText(), nodeIdentifier);
}
}
@Test
public void testNoPredicate() throws Exception {
ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, TupleDomain.all());
ConnectorSplitSource splitSource =
splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout);
List<ConnectorSplit> allSplits = getAllSplits(splitSource);
assertEquals(allSplits.size(), nodes.size());
Set<String> actualNodes = nodes.stream().map(Node::getNodeIdentifier).collect(toSet());
Set<String> expectedNodes = new HashSet<>();
for (ConnectorSplit split : allSplits) {
List<HostAddress> addresses = split.getAddresses();
assertEquals(addresses.size(), 1);
expectedNodes.add(addresses.get(0).getHostText());
}
assertEquals(actualNodes, expectedNodes);
}
@Test
public void testRecordSetProvider() throws Exception {
for (SchemaTableName schemaTableName : metadata.listTables(SESSION, "jmx")) {
JmxTableHandle tableHandle = metadata.getTableHandle(SESSION, schemaTableName);
List<ColumnHandle> columnHandles =
ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());
ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, TupleDomain.all());
ConnectorSplitSource splitSource =
splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout);
List<ConnectorSplit> allSplits = getAllSplits(splitSource);
assertEquals(allSplits.size(), nodes.size());
ConnectorSplit split = allSplits.get(0);
RecordSet recordSet =
recordSetProvider.getRecordSet(
JmxTransactionHandle.INSTANCE, SESSION, split, columnHandles);
try (RecordCursor cursor = recordSet.cursor()) {
while (cursor.advanceNextPosition()) {
for (int i = 0; i < recordSet.getColumnTypes().size(); i++) {
cursor.isNull(i);
}
}
}
}
}
private static List<ConnectorSplit> getAllSplits(ConnectorSplitSource splitSource)
throws InterruptedException, ExecutionException {
ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
while (!splitSource.isFinished()) {
List<ConnectorSplit> batch = splitSource.getNextBatch(1000).get();
splits.addAll(batch);
}
return splits.build();
}
private class TestingNodeManager implements NodeManager {
@Override
public Set<Node> getNodes(NodeState state) {
return nodes;
}
@Override
public Set<Node> getActiveDatasourceNodes(String datasourceName) {
return nodes;
}
@Override
public Node getCurrentNode() {
return localNode;
}
@Override
public Set<Node> getCoordinators() {
return ImmutableSet.of(localNode);
}
}
private static class TestingNode implements Node {
private final String hostname;
public TestingNode(String hostname) {
this.hostname = hostname;
}
@Override
public HostAddress getHostAndPort() {
return HostAddress.fromParts(hostname, 8080);
}
@Override
public URI getHttpUri() {
return URI.create(format("http://%s:8080", hostname));
}
@Override
public String getNodeIdentifier() {
return hostname;
}
}
}