This is an automated email from the ASF dual-hosted git repository. Caideyipi pushed a commit to branch fix/session-close-cleanup in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 60a4488dfab8db339bca13e4e2d03fac2b4ff64a Author: Caideyipi <[email protected]> AuthorDate: Tue Nov 3 11:08:43 2026 +0800 Make Session.close resilient to connection failures --- .../java/org/apache/iotdb/session/Session.java | 47 ++++++++++++-- .../java/org/apache/iotdb/session/SessionTest.java | 74 ++++++++++++++++++++++ 2 files changed, 116 insertions(+), 5 deletions(-) diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index de3ee417afc..9e3d41ca65c 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java @@ -87,6 +87,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.Iterator; import java.util.LinkedHashSet; import java.util.List; @@ -675,19 +676,48 @@ public class Session implements ISession { @Override public synchronized void close() throws IoTDBConnectionException { + Throwable closeException = null; try { if (isClosed) { return; } - if (enableRedirection) { + + Set<SessionConnection> closedConnections = Collections.newSetFromMap(new IdentityHashMap<>()); + if (endPointToSessionConnection != null) { for (SessionConnection sessionConnection : endPointToSessionConnection.values()) { - sessionConnection.close(); + if (sessionConnection != null && closedConnections.add(sessionConnection)) { + try { + sessionConnection.close(); + } catch (IoTDBConnectionException | RuntimeException e) { + if (closeException == null) { + closeException = e; + } else if (closeException != e) { + closeException.addSuppressed(e); + } + } + } + } + } + + SessionConnection defaultSessionConnection = getDefaultSessionConnection(); + if (defaultSessionConnection != null && closedConnections.add(defaultSessionConnection)) { + try { + defaultSessionConnection.close(); + } catch (IoTDBConnectionException | RuntimeException e) { + if (closeException == null) { + closeException = e; + } else if (closeException != e) { + closeException.addSuppressed(e); + } } - endPointToSessionConnection.clear(); } - getDefaultSessionConnection().close(); - setDefaultSessionConnection(null); } finally { + if (!isClosed) { + if (endPointToSessionConnection != null) { + endPointToSessionConnection.clear(); + } + setDefaultSessionConnection(null); + } // if executorService is null, it means that availableNodes is got from SessionPool and we // shouldn't clean that if (this.executorService != null) { @@ -698,6 +728,13 @@ public class Session implements ISession { } isClosed = true; } + + if (closeException instanceof IoTDBConnectionException) { + throw (IoTDBConnectionException) closeException; + } + if (closeException instanceof RuntimeException) { + throw (RuntimeException) closeException; + } } public SessionConnection constructSessionConnection( diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java index bf243f66d86..05575536f26 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionTest.java @@ -21,6 +21,7 @@ package org.apache.iotdb.session; import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ISession; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.util.Version; @@ -53,9 +54,11 @@ import java.util.Arrays; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Random; +import java.util.concurrent.ScheduledExecutorService; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; @@ -99,6 +102,77 @@ public class SessionTest { } } + @Test + public void testCloseAttemptsEveryConnectionAndCleansUpAfterFailure() throws Exception { + SessionConnection redirectedConnection = Mockito.mock(SessionConnection.class); + IoTDBConnectionException defaultCloseException = new IoTDBConnectionException("default"); + IoTDBConnectionException redirectedCloseException = new IoTDBConnectionException("redirected"); + Mockito.doThrow(defaultCloseException).when(sessionConnection).close(); + Mockito.doThrow(redirectedCloseException).when(redirectedConnection).close(); + + Map<TEndPoint, SessionConnection> connections = new LinkedHashMap<>(); + connections.put(new TEndPoint("default", 1), sessionConnection); + connections.put(new TEndPoint("redirected", 2), redirectedConnection); + Whitebox.setInternalState(session, "endPointToSessionConnection", connections); + Whitebox.setInternalState(session, "enableRedirection", true); + Whitebox.setInternalState(session, "isClosed", false); + + ScheduledExecutorService executorService = Mockito.mock(ScheduledExecutorService.class); + INodeSupplier availableNodes = Mockito.mock(INodeSupplier.class); + Whitebox.setInternalState(session, "executorService", executorService); + Whitebox.setInternalState(session, "availableNodes", availableNodes); + + try { + session.close(); + fail("Expected the first connection close failure to be propagated"); + } catch (IoTDBConnectionException e) { + Assert.assertSame(defaultCloseException, e); + Assert.assertArrayEquals( + new Throwable[] {redirectedCloseException}, defaultCloseException.getSuppressed()); + } + + Mockito.verify(sessionConnection).close(); + Mockito.verify(redirectedConnection).close(); + Assert.assertTrue(connections.isEmpty()); + Assert.assertNull(Whitebox.getInternalState(session, "defaultSessionConnection")); + Assert.assertNull(Whitebox.getInternalState(session, "executorService")); + Assert.assertNull(Whitebox.getInternalState(session, "availableNodes")); + Mockito.verify(executorService).shutdown(); + Mockito.verify(availableNodes).close(); + + session.close(); + Mockito.verify(sessionConnection).close(); + Mockito.verify(redirectedConnection).close(); + } + + @Test + public void testCloseClosesQueryRedirectedConnections() throws Exception { + SessionConnection redirectedConnection = Mockito.mock(SessionConnection.class); + Map<TEndPoint, SessionConnection> connections = new LinkedHashMap<>(); + connections.put(new TEndPoint("default", 1), sessionConnection); + connections.put(new TEndPoint("redirected", 2), redirectedConnection); + Whitebox.setInternalState(session, "endPointToSessionConnection", connections); + Whitebox.setInternalState(session, "enableQueryRedirection", true); + Whitebox.setInternalState(session, "isClosed", false); + + session.close(); + + Mockito.verify(sessionConnection).close(); + Mockito.verify(redirectedConnection).close(); + Assert.assertTrue(connections.isEmpty()); + } + + @Test + public void testCloseDoesNotCloseBorrowedNodeSupplier() throws Exception { + INodeSupplier availableNodes = Mockito.mock(INodeSupplier.class); + Whitebox.setInternalState(session, "availableNodes", availableNodes); + Whitebox.setInternalState(session, "isClosed", false); + + session.close(); + + Mockito.verify(availableNodes, Mockito.never()).close(); + } + @Test public void testBuildSession() { Session session1 =
