Add more fine-grained State to MySqlSession.

This will track that the connection to the MySQL Server is in a known state, and help implement cancellation.

Author: bgrainger

src/MySqlConnector/MySqlClient/CommandExecutors/TextCommandExecutor.cs

@@ -50,6 +50,7 @@ public virtual async Task<object> ExecuteScalarAsync(string commandText, MySqlPa
 		public virtual async Task<DbDataReader> ExecuteReaderAsync(string commandText, MySqlParameterCollection parameterCollection,
 			CommandBehavior behavior, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
+			m_command.Connection.Session.StartQuerying();
 			m_command.LastInsertedId = -1;
 			var statementPreparerOptions = StatementPreparerOptions.None;
 			if (m_command.Connection.AllowUserVariables || m_command.CommandType == CommandType.StoredProcedure)

src/MySqlConnector/MySqlClient/MySqlCommand.cs

@@ -177,8 +177,6 @@ private void VerifyValid()
 				throw new InvalidOperationException("The transaction associated with this command is not the connection's active transaction.");
 			if (string.IsNullOrWhiteSpace(CommandText))
 				throw new InvalidOperationException("CommandText must be specified");
-			if (Connection.ActiveReader != null)
-				throw new MySqlException("There is already an open DataReader associated with this Connection which must be closed first.");
 		}
 
 		internal void ReaderClosed() => (m_commandExecutor as StoredProcedureCommandExecutor)?.SetParams();

src/MySqlConnector/MySqlClient/MySqlConnection.cs

@@ -230,7 +230,7 @@ internal async Task<CachedProcedure> GetCachedProcedure(IOBehavior ioBehavior, s
 		}
 
 		internal MySqlTransaction CurrentTransaction { get; set; }
-		internal MySqlDataReader ActiveReader { get; set; }
+		internal MySqlDataReader ActiveReader => m_session.ActiveReader;
 		internal bool AllowUserVariables => m_connectionSettings.AllowUserVariables;
 		internal bool BufferResultSets => m_connectionSettings.BufferResultSets;
 		internal bool ConvertZeroDateTime => m_connectionSettings.ConvertZeroDateTime;
@@ -309,11 +309,8 @@ private void DoClose()
 		private void CloseDatabase()
 		{
 			m_cachedProcedures = null;
-			if (ActiveReader != null)
-			{
-				ActiveReader.Dispose();
-				ActiveReader = null;
-			}
+			if (Session.ActiveReader != null)
+				Session.ActiveReader.Dispose();
 			if (CurrentTransaction != null && m_session.IsConnected)
 			{
 				CurrentTransaction.Dispose();

src/MySqlConnector/MySqlClient/MySqlDataReader.cs

@@ -215,14 +215,29 @@ protected override void Dispose(bool disposing)
 		internal static async Task<MySqlDataReader> CreateAsync(MySqlCommand command, CommandBehavior behavior, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
 			var dataReader = new MySqlDataReader(command, behavior);
-			await dataReader.ReadFirstResultSetAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
-			command.Connection.ActiveReader = dataReader;
-			if (command.Connection.BufferResultSets)
+			command.Connection.Session.SetActiveReader(dataReader);
+
+			try
+			{
+				await dataReader.ReadFirstResultSetAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
+				if (command.Connection.BufferResultSets)
+				{
+					while (await dataReader.BufferNextResultAsync(ioBehavior, cancellationToken).ConfigureAwait(false) != null)
+					{
+					}
+				}
+				return dataReader;
+			}
+			catch (Exception)
+			{
+				dataReader.Dispose();
+				throw;
+			}
+			finally
 			{
-				while (await dataReader.BufferNextResultAsync(ioBehavior, cancellationToken).ConfigureAwait(false) != null);
-				command.Connection.ActiveReader = null;
+				if (command.Connection.BufferResultSets)
+					command.Connection.Session.FinishQuerying();
 			}
-			return dataReader;
 		}
 
 		internal async Task ReadFirstResultSetAsync(IOBehavior ioBehavior, CancellationToken cancellationToken)
@@ -250,8 +265,8 @@ private void DoClose()
 				m_nextResultSetBuffer.Clear();
 
 				var connection = Command.Connection;
-				if (connection.ActiveReader == this)
-					connection.ActiveReader = null;
+				if (!connection.BufferResultSets)
+					connection.Session.FinishQuerying();
 				Command.ReaderClosed();
 				if ((m_behavior & CommandBehavior.CloseConnection) != 0)
 				{

src/MySqlConnector/Serialization/MySqlSession.cs

@@ -23,6 +23,7 @@ public MySqlSession()
 
 		public MySqlSession(ConnectionPool pool, int poolGeneration)
 		{
+			m_lock = new object();
 			CreatedUtc = DateTime.UtcNow;
 			Pool = pool;
 			PoolGeneration = poolGeneration;
@@ -38,13 +39,59 @@ public MySqlSession(ConnectionPool pool, int poolGeneration)
 
 		public void ReturnToPool() => Pool?.Return(this);
 
-		public bool IsConnected => m_state == State.Connected;
+		public bool IsConnected
+		{
+			get
+			{
+				lock (m_lock)
+					return m_state == State.Connected;
+			}
+		}
+
+		public void StartQuerying()
+		{
+			lock (m_lock)
+			{
+				if (m_state == State.Querying)
+					throw new MySqlException("There is already an open DataReader associated with this Connection which must be closed first.");
+
+				VerifyState(State.Connected);
+				m_state = State.Querying;
+			}
+		}
+
+		public MySqlDataReader ActiveReader => m_activeReader;
+
+		public void SetActiveReader(MySqlDataReader dataReader)
+		{
+			VerifyState(State.Querying);
+			if (dataReader == null)
+				throw new ArgumentNullException(nameof(dataReader));
+			if (m_activeReader != null)
+				throw new InvalidOperationException("Can't replace active reader.");
+			m_activeReader = dataReader;
+		}
+
+		public void FinishQuerying()
+		{
+			lock (m_lock)
+			{
+				VerifyState(State.Querying);
+				m_state = State.Connected;
+				m_activeReader = null;
+			}
+		}
 
 		public async Task DisposeAsync(IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
 			if (m_payloadHandler != null)
 			{
 				// attempt to gracefully close the connection, ignoring any errors (it may have been closed already by the server, etc.)
+				lock (m_lock)
+				{
+					VerifyState(State.Connected, State.Failed);
+					m_state = State.Closing;
+				}
 				try
 				{
 					m_payloadHandler.StartNewConversation();
@@ -59,18 +106,28 @@ public async Task DisposeAsync(IOBehavior ioBehavior, CancellationToken cancella
 				}
 			}
 			ShutdownSocket();
-			m_state = State.Closed;
+			lock (m_lock)
+				m_state = State.Closed;
 		}
 
 		public async Task ConnectAsync(ConnectionSettings cs, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
+			lock (m_lock)
+			{
+				VerifyState(State.Created);
+				m_state = State.Connecting;
+			}
 			var connected = false;
 			if (cs.ConnectionType == ConnectionType.Tcp)
 				connected = await OpenTcpSocketAsync(cs, ioBehavior, cancellationToken).ConfigureAwait(false);
 			else if (cs.ConnectionType == ConnectionType.Unix)
 				connected = await OpenUnixSocketAsync(cs, ioBehavior, cancellationToken).ConfigureAwait(false);
 			if (!connected)
+			{
+				lock (m_lock)
+					m_state = State.Failed;
 				throw new MySqlException("Unable to connect to any of the specified MySQL hosts.");
+			}
 
 			var socketByteHandler = new SocketByteHandler(m_socket);
 			m_payloadHandler = new StandardPayloadHandler(socketByteHandler);
@@ -114,6 +171,7 @@ public async Task ConnectAsync(ConnectionSettings cs, IOBehavior ioBehavior, Can
 
 		public async Task ResetConnectionAsync(ConnectionSettings cs, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
+			VerifyState(State.Connected);
 			if (ServerVersion.Version.CompareTo(ServerVersions.SupportsResetConnection) >= 0)
 			{
 				await SendAsync(ResetConnectionPayload.Create(), ioBehavior, cancellationToken).ConfigureAwait(false);
@@ -150,6 +208,8 @@ public async Task ResetConnectionAsync(ConnectionSettings cs, IOBehavior ioBehav
 
 		public async Task<bool> TryPingAsync(IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
+			VerifyState(State.Connected);
+
 			// check if client socket is still connected
 			// http://stackoverflow.com/questions/2661764/how-to-check-if-a-socket-is-connected-disconnected-in-c
 			if (m_socket.Poll(1, SelectMode.SelectRead) && m_socket.Available == 0)
@@ -169,6 +229,7 @@ public async Task<bool> TryPingAsync(IOBehavior ioBehavior, CancellationToken ca
 			{
 			}
 
+			VerifyState(State.Failed);
 			return false;
 		}
 
@@ -196,10 +257,13 @@ public ValueTask<int> SendReplyAsync(PayloadData payload, IOBehavior ioBehavior,
 
 		private void VerifyConnected()
 		{
-			if (m_state == State.Closed)
-				throw new ObjectDisposedException(nameof(MySqlSession));
-			if (m_state != State.Connected)
-				throw new InvalidOperationException("MySqlSession is not connected.");
+			lock (m_lock)
+			{
+				if (m_state == State.Closed)
+					throw new ObjectDisposedException(nameof(MySqlSession));
+				if (m_state != State.Connected && m_state != State.Querying && m_state != State.CancelingQuery && m_state != State.Closing)
+					throw new InvalidOperationException("MySqlSession is not connected.");
+			}
 		}
 
 		private async Task<bool> OpenTcpSocketAsync(ConnectionSettings cs, IOBehavior ioBehavior, CancellationToken cancellationToken)
@@ -270,7 +334,8 @@ private async Task<bool> OpenTcpSocketAsync(ConnectionSettings cs, IOBehavior io
 					m_socket = m_tcpClient.Client;
 					m_networkStream = m_tcpClient.GetStream();
 					SerializationUtility.SetKeepalive(m_socket, cs.Keepalive);
-					m_state = State.Connected;
+					lock (m_lock)
+						m_state = State.Connected;
 					return true;
 				}
 			}
@@ -316,7 +381,8 @@ private async Task<bool> OpenUnixSocketAsync(ConnectionSettings cs, IOBehavior i
 				m_socket = socket;
 				m_networkStream = new NetworkStream(socket);
 
-				m_state = State.Connected;
+				lock (m_lock)
+					m_state = State.Connected;
 				return true;
 			}
 
@@ -398,7 +464,8 @@ private async Task InitSslAsync(ProtocolCapabilities serverCapabilities, Connect
 			{
 				ShutdownSocket();
 				m_hostname = "";
-				m_state = State.Failed;
+				lock (m_lock)
+					m_state = State.Failed;
 				if (ex is AuthenticationException)
 					throw new MySqlException("SSL Authentication Error", ex);
 				if (ex is IOException && clientCertificates != null)
@@ -450,10 +517,10 @@ private void ShutdownSocket()
 
 		private ValueTask<int> TryAsync<TArg>(Func<TArg, IOBehavior, ValueTask<int>> func, TArg arg, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
-			VerifyConnected();
 			ValueTask<int> task;
 			try
 			{
+				VerifyConnected();
 				task = func(arg, ioBehavior);
 			}
 			catch (Exception ex)
@@ -480,10 +547,10 @@ private int TryAsyncContinuation(Task<int> task)
 
 		private ValueTask<PayloadData> TryAsync(Func<ProtocolErrorBehavior, IOBehavior, ValueTask<ArraySegment<byte>>> func, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
-			VerifyConnected();
 			ValueTask<ArraySegment<byte>> task;
 			try
 			{
+				VerifyConnected();
 				task = func(ProtocolErrorBehavior.Throw, ioBehavior);
 			}
 			catch (Exception ex)
@@ -516,22 +583,57 @@ private PayloadData TryAsyncContinuation(Task<ArraySegment<byte>> task)
 
 		private void SetFailed()
 		{
-			m_state = State.Failed;
+			lock (m_lock)
+				m_state = State.Failed;
+		}
+
+
+		private void VerifyState(State state)
+		{
+			if (m_state != state)
+				throw new InvalidOperationException("Expected state to be {0} but was {1}.".FormatInvariant(state, m_state));
+		}
+
+		private void VerifyState(State state1, State state2)
+		{
+			if (m_state != state1 && m_state != state2)
+				throw new InvalidOperationException("Expected state to be ({0}|{1}) but was {2}.".FormatInvariant(state1, state2, m_state));
 		}
 
 		private enum State
 		{
+			// The session has been created; no connection has been made.
 			Created,
+
+			// The session is attempting to connect to a server.
+			Connecting,
+
+			// The session is connected to a server; there is no active query.
 			Connected,
+
+			// The session is connected to a server and a query is being made.
+			Querying,
+
+			// The session is connected to a server and the active query is being cancelled.
+			CancelingQuery,
+
+			// The session is closing.
+			Closing,
+
+			// The session is closed.
 			Closed,
+
+			// An unexpected error occurred; the session is in an unusable state.
 			Failed,
 		}
 
+		readonly object m_lock;
 		State m_state;
 		string m_hostname = "";
 		TcpClient m_tcpClient;
 		Socket m_socket;
 		NetworkStream m_networkStream;
 		IPayloadHandler m_payloadHandler;
+		MySqlDataReader m_activeReader;
 	}
 }