Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -51,91 +51,112 @@ public void initialize(Map<String, Object> config) throws PluginConfigurationExc

@Override
public void sense(MemqCluster cluster) throws Exception {
sense(cluster, new MemqZookeeperClient(cluster));
}

/**
* Reads cluster state from ZooKeeper and publishes it as cluster attributes.
*
* Reads are defensive: the cluster's attributes and node map are only mutated once a complete,
* trustworthy snapshot has been read. If the client is unhealthy or any read fails partway
* through, this method leaves the previously published (known-good) state untouched instead of
* overwriting it with empty, partial, or cross-cluster data.
*/
void sense(MemqCluster cluster, MemqZookeeperClient memqZookeeperClient) throws Exception {
if (!memqZookeeperClient.isConnected()) {
logger.warning("ZooKeeper client for cluster " + cluster.getClusterId()
+ " is not connected; skipping update to avoid showing incorrect data.");
return;
}

List<String> brokerNames;
try {
MemqZookeeperClient memqZookeeperClient = new MemqZookeeperClient(cluster);

List<String> brokerNames = memqZookeeperClient.getBrokerNames();
Map<String, List<String>> writeBrokerAssignments = new HashMap<>();
Map<String, Broker> rawBrokerMap = new HashMap<>();
Gson gson = new Gson();
Set<String> brokersInZookeeper = new HashSet<>();
for (String brokerName : brokerNames) {
String brokerDataJsonString = null;
try {
brokerDataJsonString = memqZookeeperClient.getBrokerData(brokerName);
} catch (KeeperException.NoNodeException e) {
cluster.getNodeMap().remove(brokerName);
logger.info(
"Broker data of " + brokerName + " is not available in zookeeper. The broker might be removed.");
continue;
} catch (Exception e) {
logger.severe(
"Faced an unknown exception when getting broker data for " + brokerName +" from zookeeper:" + e);
continue;
}
Broker broker = gson.fromJson(brokerDataJsonString, Broker.class);
NodeInfo info = new NodeInfo();
info.setClusterId(cluster.getClusterId());
String hostname = NetworkUtils.getHostnameFromIpIfAvailable(broker.getBrokerIP());
info.setHostname(hostname);
info.setIp(broker.getBrokerIP());
info.setNodeType(broker.getInstanceType());
info.setNodeId(broker.getBrokerIP());
info.setRack(broker.getLocality());
info.setServicePort(broker.getBrokerPort());
info.setTimestamp(System.currentTimeMillis());
cluster.addNodeWithoutAgent(info);

rawBrokerMap.put(broker.getBrokerIP(), broker);
for (TopicConfig topicConfig : broker.getAssignedTopics()) {
String topicName = topicConfig.getTopic();
List<String> hostnames = writeBrokerAssignments.get(topicName);
if (hostnames == null) {
hostnames = new ArrayList<>();
writeBrokerAssignments.put(topicName, hostnames);
}
hostnames.add(hostname);
}
brokersInZookeeper.add(broker.getBrokerIP());
brokerNames = memqZookeeperClient.getBrokerNames();
} catch (Exception e) {
logger.warning("Failed to read broker list for cluster " + cluster.getClusterId()
+ "; skipping update to avoid showing incorrect data: " + e);
return;
}

Gson gson = new Gson();
Map<String, List<String>> writeBrokerAssignments = new HashMap<>();
Map<String, Broker> rawBrokerMap = new HashMap<>();
Set<String> brokersInZookeeper = new HashSet<>();
List<NodeInfo> discoveredNodes = new ArrayList<>();
Set<String> removedBrokers = new HashSet<>();

for (String brokerName : brokerNames) {
String brokerDataJsonString;
try {
brokerDataJsonString = memqZookeeperClient.getBrokerData(brokerName);
} catch (KeeperException.NoNodeException e) {
removedBrokers.add(brokerName);
logger.info(
"Broker data of " + brokerName + " is not available in zookeeper. The broker might be removed.");
continue;
} catch (Exception e) {
// A transient read failure means we have an incomplete snapshot. Bail out rather than
// silently dropping brokers and reporting NO BROKER / partial assignments.
logger.warning("Failed to read broker data for " + brokerName + " in cluster "
+ cluster.getClusterId() + "; skipping update to avoid showing incorrect data: " + e);
return;
}
Broker broker = gson.fromJson(brokerDataJsonString, Broker.class);
NodeInfo info = new NodeInfo();
info.setClusterId(cluster.getClusterId());
String hostname = NetworkUtils.getHostnameFromIpIfAvailable(broker.getBrokerIP());
info.setHostname(hostname);
info.setIp(broker.getBrokerIP());
info.setNodeType(broker.getInstanceType());
info.setNodeId(broker.getBrokerIP());
info.setRack(broker.getLocality());
info.setServicePort(broker.getBrokerPort());
info.setTimestamp(System.currentTimeMillis());
discoveredNodes.add(info);

boolean noBrokerInZookeeper = false;
if (brokersInZookeeper.isEmpty()) {
logger.warning("No broker found in zookeeper for cluster " + cluster.getClusterId());
noBrokerInZookeeper = true;
} else {
// Remove brokers that are not in zookeeper from the cluster node map
for (String nodeId : cluster.getNodeMap().keySet()) {
if (!brokersInZookeeper.contains(nodeId)) {
cluster.getNodeMap().remove(nodeId);
}
}
rawBrokerMap.put(broker.getBrokerIP(), broker);
for (TopicConfig topicConfig : broker.getAssignedTopics()) {
writeBrokerAssignments.computeIfAbsent(topicConfig.getTopic(), k -> new ArrayList<>())
.add(hostname);
}
brokersInZookeeper.add(broker.getBrokerIP());
}

Map<String, TopicConfig> topicConfigMap = new HashMap<>();
List<String> topics = memqZookeeperClient.getTopics();
for (String topicName : topics) {
Map<String, TopicConfig> topicConfigMap = new HashMap<>();
try {
for (String topicName : memqZookeeperClient.getTopics()) {
String topicDataJsonString = memqZookeeperClient.getTopicData(topicName);
TopicConfig topicConfig = gson.fromJson(topicDataJsonString, TopicConfig.class);
topicConfigMap.put(topicName, topicConfig);
topicConfigMap.put(topicName, gson.fromJson(topicDataJsonString, TopicConfig.class));
}
} catch (Exception e) {
logger.warning("Failed to read topics for cluster " + cluster.getClusterId()
+ "; skipping update to avoid showing incorrect data: " + e);
return;
}

String clusterContext = "NO BROKER";
if (!noBrokerInZookeeper) {
String governorIp = memqZookeeperClient.getGovernorIp();
if (governorIp != null) {
clusterContext = "Governor: " + governorIp + "\n";
}
}
// Snapshot is complete and trustworthy: apply it to the cluster.
for (NodeInfo info : discoveredNodes) {
cluster.addNodeWithoutAgent(info);
}
for (String removed : removedBrokers) {
cluster.getNodeMap().remove(removed);
}

setAttribute(cluster, TOPIC_CONFIG, topicConfigMap);
setAttribute(cluster, RAW_BROKER_INFO, rawBrokerMap);
setAttribute(cluster, WRITE_ASSIGNMENTS, writeBrokerAssignments);
setAttribute(cluster, CLUSTER_CONTEXT, clusterContext);
} catch (Exception e) {
e.printStackTrace();
throw e;
String clusterContext;
if (brokersInZookeeper.isEmpty()) {
logger.warning("No broker found in zookeeper for cluster " + cluster.getClusterId());
clusterContext = "NO BROKER";
Comment thread
jeffxiang marked this conversation as resolved.
} else {
// Prune stale nodes only when we have a confirmed, non-empty broker list.
cluster.getNodeMap().keySet().removeIf(nodeId -> !brokersInZookeeper.contains(nodeId));
String governorIp = memqZookeeperClient.getGovernorIp();
clusterContext = governorIp != null ? "Governor: " + governorIp + "\n" : "";
}

setAttribute(cluster, TOPIC_CONFIG, topicConfigMap);
setAttribute(cluster, RAW_BROKER_INFO, rawBrokerMap);
setAttribute(cluster, WRITE_ASSIGNMENTS, writeBrokerAssignments);
setAttribute(cluster, CLUSTER_CONTEXT, clusterContext);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.utils.CloseableUtils;

import java.util.List;

Expand Down Expand Up @@ -36,27 +37,69 @@ public void disableRefreshZkClientOnException() {

/**
* Create a new Zookeeper client using the connection string provided in the cluster configuration.
*
* The chroot suffix is applied as a Curator namespace rather than as a
* ZooKeeper connection-string chroot. A connection-string chroot is dropped by Curator's
* EnsembleTracker whenever the ensemble emits a dynamic reconfiguration event, which silently
* re-points reads at the root of the (shared) ensemble and leaks data across clusters. A Curator
* namespace is applied client-side and survives reconfiguration/session loss, keeping every read
* scoped to this cluster.
* @return CuratorFramework
* @throws Exception
*/
private CuratorFramework createZkClient() throws Exception {
CuratorFramework curator = CuratorFrameworkFactory.newClient(
zkUrl,
new ExponentialBackoffRetry(1000, 3)
);
CuratorFramework curator = CuratorFrameworkFactory.builder()
.connectString(parseConnectString(zkUrl))
.retryPolicy(new ExponentialBackoffRetry(1000, 3))
.namespace(parseNamespace(zkUrl))
.build();
curator.start();
curator.blockUntilConnected();
return curator;
}

/**
* Extract the host:port list from a zk connection string, dropping any chroot suffix.
* e.g. "h1:2181,h2:2181/memq/cluster01" -> "h1:2181,h2:2181"
*/
static String parseConnectString(String zkUrl) {
int slashIndex = zkUrl.indexOf('/');
return slashIndex >= 0 ? zkUrl.substring(0, slashIndex) : zkUrl;
}

/**
* Extract the chroot path from a zk connection string as a Curator namespace (no leading slash),
* or null when there is no chroot. e.g. "h1:2181/memq/cluster01" -> "memq/cluster01"
*/
static String parseNamespace(String zkUrl) {
int slashIndex = zkUrl.indexOf('/');
if (slashIndex < 0) {
return null;
}
String namespace = zkUrl.substring(slashIndex + 1);
return namespace.isEmpty() ? null : namespace;
}

/**
* Refresh the Zookeeper client by creating a new one.
* The new client is then set in the cluster object.
* The new client is then set in the cluster object and the previous client is closed to avoid
* leaking ZooKeeper connections/sessions across refreshes.
* @throws Exception
*/
public void refreshZkClient() throws Exception {
CuratorFramework oldClient = this.zkClient;
this.zkClient = createZkClient();
cluster.setZkClient(this.zkClient);
if (oldClient != null) {
CloseableUtils.closeQuietly(oldClient);
}
}

/**
* @return true if the underlying client currently has a live connection to ZooKeeper.
*/
public boolean isConnected() {
return zkClient != null && zkClient.getZookeeperClient().isConnected();
}

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
package com.pinterest.orion.core.automation.sensor.memq;

import com.google.gson.Gson;
import com.pinterest.orion.core.Node;
import com.pinterest.orion.core.memq.MemqCluster;
import com.pinterest.orion.core.utils.memq.zookeeper.MemqZookeeperClient;
import org.junit.Before;
import org.junit.Test;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
* Verifies the defensive read behavior of {@link MemqClusterSensor}: cluster attributes are only
* published when a complete, trustworthy snapshot is read, so transient ZooKeeper failures never
* overwrite known-good state with empty / partial / cross-cluster data.
*/
public class MemqClusterSensorTest {

private MemqClusterSensor sensor;
private MemqCluster cluster;
private MemqZookeeperClient client;
private Map<String, Node> nodeMap;

@Before
public void setUp() {
sensor = spy(new MemqClusterSensor());
cluster = mock(MemqCluster.class);
client = mock(MemqZookeeperClient.class);
nodeMap = new HashMap<>();
when(cluster.getClusterId()).thenReturn("scorpion01");
when(cluster.getNodeMap()).thenReturn(nodeMap);
}

@Test
public void whenClientNotConnectedSkipsUpdate() throws Exception {
when(client.isConnected()).thenReturn(false);

sensor.sense(cluster, client);

verify(client, never()).getBrokerNames();
verify(sensor, never()).setAttribute(any(), anyString(), any());
}

@Test
public void whenBrokerListReadFailsSkipsUpdate() throws Exception {
when(client.isConnected()).thenReturn(true);
when(client.getBrokerNames()).thenThrow(new RuntimeException("zk read failed"));

sensor.sense(cluster, client);

verify(sensor, never()).setAttribute(any(), anyString(), any());
}

@Test
public void whenBrokerDataReadFailsSkipsUpdate() throws Exception {
when(client.isConnected()).thenReturn(true);
when(client.getBrokerNames()).thenReturn(Arrays.asList("broker0"));
when(client.getBrokerData("broker0")).thenThrow(new RuntimeException("partial read"));

sensor.sense(cluster, client);

// Incomplete snapshot: must not read topics or publish anything.
verify(client, never()).getTopics();
verify(sensor, never()).setAttribute(any(), anyString(), any());
}

@Test
public void whenTopicReadFailsSkipsUpdate() throws Exception {
when(client.isConnected()).thenReturn(true);
when(client.getBrokerNames()).thenReturn(Collections.emptyList());
when(client.getTopics()).thenThrow(new RuntimeException("topic read failed"));

sensor.sense(cluster, client);

verify(sensor, never()).setAttribute(any(), anyString(), any());
}

@Test
public void whenNoBrokersPublishesNoBrokerContext() throws Exception {
when(client.isConnected()).thenReturn(true);
when(client.getBrokerNames()).thenReturn(Collections.emptyList());
when(client.getTopics()).thenReturn(Collections.emptyList());

sensor.sense(cluster, client);

verify(sensor).setAttribute(cluster, MemqCluster.CLUSTER_CONTEXT, "NO BROKER");
}

@Test
public void whenBrokersPresentPublishesGovernorContext() throws Exception {
when(client.isConnected()).thenReturn(true);
when(client.getBrokerNames()).thenReturn(Arrays.asList("broker0"));
Broker broker = new Broker("10.0.0.1", (short) 9092, "c5.large", "us-east-1a", new HashSet<>());
when(client.getBrokerData("broker0")).thenReturn(new Gson().toJson(broker));
when(client.getTopics()).thenReturn(Collections.emptyList());
when(client.getGovernorIp()).thenReturn("10.0.0.1");

sensor.sense(cluster, client);

verify(sensor).setAttribute(cluster, MemqCluster.CLUSTER_CONTEXT, "Governor: 10.0.0.1\n");
}
}
Loading
Loading