Improve connection keepalive, disconnect, and connection loss after reviewing crash logs and testing a number of situations:

 - Improve stability of closing connections after a connection loss
 - Minimise prompting a user for connection state restore if closing windows/tabs
 - Allow cancellation of keepalive ping threads to prevent crashes after deallocation of parent
 - Manually handle ping thread state struct memory to avoid cross-thread deallocation issues
 - Improve disconnection speed and resilience

4 files changed, 87 insertions, 20 deletions

diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.h b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.h
index e84c4ca6..0bbac2d2 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.h
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.h
@@ -52,4 +52,7 @@ void _backgroundPingTask(void *ptr);
 void _forceThreadExit(int signalNumber);
 void _pingThreadCleanup(void *pingDetails);
 
+// Cancellation
+- (void)_cancelKeepAlives;
+
 @end
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.m
index 9a28bf57..fb8a984f 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection Categories/Ping & KeepAlive.m
@@ -49,7 +49,8 @@
 - (void)_keepAlive
 {
 
-	// Do nothing if not connected or if keepalive is disabled
+	// Do nothing if not connected, if keepalive is disabled, or a keepalive is in
+	// progress.
 	if (state != SPMySQLConnected || !useKeepAlive) return;
 
 	// Check to see whether a ping is required.  First, compare the last query
@@ -81,6 +82,7 @@
  */
 - (void)_threadedKeepAlive
 {
+	keepAliveThread = [NSThread currentThread];
 
 	// If the maximum number of ping failures has been reached, determine whether to reconnect.
 	if (keepAliveLastPingBlocked || keepAlivePingFailures >= 3) {
@@ -97,6 +99,7 @@
 		}
 
 		// Return as no further ping action required this cycle.
+		keepAliveThread = nil;
 		return;
 	}
 
@@ -107,6 +110,7 @@
 	} else {
 		keepAlivePingFailures++;
 	}
+	keepAliveThread = nil;
 }
 
 #pragma mark -
@@ -142,16 +146,16 @@
 	if (timeout > 0) pingTimeout = timeout;
 
 	// Set up a struct containing details the ping task will need
-	SPMySQLConnectionPingDetails pingDetails;
-	pingDetails.mySQLConnection = mySQLConnection;
-	pingDetails.keepAliveLastPingSuccessPointer = &keepAliveLastPingSuccess;
-	pingDetails.keepAlivePingActivePointer = &keepAlivePingThreadActive;
+	SPMySQLConnectionPingDetails *pingDetails = malloc(sizeof(SPMySQLConnectionPingDetails));
+	pingDetails->mySQLConnection = mySQLConnection;
+	pingDetails->keepAliveLastPingSuccessPointer = &keepAliveLastPingSuccess;
+	pingDetails->keepAlivePingActivePointer = &keepAlivePingThreadActive;
 
 	// Create a pthread for the ping
 	pthread_attr_t attr;
 	pthread_attr_init(&attr);
 	pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
-	pthread_create(&keepAlivePingThread, &attr, (void *)&_backgroundPingTask, &pingDetails);
+	pthread_create(&keepAlivePingThread_t, &attr, (void *)&_backgroundPingTask, pingDetails);
 
 	// Record the ping start time
 	pingStartTime_t = mach_absolute_time();
@@ -161,25 +165,29 @@
 		usleep((useconds_t)loopDelay);
 		pingElapsedTime = _elapsedSecondsSinceAbsoluteTime(pingStartTime_t);
 
-		// If the ping timeout has been exceeded, force a timeout; double-check that the
-		// thread is still active.
-		if (pingElapsedTime > pingTimeout && keepAlivePingThreadActive && !threadCancelled) {
-			pthread_cancel(keepAlivePingThread);
+		// If the ping timeout has been exceeded, or the ping thread has been
+		// cancelled, force a timeout; double-check that the thread is still active.
+		if (([keepAliveThread isCancelled] || pingElapsedTime > pingTimeout)
+			&& keepAlivePingThreadActive
+			&& !threadCancelled)
+		{
+			pthread_cancel(keepAlivePingThread_t);
 			threadCancelled = YES;
 
 		// If the timeout has been exceeded by an additional two seconds, and the thread is
 		// still active, kill the thread.  This can occur in certain network conditions causing
 		// a blocking read.
 		} else if (pingElapsedTime > (pingTimeout + 2) && keepAlivePingThreadActive) {
-			pthread_kill(keepAlivePingThread, SIGUSR1);	
+			pthread_kill(keepAlivePingThread_t, SIGUSR1);	
 			keepAlivePingThreadActive = NO;
 			keepAliveLastPingBlocked = YES;
 		}
 	} while (keepAlivePingThreadActive);
 
 	// Clean up
-	keepAlivePingThread = NULL;
+	keepAlivePingThread_t = NULL;
 	pthread_attr_destroy(&attr);
+	free(pingDetails);
 
     // Unlock the connection
 	[self _unlockConnection];
@@ -221,6 +229,11 @@ void _forceThreadExit(int signalNumber)
 	pthread_exit(NULL);
 }
 
+/**
+ * A thread cleanup routine.  This is added to the thread using a
+ * pthread_cleanup_push call; a pthread_exit or a pthread_cleanup_pop
+ * both execute this function.
+ */
 void _pingThreadCleanup(void *pingDetails)
 {
 	SPMySQLConnectionPingDetails *pingDetailsStruct = pingDetails;
@@ -230,4 +243,30 @@ void _pingThreadCleanup(void *pingDetails)
 	mysql_thread_end();
 }
 
