diff --git a/src/main/java/com/alibaba/druid/pool/ha/HighAvailableDataSource.java b/src/main/java/com/alibaba/druid/pool/ha/HighAvailableDataSource.java index 25db173583..52db11f29c 100644 --- a/src/main/java/com/alibaba/druid/pool/ha/HighAvailableDataSource.java +++ b/src/main/java/com/alibaba/druid/pool/ha/HighAvailableDataSource.java @@ -97,6 +97,7 @@ public class HighAvailableDataSource extends WrapperAdapter implements DataSourc private String dataSourceFile = DEFAULT_DATA_SOURCE_FILE; private String propertyPrefix = ""; private int poolPurgeIntervalSeconds = PoolUpdater.DEFAULT_INTERVAL; + private boolean allowEmptyPoolWhenUpdate = false; private volatile boolean inited = false; @@ -113,6 +114,7 @@ public class HighAvailableDataSource extends WrapperAdapter implements DataSourc } if (dataSourceMap == null || dataSourceMap.isEmpty()) { poolUpdater.setIntervalSeconds(poolPurgeIntervalSeconds); + poolUpdater.setAllowEmptyPool(allowEmptyPoolWhenUpdate); poolUpdater.init(); createNodeMap(); } @@ -304,6 +306,14 @@ public class HighAvailableDataSource extends WrapperAdapter implements DataSourc nodeListener.update(); // Do update in the current Thread at the startup } + public boolean isAllowEmptyPoolWhenUpdate() { + return allowEmptyPoolWhenUpdate; + } + + public void setAllowEmptyPoolWhenUpdate(boolean allowEmptyPoolWhenUpdate) { + this.allowEmptyPoolWhenUpdate = allowEmptyPoolWhenUpdate; + } + // Getters & Setters public int getPoolPurgeIntervalSeconds() { diff --git a/src/main/java/com/alibaba/druid/pool/ha/node/PoolUpdater.java b/src/main/java/com/alibaba/druid/pool/ha/node/PoolUpdater.java index 7a1f50b41d..f46e489d8e 100644 --- a/src/main/java/com/alibaba/druid/pool/ha/node/PoolUpdater.java +++ b/src/main/java/com/alibaba/druid/pool/ha/node/PoolUpdater.java @@ -50,6 +50,7 @@ public class PoolUpdater implements Observer { private ScheduledExecutorService executor; private int intervalSeconds = DEFAULT_INTERVAL; private volatile boolean inited = false; + private boolean allowEmptyPool = false; public PoolUpdater(HighAvailableDataSource highAvailableDataSource) { setHighAvailableDataSource(highAvailableDataSource); @@ -213,6 +214,11 @@ public class PoolUpdater implements Observer { if (nodeName == null || nodeName.isEmpty() || !map.containsKey(nodeName)) { return; } + map = highAvailableDataSource.getAvailableDataSourceMap(); + if (!allowEmptyPool && map.size() == 1 && map.containsKey(nodeName)) { + LOG.warn(nodeName + " is the only DataSource left, don't remove it."); + return; + } blacklistNode(nodeName); } @@ -247,4 +253,12 @@ public class PoolUpdater implements Observer { public void setIntervalSeconds(int intervalSeconds) { this.intervalSeconds = intervalSeconds; } + + public boolean isAllowEmptyPool() { + return allowEmptyPool; + } + + public void setAllowEmptyPool(boolean allowEmptyPool) { + this.allowEmptyPool = allowEmptyPool; + } } diff --git a/src/test/java/com/alibaba/druid/pool/ha/node/FileNodeListenerTest.java b/src/test/java/com/alibaba/druid/pool/ha/node/FileNodeListenerTest.java index de6255e34c..e77e31b0b7 100644 --- a/src/test/java/com/alibaba/druid/pool/ha/node/FileNodeListenerTest.java +++ b/src/test/java/com/alibaba/druid/pool/ha/node/FileNodeListenerTest.java @@ -55,11 +55,11 @@ public class FileNodeListenerTest { DruidDataSource foo = (DruidDataSource) dataSource.getAvailableDataSourceMap().get("foo"); assertEquals("jdbc:derby:memory:foo;create=true", foo.getUrl()); - // Remove all + // try to emove all but we have to keep one left writePropertiesFile(file, new Properties()); Thread.sleep(6000); - assertTrue(dataSource.getAvailableDataSourceMap().isEmpty()); + assertEquals(1, dataSource.getAvailableDataSourceMap().size()); dataSource.destroy(); } diff --git a/src/test/java/com/alibaba/druid/pool/ha/node/PoolUpdaterTest.java b/src/test/java/com/alibaba/druid/pool/ha/node/PoolUpdaterTest.java index ac1adcf4c3..70a8215bee 100644 --- a/src/test/java/com/alibaba/druid/pool/ha/node/PoolUpdaterTest.java +++ b/src/test/java/com/alibaba/druid/pool/ha/node/PoolUpdaterTest.java @@ -39,6 +39,7 @@ public class PoolUpdaterTest { @Test public void testUpdate() { haDataSource.getDataSourceMap().put("foo", new MockDataSource("foo")); + haDataSource.getDataSourceMap().put("bar", new MockDataSource("bar")); NodeEvent event = new NodeEvent(); event.setNodeName("foo"); @@ -91,6 +92,18 @@ public class PoolUpdaterTest { assertTrue(haDataSource.getDataSourceMap().isEmpty()); } + @Test + public void testUpdate_onlyOneLeftToRemove() { + haDataSource.getDataSourceMap().put("foo", new MockDataSource("foo")); + + NodeEvent event = new NodeEvent(); + event.setNodeName("foo"); + event.setType(NodeEventTypeEnum.DELETE); + + updater.update(new FileNodeListener(), new NodeEvent[] { event }); + assertTrue(haDataSource.getDataSourceMap().containsKey("foo")); + } + @Test public void testAddNode() { String url = "jdbc:derby:memory:foo;create=true"; @@ -127,6 +140,7 @@ public class PoolUpdaterTest { @Test public void testDeleteNode() { haDataSource.getDataSourceMap().put("foo", new MockDataSource("foo")); + haDataSource.getDataSourceMap().put("bar", new MockDataSource("bar")); NodeEvent event = new NodeEvent(); event.setNodeName("foo"); diff --git a/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeListenerTest.java b/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeListenerTest.java index 619e111e5a..62fc6f02d8 100644 --- a/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeListenerTest.java +++ b/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeListenerTest.java @@ -56,10 +56,11 @@ public class ZookeeperNodeListenerTest { // 2. Register one Node ZookeeperNodeRegister register = registerNodeAndReturnRegister(); + ZookeeperNodeRegister register2 = registerAnotherNodeAndReturnRegister(); Thread.sleep(3000); // Wait for the Node to be created. assertFalse(dataSource.getDataSourceMap().isEmpty()); - assertEquals(1, dataSource.getAvailableDataSourceMap().size()); + assertEquals(2, dataSource.getAvailableDataSourceMap().size()); DruidDataSource ds = (DruidDataSource) dataSource.getAvailableDataSourceMap().get("foo.test-foo"); assertEquals("jdbc:derby:memory:foo;create=true", ds.getUrl()); assertNotNull(dataSource.getConnection()); @@ -70,15 +71,16 @@ public class ZookeeperNodeListenerTest { for (int i = 0; i < 10; i++) { Thread.sleep(500); if (dataSource.isInBlackList("foo.test-foo")) { - assertTrue(dataSource.getAvailableDataSourceMap().isEmpty()); + assertEquals(1, dataSource.getAvailableDataSourceMap().size()); assertFalse(dataSource.getDataSourceMap().isEmpty()); - assertNull(dataSource.getConnection()); + assertNotNull(dataSource.getConnection()); checkBlackList = true; break; } } assertTrue(checkBlackList); - + register2.destroy(); + Thread.sleep(3000); // Wait for the destory to be finished dataSource.destroy(); } @@ -187,6 +189,25 @@ public class ZookeeperNodeListenerTest { return register; } + private ZookeeperNodeRegister registerAnotherNodeAndReturnRegister() { + ZookeeperNodeRegister register = new ZookeeperNodeRegister(); + register.setZkConnectString(server.getConnectString()); + register.setPath(PATH); + register.init(); + + List payload = new ArrayList(); + ZookeeperNodeInfo node = new ZookeeperNodeInfo(); + node.setPrefix("bar"); + node.setHost("127.0.0.1"); + node.setPort(5678); + node.setDatabase("bar"); + node.setUsername("sa"); + node.setPassword(""); + payload.add(node); + register.register("test-bar", payload); + return register; + } + @Test public void testInitFailure() { ZookeeperNodeListener listener = new ZookeeperNodeListener();