- In the SPMySQL.framework, separate framework-triggered connections and disconnections from external actions, and use that separation to perform safer disconnects

 - When closing a database document, add a new notification, and use that to resolve retain cycles affecting connection processes
 - Improve connection controller disconnection when the document is closed, fixing crashes, by building on those two features (addresses Issue #1396)
 - Use some of the new functionality to improve SSH and MySQL connection cancellation, making both cancelable in the interface and making both respond much more quickly

4 files changed, 146 insertions, 109 deletions

diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h b/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h
index 64fa0bc3..c3c57e3f 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQL Private APIs.h
@@ -41,10 +41,12 @@
 
 @interface SPMySQLConnection (PrivateAPI)
 
+- (BOOL)_connect;
 - (MYSQL *)_makeRawMySQLConnectionWithEncoding:(NSString *)encodingName isMasterConnection:(BOOL)isMaster;
 - (BOOL)_reconnectAllowingRetries:(BOOL)canRetry;
 - (BOOL)_reconnectAfterBackgroundConnectionLoss;
 - (BOOL)_waitForNetworkConnectionWithTimeout:(double)timeoutSeconds;
+- (void)_disconnect;
 - (void)_updateConnectionVariables;
 - (void)_restoreConnectionVariables;
 - (BOOL)_checkConnectionIfNecessary;
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m
index e0bfef52..867b7fba 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Max Packet Size.m
@@ -150,7 +150,7 @@
 		if (newSize != NSNotFound) {
 
 			// Successfully increased the global size - reconnect to use it, and return success
-			[self reconnect];
+			[self _reconnectAllowingRetries:YES];
 			return YES;
 		}
 	}
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m
index bafa1e6e..22e35648 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Querying & Preparation.m
@@ -222,6 +222,11 @@
 	lastQueryWasCancelled = NO;
 	lastQueryWasCancelledUsingReconnect = NO;
 
+	// If a disconnect was requested, cancel the action
+	if (userTriggeredDisconnect) {
+		return nil;
+	}
+
 	// Check the connection state - if no connection is available, log an
 	// error and return.
 	if (state == SPMySQLDisconnected || state == SPMySQLConnecting) {
@@ -368,6 +373,7 @@
 
 	// Unlock the connection if appropriate - if not a streaming result type.
 	if (![theResult isKindOfClass:[SPMySQLStreamingResult class]]) {
+		[self _tryLockConnection];
 		[self _unlockConnection];
 
 		// Also perform restore if appropriate
@@ -514,7 +520,6 @@
 	// so set up a new connection to run the KILL command.
 	MYSQL *killerConnection = [self _makeRawMySQLConnectionWithEncoding:@"utf8" isMasterConnection:NO];
 
-
 	// If the new connection was successfully set up, use it to run a KILL command.
 	if (killerConnection) {
 		NSStringEncoding aStringEncoding = [SPMySQLConnection stringEncodingForMySQLCharset:mysql_character_set_name(killerConnection)];
@@ -553,7 +558,7 @@
 		} else {
 			NSLog(@"SPMySQL Framework: query cancellation failed due to cancellation query error (status %d)", killQueryStatus);
 		}
-	} else {
+	} else if (!userTriggeredDisconnect) {
 		NSLog(@"SPMySQL Framework: query cancellation failed because connection failed");
 	}
 
@@ -565,13 +570,13 @@
 		return;
 	}
 
-	if (state == SPMySQLDisconnecting) return;
+	if (state == SPMySQLDisconnecting || state == SPMySQLDisconnected) return;
 
 	// Reset the connection with a reconnect.  Unlock the connection beforehand,
 	// to allow the reconnect, but lock it again afterwards to restore the expected
 	// state (query execution process should unlock as appropriate).
 	[self _unlockConnection];
