[To dev/1.3] Fix SessionPool client leak on reconnection and query failures, and preserve server error messages (#45)

CritasWang · web-flow · commit 26d96b12ccb1 · 2026-02-10T15:33:43.000+08:00
* [To dev/1.3] Fix SessionPool client leak on reconnection and query failures, and preserve server error messages

* fix ByteBuffer build error ,use  Array.Reverse()
diff --git a/docs/SessionPool_Exception_Handling.md b/docs/SessionPool_Exception_Handling.md
diff --git a/src/Apache.IoTDB.Data/IoTDBDataReader.cs b/src/Apache.IoTDB.Data/IoTDBDataReader.cs
@@ -56,8 +56,8 @@ internal IoTDBDataReader(IoTDBCommand IoTDBCommand, SessionDataSet dataSet, bool
             _command = IoTDBCommand;
             _closeConnection = closeConnection;
             _fieldCount = dataSet.ColumnNames.Count;
-            _hasRows = dataSet.RowCount > 0;
-            _recordsAffected = dataSet.RowCount;
+            _hasRows = dataSet.CurrentBatchRowCount() > 0;
+            _recordsAffected = -1; // Total row count is unknown; use -1 per ADO.NET convention
             _closed = _closeConnection;
             _metas = dataSet.ColumnNames;
             _dataSet = dataSet;
diff --git a/src/Apache.IoTDB/ConcurrentClientQueue.cs b/src/Apache.IoTDB/ConcurrentClientQueue.cs
@@ -28,6 +28,7 @@ namespace Apache.IoTDB
     public class ConcurrentClientQueue
     {
         public ConcurrentQueue<Client> ClientQueue { get; }
+        internal IPoolDiagnosticReporter DiagnosticReporter { get; set; }
 
         public ConcurrentClientQueue(List<Client> clients)
         {
@@ -42,55 +43,62 @@ public ConcurrentClientQueue()
         public void Return(Client client)
         {
             Monitor.Enter(ClientQueue);
-            ClientQueue.Enqueue(client);
-            Monitor.PulseAll(ClientQueue); // wake up all threads waiting on the queue, refresh the waiting time
-            Monitor.Exit(ClientQueue);
-            Thread.Sleep(0);
-        }
-        int _ref = 0;
-        public void AddRef()
-        {
-            lock (this)
+            try
             {
-                _ref++;
+                ClientQueue.Enqueue(client);
+                Monitor.PulseAll(ClientQueue); // wake up all threads waiting on the queue, refresh the waiting time
             }
-        }
-        public int GetRef()
-        {
-            return _ref;
-        }
-        public void RemoveRef()
-        {
-            lock (this)
+            finally
             {
-                _ref--;
+                Monitor.Exit(ClientQueue);
             }
+            Thread.Sleep(0);
         }
+        private int _ref = 0;
+        public void AddRef() => Interlocked.Increment(ref _ref);
+        public int GetRef() => Volatile.Read(ref _ref);
+        public void RemoveRef() => Interlocked.Decrement(ref _ref);
         public int Timeout { get; set; } = 10;
         public Client Take()
         {
             Client client = null;
             Monitor.Enter(ClientQueue);
-            while (true)
+            try
             {
-                bool timeout = false;
-                if (ClientQueue.IsEmpty)
+                while (true)
                 {
-                    timeout = !Monitor.Wait(ClientQueue, TimeSpan.FromSeconds(Timeout));
-                }
-                ClientQueue.TryDequeue(out client);
+                    bool timeout = false;
+                    if (ClientQueue.IsEmpty)
+                    {
+                        timeout = !Monitor.Wait(ClientQueue, TimeSpan.FromSeconds(Timeout));
+                    }
+                    ClientQueue.TryDequeue(out client);
 
-                if (client != null || timeout)
-                {
-                    break;
+                    if (client != null || timeout)
+                    {
+                        break;
+                    }
                 }
             }
-            Monitor.Exit(ClientQueue);
+            finally
+            {
+                Monitor.Exit(ClientQueue);
+            }
             if (client == null)
             {
-                throw new TimeoutException($"Connection pool is empty and wait time out({Timeout}s)!");
+                var reasonPhrase = $"Connection pool is empty and wait time out({Timeout}s)";
+                if (DiagnosticReporter != null)
+                {
+                    throw DiagnosticReporter.BuildDepletionException(reasonPhrase);
+                }
+                throw new TimeoutException(reasonPhrase);
             }
             return client;
         }
     }
+
+    internal interface IPoolDiagnosticReporter
+    {
+        SessionPoolDepletedException BuildDepletionException(string reasonPhrase);
+    }
 }
diff --git a/src/Apache.IoTDB/DataStructure/ByteBuffer.cs b/src/Apache.IoTDB/DataStructure/ByteBuffer.cs
@@ -18,7 +18,6 @@
  */
 
 using System;
-using System.Linq;
 using System.Text;
 
 namespace Apache.IoTDB.DataStructure
@@ -70,9 +69,12 @@ public bool GetBool()
         public int GetInt()
         {
             var intBuff = _buffer[_readPos..(_readPos + 4)];
-            if (_isLittleEndian) intBuff = intBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(intBuff);
+            }
 #if NET461_OR_GREATER || NETSTANDARD2_0
-            var intValue = BitConverter.ToInt32(intBuff,0);
+            var intValue = BitConverter.ToInt32(intBuff, 0);
 #else
             var intValue = BitConverter.ToInt32(intBuff);
 #endif
@@ -84,10 +86,12 @@ public int GetInt()
         public long GetLong()
         {
             var longBuff = _buffer[_readPos..(_readPos + 8)];
-
-            if (_isLittleEndian) longBuff = longBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(longBuff);
+            }
 #if NET461_OR_GREATER || NETSTANDARD2_0
-            var longValue = BitConverter.ToInt64(longBuff,0);
+            var longValue = BitConverter.ToInt64(longBuff, 0);
 #else
             var longValue = BitConverter.ToInt64(longBuff);
 #endif
@@ -99,10 +103,12 @@ public long GetLong()
         public float GetFloat()
         {
             var floatBuff = _buffer[_readPos..(_readPos + 4)];
-
-            if (_isLittleEndian) floatBuff = floatBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(floatBuff);
+            }
 #if NET461_OR_GREATER || NETSTANDARD2_0
-            var floatValue = BitConverter.ToSingle(floatBuff,0);
+            var floatValue = BitConverter.ToSingle(floatBuff, 0);
 #else
             var floatValue = BitConverter.ToSingle(floatBuff);
 #endif
@@ -113,10 +119,12 @@ public float GetFloat()
         public double GetDouble()
         {
             var doubleBuff = _buffer[_readPos..(_readPos + 8)];
-
-            if (_isLittleEndian) doubleBuff = doubleBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(doubleBuff);
+            }
 #if NET461_OR_GREATER || NETSTANDARD2_0
-            var doubleValue = BitConverter.ToDouble(doubleBuff,0);
+            var doubleValue = BitConverter.ToDouble(doubleBuff, 0);
 #else
             var doubleValue = BitConverter.ToDouble(doubleBuff);
 #endif
@@ -162,8 +170,10 @@ private void ExtendBuffer(int spaceNeed)
         public void AddBool(bool value)
         {
             var boolBuffer = BitConverter.GetBytes(value);
-
-            if (_isLittleEndian) boolBuffer = boolBuffer.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(boolBuffer);
+            }
 
             ExtendBuffer(boolBuffer.Length);
             boolBuffer.CopyTo(_buffer, _writePos);
@@ -173,8 +183,10 @@ public void AddBool(bool value)
         public void AddInt(int value)
         {
             var intBuff = BitConverter.GetBytes(value);
-
-            if (_isLittleEndian) intBuff = intBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(intBuff);
+            }
 
             ExtendBuffer(intBuff.Length);
             intBuff.CopyTo(_buffer, _writePos);
@@ -184,8 +196,10 @@ public void AddInt(int value)
         public void AddLong(long value)
         {
             var longBuff = BitConverter.GetBytes(value);
-
-            if (_isLittleEndian) longBuff = longBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(longBuff);
+            }
 
             ExtendBuffer(longBuff.Length);
             longBuff.CopyTo(_buffer, _writePos);
@@ -195,8 +209,10 @@ public void AddLong(long value)
         public void AddFloat(float value)
         {
             var floatBuff = BitConverter.GetBytes(value);
-
-            if (_isLittleEndian) floatBuff = floatBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(floatBuff);
+            }
 
             ExtendBuffer(floatBuff.Length);
             floatBuff.CopyTo(_buffer, _writePos);
@@ -206,8 +222,10 @@ public void AddFloat(float value)
         public void AddDouble(double value)
         {
             var doubleBuff = BitConverter.GetBytes(value);
-
-            if (_isLittleEndian) doubleBuff = doubleBuff.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(doubleBuff);
+            }
 
             ExtendBuffer(doubleBuff.Length);
             doubleBuff.CopyTo(_buffer, _writePos);
@@ -237,8 +255,10 @@ public void AddBinary(byte[] value)
         public void AddChar(char value)
         {
             var charBuf = BitConverter.GetBytes(value);
-
-            if (_isLittleEndian) charBuf = charBuf.Reverse().ToArray();
+            if (_isLittleEndian)
+            {
+                Array.Reverse(charBuf);
+            }
 
             ExtendBuffer(charBuf.Length);
             charBuf.CopyTo(_buffer, _writePos);
diff --git a/src/Apache.IoTDB/DataStructure/SessionDataSet.cs b/src/Apache.IoTDB/DataStructure/SessionDataSet.cs
@@ -52,7 +52,22 @@ public class SessionDataSet : System.IDisposable
         private int Flag => 0x80;
         private int DefaultTimeout => 10000;
         public int FetchSize { get; set; }
+
+        /// <summary>
+        /// Gets the number of rows in the current fetched batch.
+        /// Note: This is NOT the total row count of the query result. Use HasNext() to check for more data.
+        /// </summary>
+        /// <returns>The number of rows in the current batch.</returns>
+        public int CurrentBatchRowCount() => _currentBatchRowCount;
+
+        /// <summary>
+        /// Gets the number of rows in the current fetched batch.
+        /// </summary>
+        /// <returns>The number of rows in the current batch.</returns>
+        [Obsolete("Use CurrentBatchRowCount() instead. This property returns batch size, not total row count.")]
         public int RowCount { get; set; }
+
+        private int _currentBatchRowCount;
         public SessionDataSet(string sql, TSExecuteStatementResp resp, Client client, ConcurrentClientQueue clientQueue, long statementId)
         {
             _clientQueue = clientQueue;
@@ -74,7 +89,8 @@ public SessionDataSet(string sql, TSExecuteStatementResp resp, Client client, Co
             // some internal variable
             _hasCatchedResult = false;
             _rowIndex = 0;
-            RowCount = _queryDataset.Time.Length / sizeof(long);
+            _currentBatchRowCount = _queryDataset.Time.Length / sizeof(long);
+            RowCount = _currentBatchRowCount;
 
             _columnNames = resp.Columns;
             _columnTypeLst = resp.DataTypeList;
diff --git a/src/Apache.IoTDB/PoolHealthMetrics.cs b/src/Apache.IoTDB/PoolHealthMetrics.cs
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+using System.Threading;
+
+namespace Apache.IoTDB
+{
+    /// <summary>
+    /// Encapsulates real-time health statistics for connection pool monitoring.
+    /// Thread-safe implementation for concurrent access patterns.
+    /// </summary>
+    internal class PoolHealthMetrics
+    {
+        private int _reconnectionFailureTally;
+        private readonly int _configuredMaxSize;
+
+        public PoolHealthMetrics(int configuredMaxSize)
+        {
+            _configuredMaxSize = configuredMaxSize;
+        }
+
+        public void IncrementReconnectionFailures()
+        {
+            Interlocked.Increment(ref _reconnectionFailureTally);
+        }
+
+        public void ResetAllCounters()
+        {
+            Interlocked.Exchange(ref _reconnectionFailureTally, 0);
+        }
+
+        public int GetReconnectionFailureTally() => Volatile.Read(ref _reconnectionFailureTally);
+
+        public int GetConfiguredMaxSize() => _configuredMaxSize;
+    }
+}
diff --git a/src/Apache.IoTDB/ReconnectionFailedException.cs b/src/Apache.IoTDB/ReconnectionFailedException.cs
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+using System;
+using Thrift;
+
+namespace Apache.IoTDB
+{
+    /// <summary>
+    /// Exception thrown when all reconnection attempts to the server have failed.
+    /// This exception is used internally to distinguish reconnection failures from other errors.
+    /// </summary>
+    internal class ReconnectionFailedException : TException
+    {
+        internal ReconnectionFailedException(string message)
+            : base(message, null)
+        {
+        }
+
+        internal ReconnectionFailedException(string message, Exception innerException)
+            : base(message, innerException)
+        {
+        }
+    }
+}
diff --git a/src/Apache.IoTDB/SessionPool.cs b/src/Apache.IoTDB/SessionPool.cs
diff --git a/src/Apache.IoTDB/SessionPoolDepletedException.cs b/src/Apache.IoTDB/SessionPoolDepletedException.cs
