diff --git a/orion-server/src/main/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensor.java b/orion-server/src/main/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensor.java index d2504d2e..7429dbf0 100644 --- a/orion-server/src/main/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensor.java +++ b/orion-server/src/main/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensor.java @@ -51,91 +51,112 @@ public void initialize(Map 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 brokerNames; try { - MemqZookeeperClient memqZookeeperClient = new MemqZookeeperClient(cluster); - - List brokerNames = memqZookeeperClient.getBrokerNames(); - Map> writeBrokerAssignments = new HashMap<>(); - Map rawBrokerMap = new HashMap<>(); - Gson gson = new Gson(); - Set 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 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> writeBrokerAssignments = new HashMap<>(); + Map rawBrokerMap = new HashMap<>(); + Set brokersInZookeeper = new HashSet<>(); + List discoveredNodes = new ArrayList<>(); + Set 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 topicConfigMap = new HashMap<>(); - List topics = memqZookeeperClient.getTopics(); - for (String topicName : topics) { + Map 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"; + } 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); } } diff --git a/orion-server/src/main/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClient.java b/orion-server/src/main/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClient.java index 9a72e334..03feebb6 100644 --- a/orion-server/src/main/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClient.java +++ b/orion-server/src/main/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClient.java @@ -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; @@ -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(); } /** diff --git a/orion-server/src/test/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensorTest.java b/orion-server/src/test/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensorTest.java new file mode 100644 index 00000000..d3934965 --- /dev/null +++ b/orion-server/src/test/java/com/pinterest/orion/core/automation/sensor/memq/MemqClusterSensorTest.java @@ -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 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"); + } +} diff --git a/orion-server/src/test/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClientTest.java b/orion-server/src/test/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClientTest.java new file mode 100644 index 00000000..82a7e03d --- /dev/null +++ b/orion-server/src/test/java/com/pinterest/orion/core/utils/memq/zookeeper/MemqZookeeperClientTest.java @@ -0,0 +1,50 @@ +package com.pinterest.orion.core.utils.memq.zookeeper; + +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +/** + * Verifies that the chroot in a zk connection string is parsed out into a Curator namespace. + * This is the core of the fix: the chroot must NOT stay in the connection string (where it is + * dropped on dynamic reconfiguration) but be applied as a client-side namespace instead. + */ +public class MemqZookeeperClientTest { + + @Test + public void parsesSingleHostWithChroot() { + String zkUrl = "zk-host.example.com:2181/memq/cluster01"; + assertEquals("zk-host.example.com:2181", MemqZookeeperClient.parseConnectString(zkUrl)); + assertEquals("memq/cluster01", MemqZookeeperClient.parseNamespace(zkUrl)); + } + + @Test + public void parsesMultiHostWithChroot() { + String zkUrl = "10.12.90.128:2181,10.12.72.198:2181,10.12.7.82:2181/memq/cluster09"; + assertEquals("10.12.90.128:2181,10.12.72.198:2181,10.12.7.82:2181", + MemqZookeeperClient.parseConnectString(zkUrl)); + assertEquals("memq/cluster09", MemqZookeeperClient.parseNamespace(zkUrl)); + } + + @Test + public void parsesNestedChroot() { + String zkUrl = "host:2181/a/b/c"; + assertEquals("host:2181", MemqZookeeperClient.parseConnectString(zkUrl)); + assertEquals("a/b/c", MemqZookeeperClient.parseNamespace(zkUrl)); + } + + @Test + public void handlesNoChroot() { + String zkUrl = "host:2181"; + assertEquals("host:2181", MemqZookeeperClient.parseConnectString(zkUrl)); + assertNull(MemqZookeeperClient.parseNamespace(zkUrl)); + } + + @Test + public void handlesTrailingSlashAsNoChroot() { + String zkUrl = "host:2181/"; + assertEquals("host:2181", MemqZookeeperClient.parseConnectString(zkUrl)); + assertNull(MemqZookeeperClient.parseNamespace(zkUrl)); + } +}