-	[self reconnect];
+	[self _reconnectAllowingRetries:YES];
 	[self _lockConnection];
 
 	// Reset tracking bools to cover encompassed queries
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
index 1470d1bf..e12063a2 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
@@ -250,66 +250,8 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (BOOL)connect
 {
-
-	// If a connection is already active in some form, throw an exception
-	if (state != SPMySQLDisconnected && state != SPMySQLConnectionLostInBackground) {
-		[NSException raise:NSInternalInconsistencyException format:@"Attempted to connect a connection that is not disconnected (%d).", state];
-		return NO;
-	}
-	state = SPMySQLConnecting;
-
-	// Lock the connection for safety
-	[self _lockConnection];
-
-	// Attempt the connection
-	mySQLConnection = [self _makeRawMySQLConnectionWithEncoding:encoding isMasterConnection:YES];
-
-	// If the connection failed, reset state and return
-	if (!mySQLConnection) {
-		[self _unlockConnection];
-		state = SPMySQLDisconnected;
-		return NO;
-	}
-
-	// Successfully connected - record connected state and reset tracking variables
-	state = SPMySQLConnected;
 	userTriggeredDisconnect = NO;
-	initialConnectTime = mach_absolute_time();
-	mysqlConnectionThreadId = mySQLConnection->thread_id;
-	lastConnectionUsedTime = 0;
-
-	// Update SSL state
-	connectedWithSSL = NO;
-	if (useSSL) connectedWithSSL = (mysql_get_ssl_cipher(mySQLConnection))?YES:NO;
-	if (useSSL && !connectedWithSSL) {
-		if ([delegate respondsToSelector:@selector(connectionFellBackToNonSSL:)]) {
-			[delegate connectionFellBackToNonSSL:self];
-		}
-	}
-
-	// Reset keepalive variables
-	lastKeepAliveTime = 0;
-	keepAlivePingFailures = 0;
-
-	// Clear the connection error record
-	[self _updateLastErrorID:NSNotFound];
-	[self _updateLastErrorMessage:nil];
-
-	// Unlock the connection
-	[self _unlockConnection];
-
-	// Update connection variables to be in sync with the server state.  As this performs
-	// a query, ensure the connection is still up afterwards (!)
-	[self _updateConnectionVariables];
-	if (state != SPMySQLConnected) return NO;
-
-	// Now connection is established and verified, reset the counter
-	reconnectionRetryAttempts = 0;
-
-	// Update the maximum query size
-	[self _updateMaxQuerySize];
-
-	return YES;
+	return [self _connect];
 }
 
 /**
@@ -321,6 +263,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (BOOL)reconnect
 {
+	userTriggeredDisconnect = NO;
 	return [self _reconnectAllowingRetries:YES];
 }
 
@@ -329,49 +272,8 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (void)disconnect
 {
-
-	// If state is connection lost, set state directly to disconnected.
-	if (state == SPMySQLConnectionLostInBackground) {
-		state = SPMySQLDisconnected;
-	}
-
-	// Only continue if a connection is active
-	if (state != SPMySQLConnected && state != SPMySQLConnecting) {
-		return;
-	}
-
-	state = SPMySQLDisconnecting;
-
-	// If a query is active, cancel it
-	[self cancelCurrentQuery];
-
-	// Allow any pings or cancelled queries  to complete, inside a time limit of ten seconds
-	uint64_t disconnectStartTime_t = mach_absolute_time();
-	while (![self _tryLockConnection]) {
-		usleep(100000);
-		if (_elapsedSecondsSinceAbsoluteTime(disconnectStartTime_t) > 10) break;
-	}
-	[self _unlockConnection];
-	[self _cancelKeepAlives];
-
-	// Close the underlying MySQL connection if it still appears to be active, and not reading
-	// or writing.  While this may result in a leak of the MySQL object, it prevents crashes
-	// due to attempts to close a blocked/stuck connection.
-	if (!mySQLConnection->net.reading_or_writing && mySQLConnection->net.vio && mySQLConnection->net.buff) {
-		mysql_close(mySQLConnection);
-	}
-	mySQLConnection = NULL;
-
-	// If using a connection proxy, disconnect that too
-	if (proxy) {
-		[proxy performSelectorOnMainThread:@selector(disconnect) withObject:nil waitUntilDone:YES];
-	}
-
-	// Clear host-specific information
-	if (serverVersionString) [serverVersionString release], serverVersionString = nil;
-	if (database) [database release], database = nil;
-
-	state = SPMySQLDisconnected;
+	userTriggeredDisconnect = YES;
+	[self _disconnect];
 }
 
 #pragma mark -
@@ -433,7 +335,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 	// attempt to reconnect once, and if that fails will ask the user how to proceed - whether
 	// to keep reconnecting, or whether to disconnect.
 	if (!connectionVerified) {
-		connectionVerified = [self reconnect];
+		connectionVerified = [self _reconnectAllowingRetries:YES];
 	}
 
 	// Update the connection tracking use variable if the connection was confirmed,
@@ -506,6 +408,84 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 @implementation SPMySQLConnection (PrivateAPI)
 
 /**
+ * Handle a connection using previously set parameters, returning success or failure.
+ */
+- (BOOL)_connect
+{
+
+	// If a connection is already active in some form, throw an exception
+	if (state != SPMySQLDisconnected && state != SPMySQLConnectionLostInBackground) {
+		[NSException raise:NSInternalInconsistencyException format:@"Attempted to connect a connection that is not disconnected (%d).", state];
+		return NO;
+	}
+	state = SPMySQLConnecting;
+
+	if (userTriggeredDisconnect) {
+		return NO;
+	}
+
+	// Lock the connection for safety
+	[self _lockConnection];
+
+	// Attempt the connection
+	mySQLConnection = [self _makeRawMySQLConnectionWithEncoding:encoding isMasterConnection:YES];
+
+	// If the connection failed, reset state and return
+	if (!mySQLConnection) {
+		[self _unlockConnection];
+		state = SPMySQLDisconnected;
+		return NO;
+	}
+
+	// If the connection was cancelled, clean up and don't continue
+	if (userTriggeredDisconnect) {
+		mysql_close(mySQLConnection);
+		[self _unlockConnection];
+		mySQLConnection = NULL;
+		return NO;
+	}
+
+	// Successfully connected - record connected state and reset tracking variables
+	state = SPMySQLConnected;
+	initialConnectTime = mach_absolute_time();
+	mysqlConnectionThreadId = mySQLConnection->thread_id;
+	lastConnectionUsedTime = 0;
+
+	// Update SSL state
+	connectedWithSSL = NO;
+	if (useSSL) connectedWithSSL = (mysql_get_ssl_cipher(mySQLConnection))?YES:NO;
+	if (useSSL && !connectedWithSSL) {
+		if ([delegate respondsToSelector:@selector(connectionFellBackToNonSSL:)]) {
+			[delegate connectionFellBackToNonSSL:self];
+		}
+	}
+
+	// Reset keepalive variables
+	lastKeepAliveTime = 0;
+	keepAlivePingFailures = 0;
+
+	// Clear the connection error record
+	[self _updateLastErrorID:NSNotFound];
+	[self _updateLastErrorMessage:nil];
+
+	// Unlock the connection
+	[self _unlockConnection];
+
+	// Update connection variables to be in sync with the server state.  As this performs
+	// a query, ensure the connection is still up afterwards (!)
+	[self _updateConnectionVariables];
+	if (state != SPMySQLConnected) return NO;
+
+	// Now connection is established and verified, reset the counter
+	reconnectionRetryAttempts = 0;
+
+	// Update the maximum query size
+	[self _updateMaxQuerySize];
+
+	return YES;
+}
+
+/**
  * Make a connection using the class connection settings, returning a MySQL
  * connection object on success.
  */
