2 files changed, 50 insertions, 3 deletions

diff --git a/include/sound/compress_driver.h b/include/sound/compress_driver.h
index 9031a26..175ab32 100644
--- a/include/sound/compress_driver.h
+++ b/include/sound/compress_driver.h
@@ -48,6 +48,8 @@ struct snd_compr_ops;
  *	the ring buffer
  * @total_bytes_transferred: cumulative bytes transferred by offload DSP
  * @sleep: poll sleep
+ * @wait: drain wait queue
+ * @drain_wake: condition for drain wake
  */
 struct snd_compr_runtime {
 	snd_pcm_state_t state;
@@ -59,6 +61,8 @@ struct snd_compr_runtime {
 	u64 total_bytes_available;
 	u64 total_bytes_transferred;
 	wait_queue_head_t sleep;
+	wait_queue_head_t wait;
+	unsigned int drain_wake;
 	void *private_data;
 };
 
@@ -171,4 +175,12 @@ static inline void snd_compr_fragment_elapsed(struct snd_compr_stream *stream)
 	wake_up(&stream->runtime->sleep);
 }
 
+static inline void snd_compr_drain_notify(struct snd_compr_stream *stream)
+{
+	snd_BUG_ON(!stream);
+
+	stream->runtime->drain_wake = 1;
+	wake_up(&stream->runtime->wait);
+}
+
 #endif
diff --git a/sound/core/compress_offload.c b/sound/core/compress_offload.c
index bea523a..3eb47d0 100644
--- a/sound/core/compress_offload.c
+++ b/sound/core/compress_offload.c
@@ -123,6 +123,7 @@ static int snd_compr_open(struct inode *inode, struct file *f)
 	}
 	runtime->state = SNDRV_PCM_STATE_OPEN;
 	init_waitqueue_head(&runtime->sleep);
+	init_waitqueue_head(&runtime->wait);
 	data->stream.runtime = runtime;
 	f->private_data = (void *)data;
 	mutex_lock(&compr->lock);
@@ -682,12 +683,34 @@ static int snd_compr_stop(struct snd_compr_stream *stream)
 	if (!retval) {
 		stream->runtime->state = SNDRV_PCM_STATE_SETUP;
 		wake_up(&stream->runtime->sleep);
+		snd_compr_drain_notify(stream);
 		stream->runtime->total_bytes_available = 0;
 		stream->runtime->total_bytes_transferred = 0;
 	}
 	return retval;
 }
 
+static int snd_compress_wait_for_drain(struct snd_compr_stream *stream)
+{
+	/*
+	 * We are called with lock held. So drop the lock while we wait for
+	 * drain complete notfication from the driver
+	 *
+	 * It is expected that driver will notify the drain completion and then
+	 * stream will be moved to SETUP state, even if draining resulted in an
+	 * error. We can trigger next track after this.
+	 */
+	stream->runtime->state = SNDRV_PCM_STATE_DRAINING;
+	mutex_unlock(&stream->device->lock);
+
+	wait_event(stream->runtime->wait, stream->runtime->drain_wake);
+
+	wake_up(&stream->runtime->sleep);
+	mutex_lock(&stream->device->lock);
+
+	return 0;
+}
+
 static int snd_compr_drain(struct snd_compr_stream *stream)
 {
 	int retval;
@@ -695,11 +718,17 @@ static int snd_compr_drain(struct snd_compr_stream *stream)
 	if (stream->runtime->state == SNDRV_PCM_STATE_PREPARED ||
 			stream->runtime->state == SNDRV_PCM_STATE_SETUP)
 		return -EPERM;
+
+	stream->runtime->drain_wake = 0;
 	retval = stream->ops->trigger(stream, SND_COMPR_TRIGGER_DRAIN);
-	if (!retval) {
-		stream->runtime->state = SNDRV_PCM_STATE_DRAINING;
+	if (retval) {
+		pr_err("SND_COMPR_TRIGGER_DRAIN failed %d\n", retval);
 		wake_up(&stream->runtime->sleep);
+		return retval;
 	}
+
+	retval = snd_compress_wait_for_drain(stream);
+	stream->runtime->state = SNDRV_PCM_STATE_SETUP;
 	return retval;
 }
 
@@ -735,10 +764,16 @@ static int snd_compr_partial_drain(struct snd_compr_stream *stream)
 	if (stream->next_track == false)
 		return -EPERM;
 
+	stream->runtime->drain_wake = 0;
 	retval = stream->ops->trigger(stream, SND_COMPR_TRIGGER_PARTIAL_DRAIN);
+	if (retval) {
+		pr_err("Partial drain returned failure\n");
+		wake_up(&stream->runtime->sleep);
+		return retval;
+	}
 
 	stream->next_track = false;
-	return retval;
+	return snd_compress_wait_for_drain(stream);
 }
 
 static long snd_compr_ioctl(struct file *f, unsigned int cmd, unsigned long arg)