+#pragma mark -
+#pragma mark Cancellation
+
+/**
+ * If a keepalive thread is active, cancel it, and wait a short time for it
+ * to exit.
+ */
+- (void)_cancelKeepAlives
+{
+
+	// If no keepalive thread is active, return
+	if (!keepAliveThread) {
+		return;
+	}
+
+	// Mark the thread as cancelled
+	[keepAliveThread cancel];
+
+	// Wait inside a time limit of ten seconds for it to exit
+	uint64_t threadCancelStartTime_t = mach_absolute_time();
+	do {
+		usleep(100000);
+		if (_elapsedSecondsSinceAbsoluteTime(threadCancelStartTime_t) > 10) break;
+	} while (keepAliveThread);
+}
+
 @end
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.h b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.h
index a44ae46f..3400ecfd 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.h
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.h
@@ -86,7 +86,8 @@
 	CGFloat keepAliveInterval;
 	uint64_t lastKeepAliveTime;
 	NSUInteger keepAlivePingFailures;
-	pthread_t keepAlivePingThread;
+	NSThread *keepAliveThread;
+	pthread_t keepAlivePingThread_t;
 	BOOL keepAlivePingThreadActive;
 	BOOL keepAliveLastPingSuccess;
 	BOOL keepAliveLastPingBlocked;
diff --git a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
index 26276bda..4dbed002 100644
--- a/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
+++ b/Frameworks/SPMySQLFramework/Source/SPMySQLConnection.m
@@ -131,7 +131,8 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 		keepAliveInterval = 60;
 		keepAlivePingFailures = 0;
 		lastKeepAliveTime = 0;
-		keepAlivePingThread = NULL;
+		keepAliveThread = nil;
+		keepAlivePingThread_t = NULL;
 		keepAlivePingThreadActive = NO;
 		keepAliveLastPingSuccess = NO;
 		keepAliveLastPingBlocked = NO;
@@ -206,6 +207,9 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 	[keepAliveTimer invalidate];
 	[keepAliveTimer release];
 
+	// If a keepalive thread is active, cancel it
+	[self _cancelKeepAlives];
+
 	// Disconnect if appropriate (which should also disconnect any proxy)
 	[self disconnect];
 
@@ -326,8 +330,16 @@ 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;
+	if (state != SPMySQLConnected && state != SPMySQLConnecting) {
+		return;
+	}
+
 	state = SPMySQLDisconnecting;
 
 	// If a query is active, cancel it
@@ -335,12 +347,12 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 
 	// Allow any pings or cancelled queries  to complete, inside a time limit of ten seconds
 	uint64_t disconnectStartTime_t = mach_absolute_time();
-	do {
+	while (![self _tryLockConnection]) {
 		usleep(100000);
 		if (_elapsedSecondsSinceAbsoluteTime(disconnectStartTime_t) > 10) break;
-	} while (![self _tryLockConnection]);
+	}
 	[self _unlockConnection];
-	if (keepAlivePingThread != NULL) pthread_cancel(keepAlivePingThread), keepAlivePingThread = NULL;
+	[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
@@ -499,6 +511,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
  */
 - (MYSQL *)_makeRawMySQLConnectionWithEncoding:(NSString *)encodingName isMasterConnection:(BOOL)isMaster
 {
+	if ([[NSThread currentThread] isCancelled]) return NULL;
 
 	// Set up the MySQL connection object
 	MYSQL *theConnection = mysql_init(NULL);
@@ -625,6 +638,11 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 		}
 	}
 
+	if ([[NSThread currentThread] isCancelled]) {
+		[reconnectionPool release];
+		return NO;
+	}
+
 	isReconnecting = YES;
 
 	// Store certain details about the connection, so that if the reconnection is successful
@@ -648,6 +666,12 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 	// If no network is present, wait for a short time for one to become available
 	[self _waitForNetworkConnectionWithTimeout:10];
 
+	if ([[NSThread currentThread] isCancelled]) {
+		isReconnecting = NO;
+		[reconnectionPool release];
+		return NO;
+	}
+
 	// If there is a proxy, attempt to reconnect it in blocking fashion
 	if (proxy) {
 		uint64_t loopIterationStart_t, proxyWaitStart_t;
@@ -719,7 +743,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 	}
 
 	// If the reconnection succeeded, restore the connection state as appropriate
-	if (state == SPMySQLConnected) {
+	if (state == SPMySQLConnected && ![[NSThread currentThread] isCancelled]) {
 		reconnectSucceeded = YES;
 		if (databaseToRestore) {
 			[self selectDatabase:databaseToRestore];
@@ -733,7 +757,7 @@ const char *SPMySQLSSLPermissibleCiphers = "DHE-RSA-AES256-SHA:AES256-SHA:DHE-RS
 
 	// If the connection failed and the connection is permitted to retry,
 	// then retry the reconnection.
-	} else if (canRetry) {
+	} else if (canRetry && ![[NSThread currentThread] isCancelled]) {
 
 		// Default to attempting another reconnect
 		SPMySQLConnectionLostDecision connectionLostDecision = SPMySQLConnectionLostReconnect;