Merge "diag: Release wakeup sources properly"
diff --git a/drivers/char/diag/diagchar_core.c b/drivers/char/diag/diagchar_core.c
index f4a5a5b..7ef4768 100644
--- a/drivers/char/diag/diagchar_core.c
+++ b/drivers/char/diag/diagchar_core.c
@@ -297,10 +297,12 @@
 #ifdef CONFIG_DIAG_OVER_USB
 	/* If the SD logging process exits, change logging to USB mode */
 	if (driver->logging_process_id == current->tgid) {
+		mutex_lock(&driver->diagchar_mutex);
 		driver->logging_mode = USB_MODE;
+		diag_ws_reset();
+		mutex_unlock(&driver->diagchar_mutex);
 		diag_update_proc_vote(DIAG_PROC_MEMORY_DEVICE, VOTE_DOWN);
 		diagfwd_connect();
-		diag_ws_reset();
 #ifdef CONFIG_DIAGFWD_BRIDGE_CODE
 		diag_clear_hsic_tbl();
 		diagfwd_cancel_hsic(REOPEN_HSIC);
@@ -1265,10 +1267,16 @@
 					COPY_USER_SPACE_OR_EXIT(buf+ret,
 						*(data->buf_in_1),
 						data->write_ptr_1->length);
+					diag_ws_on_copy();
+					copy_data = 1;
 					data->in_busy_1 = 0;
 				}
 			}
 		}
+		if (!copy_data) {
+			diag_ws_on_copy();
+			copy_data = 1;
+		}
 #ifdef CONFIG_DIAG_SDIO_PIPE
 		/* copy 9K data over SDIO */
 		if (driver->in_busy_sdio == 1) {
@@ -1432,6 +1440,7 @@
 exit:
 	mutex_unlock(&driver->diagchar_mutex);
 	if (copy_data) {
+		diag_ws_on_copy_complete();
 		/*
 		 * Flush any work that is currently pending on the data
 		 * channels. This will ensure that the next read is not
@@ -1440,7 +1449,6 @@
 		for (i = 0; i < NUM_SMD_DATA_CHANNELS; i++)
 			flush_workqueue(driver->smd_data[i].wq);
 		wake_up(&driver->smd_wait_q);
-		diag_ws_on_copy_complete();
 	}
 	return ret;
 }
diff --git a/drivers/char/diag/diagfwd.c b/drivers/char/diag/diagfwd.c
index edf0216..9f2d5d3 100644
--- a/drivers/char/diag/diagfwd.c
+++ b/drivers/char/diag/diagfwd.c
@@ -531,8 +531,11 @@
 
 	return 0;
 err:
-	if (driver->logging_mode == MEMORY_DEVICE_MODE)
+	if ((smd_info->type == SMD_DATA_TYPE ||
+	     smd_info->type == SMD_CMD_TYPE) &&
+	     driver->logging_mode == MEMORY_DEVICE_MODE)
 		diag_ws_on_read(0);
+
 	return 0;
 }
 
@@ -541,6 +544,12 @@
 	if (!smd_info || !smd_info->ch)
 		return;
 
+	if ((smd_info->type == SMD_DATA_TYPE ||
+	     smd_info->type == SMD_CMD_TYPE) &&
+	     driver->logging_mode == MEMORY_DEVICE_MODE) {
+		diag_ws_on_notify();
+	}
+
 	switch (smd_info->type) {
 	case SMD_DCI_TYPE:
 	case SMD_DCI_CMD_TYPE:
@@ -559,12 +568,12 @@
 	default:
 		pr_err("diag: In %s, invalid type: %d\n", __func__,
 			smd_info->type);
+		if ((smd_info->type == SMD_DATA_TYPE ||
+		     smd_info->type == SMD_CMD_TYPE) &&
+		     driver->logging_mode == MEMORY_DEVICE_MODE)
+			diag_ws_on_read(0);
 		return;
 	}
-
-	if (driver->logging_mode == MEMORY_DEVICE_MODE &&
-	    smd_info->type == SMD_DATA_TYPE)
-		diag_ws_on_notify();
 }
 
 static int diag_smd_resize_buf(struct diag_smd_info *smd_info, void **buf,
@@ -787,9 +796,11 @@
 				}
 			}
 		}
-		if (smd_info->type == SMD_DATA_TYPE &&
-		    driver->logging_mode == MEMORY_DEVICE_MODE)
-			diag_ws_on_read(pkt_len);
+
+		if ((smd_info->type == SMD_DATA_TYPE ||
+		     smd_info->type == SMD_CMD_TYPE) &&
+		     driver->logging_mode == MEMORY_DEVICE_MODE)
+			diag_ws_on_read(total_recd);
 
 		if (total_recd > 0) {
 			if (!buf) {
@@ -814,12 +825,19 @@
 	} else if (smd_info->ch && !buf &&
 		(driver->logging_mode == MEMORY_DEVICE_MODE)) {
 			chk_logging_wakeup();
+	} else {
+		if ((smd_info->type == SMD_DATA_TYPE ||
+		     smd_info->type == SMD_CMD_TYPE) &&
+		     driver->logging_mode == MEMORY_DEVICE_MODE) {
+			diag_ws_on_read(0);
+		}
 	}
 	return;
 
 fail_return:
-	if (smd_info->type == SMD_DATA_TYPE &&
-	    driver->logging_mode == MEMORY_DEVICE_MODE)
+	if ((smd_info->type == SMD_DATA_TYPE ||
+	     smd_info->type == SMD_CMD_TYPE) &&
+	     driver->logging_mode == MEMORY_DEVICE_MODE)
 		diag_ws_on_read(0);
 
 	if (smd_info->type == SMD_DCI_TYPE ||
@@ -2214,11 +2232,11 @@
 			diag_dci_notify_client(smd_info->peripheral_mask,
 							DIAG_STATUS_OPEN);
 		}
-		wake_up(&driver->smd_wait_q);
 		diag_smd_queue_read(smd_info);
+		wake_up(&driver->smd_wait_q);
 	} else if (event == SMD_EVENT_DATA) {
-		wake_up(&driver->smd_wait_q);
 		diag_smd_queue_read(smd_info);
+		wake_up(&driver->smd_wait_q);
 		if (smd_info->type == SMD_DCI_TYPE ||
 		    smd_info->type == SMD_DCI_CMD_TYPE) {
 			diag_dci_try_activate_wakeup_source();