@@ -740,7 +720,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 
 	// If not using a proxy, or if the proxy successfully connected, trigger a connection
 	if (!proxy || [proxy state] == SPMySQLProxyConnected) {
-		[self connect];
+		[self _connect];
 	}
 
 	// If the reconnection succeeded, restore the connection state as appropriate
@@ -856,6 +836,56 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 }
 
 /**
+ * Perform a disconnect of any active connections, cleaning up state to match.
+ */
+- (void)_disconnect
+{
+
+	// If state is connection lost, set state directly to disconnected.
+	if (state == SPMySQLConnectionLostInBackground) {
+		state = SPMySQLDisconnected;
+	}
+
+	// Only continue if a connection is active
+	if (state != SPMySQLConnected && state != SPMySQLConnecting) {
+		return;
+	}
+
+	state = SPMySQLDisconnecting;
+
+	// If a query is active, cancel it
+	[self cancelCurrentQuery];
+
+	// Allow any pings or cancelled queries  to complete, inside a time limit of ten seconds
+	uint64_t disconnectStartTime_t = mach_absolute_time();
+	while (![self _tryLockConnection]) {
+		usleep(100000);
+		if (_elapsedSecondsSinceAbsoluteTime(disconnectStartTime_t) > 10) break;
+	}
+	[self _unlockConnection];
+	[self _cancelKeepAlives];
+
+	// Close the underlying MySQL connection if it still appears to be active, and not reading
+	// or writing.  While this may result in a leak of the MySQL object, it prevents crashes
+	// due to attempts to close a blocked/stuck connection.
+	if (mySQLConnection && !mySQLConnection->net.reading_or_writing && mySQLConnection->net.vio && mySQLConnection->net.buff) {
+		mysql_close(mySQLConnection);
+	}
+	mySQLConnection = NULL;
+
+	// If using a connection proxy, disconnect that too
+	if (proxy) {
+		[proxy performSelectorOnMainThread:@selector(disconnect) withObject:nil waitUntilDone:YES];
+	}
+
+	// Clear host-specific information
+	if (serverVersionString) [serverVersionString release], serverVersionString = nil;
+	if (database) [database release], database = nil;
+
+	state = SPMySQLDisconnected;
+}
+
+/**
  * Update connection variables from the server, collecting state and ensuring
  * settings like encoding are in sync.
  */