Various improvements to a cleaner shutdown of the wifi display connection.

Change-Id: Id029a89939e53c2bd1d20e572d3975ec2795c239
related-to-bug: 7247918

13 files changed, 326 insertions, 172 deletions

diff --git a/include/media/stagefright/SurfaceMediaSource.h b/include/media/stagefright/SurfaceMediaSource.h
index 9e07ea4..f60a535 100644
--- a/include/media/stagefright/SurfaceMediaSource.h
+++ b/include/media/stagefright/SurfaceMediaSource.h
@@ -52,6 +52,8 @@ class GraphicBuffer;
 // may be dropped.  It is possible to wait for the buffers to be
 // returned (but not implemented)
 
+#define DEBUG_PENDING_BUFFERS   0
+
 class SurfaceMediaSource : public MediaSource,
                                 public MediaBufferObserver,
                                 protected BufferQueue::ConsumerListener {
@@ -169,6 +171,10 @@ private:
 
     size_t mNumPendingBuffers;
 
+#if DEBUG_PENDING_BUFFERS
+    Vector<MediaBuffer *> mPendingBuffers;
+#endif
+
     // mCurrentTimestamp is the timestamp for the current texture. It
     // gets set to mLastQueuedTimestamp each time updateTexImage is called.
     int64_t mCurrentTimestamp;
diff --git a/media/libstagefright/ACodec.cpp b/media/libstagefright/ACodec.cpp
index 8190498..af7a459 100644
--- a/media/libstagefright/ACodec.cpp
+++ b/media/libstagefright/ACodec.cpp
@@ -2471,6 +2471,9 @@ bool ACodec::BaseState::onOMXEmptyBufferDone(IOMX::buffer_id bufferID) {
         // by this "mediaBuffer" object. Now that the OMX component has
         // told us that it's done with the input buffer, we can decrement
         // the mediaBuffer's reference count.
+
+        ALOGV("releasing mbuf %p", mediaBuffer);
+
         ((MediaBuffer *)mediaBuffer)->release();
         mediaBuffer = NULL;
 
diff --git a/media/libstagefright/SurfaceMediaSource.cpp b/media/libstagefright/SurfaceMediaSource.cpp
index 3d3f421..9d39d0e 100644
--- a/media/libstagefright/SurfaceMediaSource.cpp
+++ b/media/libstagefright/SurfaceMediaSource.cpp
@@ -191,6 +191,23 @@ status_t SurfaceMediaSource::stop()
     ALOGV("stop");
     Mutex::Autolock lock(mMutex);
 
+    if (mStopped) {
+        return OK;
+    }
+
+    while (mNumPendingBuffers > 0) {
+        ALOGI("Still waiting for %d buffers to be returned.",
+                mNumPendingBuffers);
+
+#if DEBUG_PENDING_BUFFERS
+        for (size_t i = 0; i < mPendingBuffers.size(); ++i) {
+            ALOGI("%d: %p", i, mPendingBuffers.itemAt(i));
+        }
+#endif
+
+        mMediaBuffersAvailableCondition.wait(mMutex);
+    }
+
     mStopped = true;
     mFrameAvailableCondition.signal();
     mMediaBuffersAvailableCondition.signal();
@@ -335,6 +352,12 @@ status_t SurfaceMediaSource::read( MediaBuffer **buffer,
 
     ++mNumPendingBuffers;
 
+#if DEBUG_PENDING_BUFFERS
+    mPendingBuffers.push_back(*buffer);
+#endif
+
+    ALOGV("returning mbuf %p", *buffer);
+
     return OK;
 }
 
@@ -391,6 +414,15 @@ void SurfaceMediaSource::signalBufferReturned(MediaBuffer *buffer) {
         CHECK(!"signalBufferReturned: bogus buffer");
     }
 
+#if DEBUG_PENDING_BUFFERS
+    for (size_t i = 0; i < mPendingBuffers.size(); ++i) {
+        if (mPendingBuffers.itemAt(i) == buffer) {
+            mPendingBuffers.removeAt(i);
+            break;
+        }
+    }
+#endif
+
     --mNumPendingBuffers;
     mMediaBuffersAvailableCondition.broadcast();
 }
diff --git a/media/libstagefright/wifi-display/ANetworkSession.cpp b/media/libstagefright/wifi-display/ANetworkSession.cpp
index 90db758..0279c34 100644
--- a/media/libstagefright/wifi-display/ANetworkSession.cpp
+++ b/media/libstagefright/wifi-display/ANetworkSession.cpp
@@ -176,7 +176,7 @@ ANetworkSession::Session::Session(
 }
 
 ANetworkSession::Session::~Session() {
-    ALOGI("Session %d gone", mSessionID);
+    ALOGV("Session %d gone", mSessionID);
 
     close(mSocket);
     mSocket = -1;
@@ -1084,7 +1084,7 @@ void ANetworkSession::threadLoop() {
                 } else {
                     status_t err = session->readMore();
                     if (err != OK) {
-                        ALOGI("readMore on socket %d failed w/ error %d (%s)",
+                        ALOGE("readMore on socket %d failed w/ error %d (%s)",
                               s, err, strerror(-err));
                     }
                 }
@@ -1093,7 +1093,7 @@ void ANetworkSession::threadLoop() {
             if (FD_ISSET(s, &ws)) {
                 status_t err = session->writeMore();
                 if (err != OK) {
-                    ALOGI("writeMore on socket %d failed w/ error %d (%s)",
+                    ALOGE("writeMore on socket %d failed w/ error %d (%s)",
                           s, err, strerror(-err));
                 }
             }
diff --git a/media/libstagefright/wifi-display/source/Converter.cpp b/media/libstagefright/wifi-display/source/Converter.cpp
index c6118d4..0cdff6b 100644
--- a/media/libstagefright/wifi-display/source/Converter.cpp
+++ b/media/libstagefright/wifi-display/source/Converter.cpp
@@ -20,12 +20,15 @@
 
 #include "Converter.h"
 
+#include "MediaPuller.h"
+
 #include <cutils/properties.h>
 #include <gui/SurfaceTextureClient.h>
 #include <media/ICrypto.h>
 #include <media/stagefright/foundation/ABuffer.h>
 #include <media/stagefright/foundation/ADebug.h>
 #include <media/stagefright/foundation/AMessage.h>
+#include <media/stagefright/MediaBuffer.h>
 #include <media/stagefright/MediaCodec.h>
 #include <media/stagefright/MediaDefs.h>
 #include <media/stagefright/MediaErrors.h>
@@ -53,14 +56,12 @@ Converter::Converter(
 }
 
 Converter::~Converter() {
-    if (mEncoder != NULL) {
-        mEncoder->release();
-        mEncoder.clear();
-    }
+    CHECK(mEncoder == NULL);
+}
 
-    AString mime;
-    CHECK(mInputFormat->findString("mime", &mime));
-    ALOGI("encoder (%s) shut down.", mime.c_str());
+void Converter::shutdownAsync() {
+    ALOGV("shutdown");
+    (new AMessage(kWhatShutdown, id()))->post();
 }
 
 status_t Converter::initCheck() const {
@@ -155,16 +156,6 @@ status_t Converter::initEncoder() {
     return mEncoder->getOutputBuffers(&mEncoderOutputBuffers);
 }
 
-void Converter::feedAccessUnit(const sp<ABuffer> &accessUnit) {
-    sp<AMessage> msg = new AMessage(kWhatFeedAccessUnit, id());
-    msg->setBuffer("accessUnit", accessUnit);
-    msg->post();
-}
-
-void Converter::signalEOS() {
-    (new AMessage(kWhatInputEOS, id()))->post();
-}
-
 void Converter::notifyError(status_t err) {
     sp<AMessage> notify = mNotify->dup();
     notify->setInt32("what", kWhatError);
@@ -174,32 +165,70 @@ void Converter::notifyError(status_t err) {
 
 void Converter::onMessageReceived(const sp<AMessage> &msg) {
     switch (msg->what()) {
-        case kWhatFeedAccessUnit:
+        case kWhatMediaPullerNotify:
         {
-            sp<ABuffer> accessUnit;
-            CHECK(msg->findBuffer("accessUnit", &accessUnit));
+            int32_t what;
+            CHECK(msg->findInt32("what", &what));
 
-            mInputBufferQueue.push_back(accessUnit);
+            if (mEncoder == NULL) {
+                ALOGV("got msg '%s' after encoder shutdown.",
+                      msg->debugString().c_str());
 
-            feedEncoderInputBuffers();
+                if (what == MediaPuller::kWhatAccessUnit) {
+                    sp<ABuffer> accessUnit;
+                    CHECK(msg->findBuffer("accessUnit", &accessUnit));
 
-            scheduleDoMoreWork();
-            break;
-        }
+                    void *mbuf;
+                    if (accessUnit->meta()->findPointer("mediaBuffer", &mbuf)
+                            && mbuf != NULL) {
+                        ALOGV("releasing mbuf %p", mbuf);
 
-        case kWhatInputEOS:
-        {
-            mInputBufferQueue.push_back(NULL);
+                        accessUnit->meta()->setPointer("mediaBuffer", NULL);
+
+                        static_cast<MediaBuffer *>(mbuf)->release();
+                        mbuf = NULL;
+                    }
+                }
+                break;
+            }
+
+            if (what == MediaPuller::kWhatEOS) {
+                mInputBufferQueue.push_back(NULL);
+
+                feedEncoderInputBuffers();
+
+                scheduleDoMoreWork();
+            } else {
+                CHECK_EQ(what, MediaPuller::kWhatAccessUnit);
+
+                sp<ABuffer> accessUnit;
+                CHECK(msg->findBuffer("accessUnit", &accessUnit));
+
+#if 0
+                void *mbuf;
+                if (accessUnit->meta()->findPointer("mediaBuffer", &mbuf)
+                        && mbuf != NULL) {
+                    ALOGI("queueing mbuf %p", mbuf);
+                }
+#endif
+
+                mInputBufferQueue.push_back(accessUnit);
 
-            feedEncoderInputBuffers();
+                feedEncoderInputBuffers();
 
-            scheduleDoMoreWork();
+                scheduleDoMoreWork();
+            }
             break;
         }
 
         case kWhatDoMoreWork:
         {
             mDoMoreWorkPending = false;
+
+            if (mEncoder == NULL) {
+                break;
+            }
+
             status_t err = doMoreWork();
 
             if (err != OK) {
@@ -212,6 +241,10 @@ void Converter::onMessageReceived(const sp<AMessage> &msg) {
 
         case kWhatRequestIDRFrame:
         {
+            if (mEncoder == NULL) {
+                break;
+            }
+
             if (mIsVideo) {
                 ALOGI("requesting IDR frame");
                 mEncoder->requestIDRFrame();
@@ -219,6 +252,18 @@ void Converter::onMessageReceived(const sp<AMessage> &msg) {
             break;
         }
 
+        case kWhatShutdown:
+        {
+            ALOGI("shutting down encoder");
+            mEncoder->release();
+            mEncoder.clear();
+
+            AString mime;
+            CHECK(mInputFormat->findString("mime", &mime));
+            ALOGI("encoder (%s) shut down.", mime.c_str());
+            break;
+        }
+
         default:
             TRESPASS();
     }
diff --git a/media/libstagefright/wifi-display/source/Converter.h b/media/libstagefright/wifi-display/source/Converter.h
index 901ae2e..9f54523 100644
--- a/media/libstagefright/wifi-display/source/Converter.h
+++ b/media/libstagefright/wifi-display/source/Converter.h
@@ -51,18 +51,20 @@ struct Converter : public AHandler {
         kWhatError,
     };
 
-protected:
-    virtual ~Converter();
-    virtual void onMessageReceived(const sp<AMessage> &msg);
-
-private:
     enum {
-        kWhatFeedAccessUnit,
-        kWhatInputEOS,
         kWhatDoMoreWork,
         kWhatRequestIDRFrame,
+        kWhatShutdown,
+        kWhatMediaPullerNotify,
     };
 
+    void shutdownAsync();
+
+protected:
+    virtual ~Converter();
+    virtual void onMessageReceived(const sp<AMessage> &msg);
+
+private:
     status_t mInitCheck;
     sp<AMessage> mNotify;
     sp<ALooper> mCodecLooper;
diff --git a/media/libstagefright/wifi-display/source/MediaPuller.cpp b/media/libstagefright/wifi-display/source/MediaPuller.cpp
index 35ae539..82ae001 100644
--- a/media/libstagefright/wifi-display/source/MediaPuller.cpp
+++ b/media/libstagefright/wifi-display/source/MediaPuller.cpp
@@ -65,33 +65,20 @@ status_t MediaPuller::start() {
     return postSynchronouslyAndReturnError(new AMessage(kWhatStart, id()));
 }
 
-status_t MediaPuller::stop() {
-    return postSynchronouslyAndReturnError(new AMessage(kWhatStop, id()));
+void MediaPuller::stopAsync(const sp<AMessage> &notify) {
+    sp<AMessage> msg = new AMessage(kWhatStop, id());
+    msg->setMessage("notify", notify);
+    msg->post();
 }
 
 void MediaPuller::onMessageReceived(const sp<AMessage> &msg) {
     switch (msg->what()) {
         case kWhatStart:
-        case kWhatStop:
         {
-            status_t err;
+            status_t err = mSource->start();
 
-            if (msg->what() == kWhatStart) {
-                err = mSource->start();
-
-                if (err == OK) {
-                    schedulePull();
-                }
-            } else {
-                sp<MetaData> meta = mSource->getFormat();
-                const char *tmp;
-                CHECK(meta->findCString(kKeyMIMEType, &tmp));
-                AString mime = tmp;
-
-                ALOGI("MediaPuller(%s) stopping.", mime.c_str());
-                err = mSource->stop();
-                ALOGI("MediaPuller(%s) stopped.", mime.c_str());
-                ++mPullGeneration;
+            if (err == OK) {
+                schedulePull();
             }
 
             sp<AMessage> response = new AMessage;
@@ -104,6 +91,24 @@ void MediaPuller::onMessageReceived(const sp<AMessage> &msg) {
             break;
         }
 
+        case kWhatStop:
+        {
+            sp<MetaData> meta = mSource->getFormat();
+            const char *tmp;
+            CHECK(meta->findCString(kKeyMIMEType, &tmp));
+            AString mime = tmp;
+
+            ALOGI("MediaPuller(%s) stopping.", mime.c_str());
+            mSource->stop();
+            ALOGI("MediaPuller(%s) stopped.", mime.c_str());
+            ++mPullGeneration;
+
+            sp<AMessage> notify;
+            CHECK(msg->findMessage("notify", &notify));
+            notify->post();
+            break;
+        }
+
         case kWhatPull:
         {
             int32_t generation;
@@ -153,6 +158,10 @@ void MediaPuller::onMessageReceived(const sp<AMessage> &msg) {
                 notify->setBuffer("accessUnit", accessUnit);
                 notify->post();
 
+                if (mbuf != NULL) {
+                    ALOGV("posted mbuf %p", mbuf);
+                }
+
                 schedulePull();
             }
             break;
diff --git a/media/libstagefright/wifi-display/source/MediaPuller.h b/media/libstagefright/wifi-display/source/MediaPuller.h
index 134e1c0..728da7b 100644
--- a/media/libstagefright/wifi-display/source/MediaPuller.h
+++ b/media/libstagefright/wifi-display/source/MediaPuller.h
@@ -33,7 +33,7 @@ struct MediaPuller : public AHandler {
     MediaPuller(const sp<MediaSource> &source, const sp<AMessage> &notify);
 
     status_t start();
-    status_t stop();
+    void stopAsync(const sp<AMessage> &notify);
 
 protected:
     virtual void onMessageReceived(const sp<AMessage> &msg);
diff --git a/media/libstagefright/wifi-display/source/PlaybackSession.cpp b/media/libstagefright/wifi-display/source/PlaybackSession.cpp
index 0facafe..9b729fc 100644
--- a/media/libstagefright/wifi-display/source/PlaybackSession.cpp
+++ b/media/libstagefright/wifi-display/source/PlaybackSession.cpp
@@ -53,14 +53,17 @@ namespace android {
 static size_t kMaxRTPPacketSize = 1500;
 static size_t kMaxNumTSPacketsPerRTPPacket = (kMaxRTPPacketSize - 12) / 188;
 
-struct WifiDisplaySource::PlaybackSession::Track : public RefBase {
-    Track(const sp<ALooper> &pullLooper,
+struct WifiDisplaySource::PlaybackSession::Track : public AHandler {
+    enum {
+        kWhatStopped,
+    };
+
+    Track(const sp<AMessage> &notify,
+          const sp<ALooper> &pullLooper,
           const sp<ALooper> &codecLooper,
           const sp<MediaPuller> &mediaPuller,
           const sp<Converter> &converter);
 
-    Track(const sp<AMessage> &format);
-
     sp<AMessage> getFormat();
     bool isAudio() const;
 
@@ -70,20 +73,27 @@ struct WifiDisplaySource::PlaybackSession::Track : public RefBase {
     void setPacketizerTrackIndex(size_t index);
 
     status_t start();
-    status_t stop();
+    void stopAsync();
+
+    bool isStopped() const { return !mStarted; }
 
     void queueAccessUnit(const sp<ABuffer> &accessUnit);
     sp<ABuffer> dequeueAccessUnit();
 
 protected:
+    virtual void onMessageReceived(const sp<AMessage> &msg);
     virtual ~Track();
 
 private:
+    enum {
+        kWhatMediaPullerStopped,
+    };
+
+    sp<AMessage> mNotify;
     sp<ALooper> mPullLooper;
     sp<ALooper> mCodecLooper;
     sp<MediaPuller> mMediaPuller;
     sp<Converter> mConverter;
-    sp<AMessage> mFormat;
     bool mStarted;
     ssize_t mPacketizerTrackIndex;
     bool mIsAudio;
@@ -95,11 +105,13 @@ private:
 };
 
 WifiDisplaySource::PlaybackSession::Track::Track(
+        const sp<AMessage> &notify,
         const sp<ALooper> &pullLooper,
         const sp<ALooper> &codecLooper,
         const sp<MediaPuller> &mediaPuller,
         const sp<Converter> &converter)
-    : mPullLooper(pullLooper),
+    : mNotify(notify),
+      mPullLooper(pullLooper),
       mCodecLooper(codecLooper),
       mMediaPuller(mediaPuller),
       mConverter(converter),
@@ -108,14 +120,8 @@ WifiDisplaySource::PlaybackSession::Track::Track(
       mIsAudio(IsAudioFormat(mConverter->getOutputFormat())) {
 }
 
-WifiDisplaySource::PlaybackSession::Track::Track(const sp<AMessage> &format)
-    : mFormat(format),
-      mPacketizerTrackIndex(-1),
-      mIsAudio(IsAudioFormat(mFormat)) {
-}
-
 WifiDisplaySource::PlaybackSession::Track::~Track() {
-    stop();
+    CHECK(!mStarted);
 }
 
 // static
@@ -128,10 +134,6 @@ bool WifiDisplaySource::PlaybackSession::Track::IsAudioFormat(
 }
 
 sp<AMessage> WifiDisplaySource::PlaybackSession::Track::getFormat() {
-    if (mFormat != NULL) {
-        return mFormat;
-    }
-
     return mConverter->getOutputFormat();
 }
 
@@ -155,9 +157,7 @@ void WifiDisplaySource::PlaybackSession::Track::setPacketizerTrackIndex(size_t i
 status_t WifiDisplaySource::PlaybackSession::Track::start() {
     ALOGV("Track::start isAudio=%d", mIsAudio);
 
-    if (mStarted) {
-        return INVALID_OPERATION;
-    }
+    CHECK(!mStarted);
 
     status_t err = OK;
 
@@ -172,24 +172,40 @@ status_t WifiDisplaySource::PlaybackSession::Track::start() {
     return err;
 }
 
-status_t WifiDisplaySource::PlaybackSession::Track::stop() {
-    ALOGV("Track::stop isAudio=%d", mIsAudio);
+void WifiDisplaySource::PlaybackSession::Track::stopAsync() {
+    ALOGV("Track::stopAsync isAudio=%d", mIsAudio);
 
-    if (!mStarted) {
-        return INVALID_OPERATION;
-    }
+    CHECK(mStarted);
 
-    status_t err = OK;
+    mConverter->shutdownAsync();
+
+    sp<AMessage> msg = new AMessage(kWhatMediaPullerStopped, id());
 
     if (mMediaPuller != NULL) {
-        err = mMediaPuller->stop();
+        mMediaPuller->stopAsync(msg);
+    } else {
+        msg->post();
     }
+}
 
-    mConverter.clear();
+void WifiDisplaySource::PlaybackSession::Track::onMessageReceived(
+        const sp<AMessage> &msg) {
+    switch (msg->what()) {
+        case kWhatMediaPullerStopped:
+        {
+            mConverter.clear();
 
-    mStarted = false;
+            mStarted = false;
 
-    return err;
+            sp<AMessage> notify = mNotify->dup();
+            notify->setInt32("what", kWhatStopped);
+            notify->post();
+            break;
+        }
+
+        default:
+            TRESPASS();
+    }
 }
 
 void WifiDisplaySource::PlaybackSession::Track::queueAccessUnit(
@@ -482,15 +498,7 @@ status_t WifiDisplaySource::PlaybackSession::onFinishPlay2() {
     }
 
     for (size_t i = 0; i < mTracks.size(); ++i) {
-        status_t err = mTracks.editValueAt(i)->start();
-
-        if (err != OK) {
-            for (size_t j = 0; j < i; ++j) {
-                mTracks.editValueAt(j)->stop();
-            }
-
-            return err;
-        }
+        CHECK_EQ((status_t)OK, mTracks.editValueAt(i)->start());
     }
 
     sp<AMessage> notify = mNotify->dup();
@@ -506,32 +514,12 @@ status_t WifiDisplaySource::PlaybackSession::pause() {
     return OK;
 }
 
-status_t WifiDisplaySource::PlaybackSession::destroy() {
-    mTracks.clear();
-
-    mPacketizer.clear();
-
-    mTracks.clear();
-
-#if ENABLE_RETRANSMISSION
-    if (mRTCPRetransmissionSessionID != 0) {
-        mNetSession->destroySession(mRTCPRetransmissionSessionID);
-    }
-
-    if (mRTPRetransmissionSessionID != 0) {
-        mNetSession->destroySession(mRTPRetransmissionSessionID);
-    }
-#endif
+void WifiDisplaySource::PlaybackSession::destroyAsync() {
+    ALOGI("destroyAsync");
 
-    if (mRTCPSessionID != 0) {
-        mNetSession->destroySession(mRTCPSessionID);
-    }
-
-    if (mRTPSessionID != 0) {
-        mNetSession->destroySession(mRTPSessionID);
+    for (size_t i = 0; i < mTracks.size(); ++i) {
+        mTracks.valueAt(i)->stopAsync();
     }
-
-    return OK;
 }
 
 void WifiDisplaySource::PlaybackSession::onMessageReceived(
@@ -669,32 +657,6 @@ void WifiDisplaySource::PlaybackSession::onMessageReceived(
             break;
         }
 
-        case kWhatMediaPullerNotify:
-        {
-            int32_t what;
-            CHECK(msg->findInt32("what", &what));
-
-            if (what == MediaPuller::kWhatEOS) {
-                ALOGI("input eos");
-
-                for (size_t i = 0; i < mTracks.size(); ++i) {
-                    mTracks.valueAt(i)->converter()->signalEOS();
-                }
-            } else {
-                CHECK_EQ(what, MediaPuller::kWhatAccessUnit);
-
-                size_t trackIndex;
-                CHECK(msg->findSize("trackIndex", &trackIndex));
-
-                sp<ABuffer> accessUnit;
-                CHECK(msg->findBuffer("accessUnit", &accessUnit));
-
-                mTracks.valueFor(trackIndex)->converter()
-                    ->feedAccessUnit(accessUnit);
-            }
-            break;
-        }
-
         case kWhatConverterNotify:
         {
             int32_t what;
@@ -776,6 +738,57 @@ void WifiDisplaySource::PlaybackSession::onMessageReceived(
             break;
         }
 
+        case kWhatTrackNotify:
+        {
+            int32_t what;
+            CHECK(msg->findInt32("what", &what));
+
+            size_t trackIndex;
+            CHECK(msg->findSize("trackIndex", &trackIndex));
+
+            if (what == Track::kWhatStopped) {
+                bool allTracksAreStopped = true;
+                for (size_t i = 0; i < mTracks.size(); ++i) {
+                    const sp<Track> &track = mTracks.valueAt(i);
+                    if (!track->isStopped()) {
+                        allTracksAreStopped = false;
+                        break;
+                    }
+                }
+
+                if (!allTracksAreStopped) {
+                    break;
+                }
+
+                mTracks.clear();
+
+                mPacketizer.clear();
+
+#if ENABLE_RETRANSMISSION
+                if (mRTCPRetransmissionSessionID != 0) {
+                    mNetSession->destroySession(mRTCPRetransmissionSessionID);
+                }
+
+                if (mRTPRetransmissionSessionID != 0) {
+                    mNetSession->destroySession(mRTPRetransmissionSessionID);
+                }
+#endif
+
+                if (mRTCPSessionID != 0) {
+                    mNetSession->destroySession(mRTCPSessionID);
+                }
+
+                if (mRTPSessionID != 0) {
+                    mNetSession->destroySession(mRTPSessionID);
+                }
+
+                sp<AMessage> notify = mNotify->dup();
+                notify->setInt32("what", kWhatSessionDestroyed);
+                notify->post();
+            }
+            break;
+        }
+
         default:
             TRESPASS();
     }
@@ -817,11 +830,6 @@ status_t WifiDisplaySource::PlaybackSession::addSource(
 
     trackIndex = mTracks.size();
 
-    notify = new AMessage(kWhatMediaPullerNotify, id());
-    notify->setSize("trackIndex", trackIndex);
-    sp<MediaPuller> puller = new MediaPuller(source, notify);
-    pullLooper->registerHandler(puller);
-
     sp<AMessage> format;
     status_t err = convertMetaDataToMessage(source->getFormat(), &format);
     CHECK_EQ(err, (status_t)OK);
@@ -842,11 +850,25 @@ status_t WifiDisplaySource::PlaybackSession::addSource(
 
     looper()->registerHandler(converter);
 
+    notify = new AMessage(Converter::kWhatMediaPullerNotify, converter->id());
+    notify->setSize("trackIndex", trackIndex);
+
+    sp<MediaPuller> puller = new MediaPuller(source, notify);
+    pullLooper->registerHandler(puller);
+
     if (numInputBuffers != NULL) {
         *numInputBuffers = converter->getInputBufferCount();
     }
 
-    mTracks.add(trackIndex, new Track(pullLooper, codecLooper, puller, converter));
+    notify = new AMessage(kWhatTrackNotify, id());
+    notify->setSize("trackIndex", trackIndex);
+
+    sp<Track> track = new Track(
+            notify, pullLooper, codecLooper, puller, converter);
+
+    looper()->registerHandler(track);
+
+    mTracks.add(trackIndex, track);
 
     if (isVideo) {
         mVideoTrackIndex = trackIndex;
diff --git a/media/libstagefright/wifi-display/source/PlaybackSession.h b/media/libstagefright/wifi-display/source/PlaybackSession.h
index 342fc85..9237a72 100644
--- a/media/libstagefright/wifi-display/source/PlaybackSession.h
+++ b/media/libstagefright/wifi-display/source/PlaybackSession.h
@@ -51,7 +51,7 @@ struct WifiDisplaySource::PlaybackSession : public AHandler {
             const char *clientIP, int32_t clientRtp, int32_t clientRtcp,
             TransportMode transportMode);
 
-    status_t destroy();
+    void destroyAsync();
 
     int32_t getRTPPort() const;
 
@@ -72,6 +72,7 @@ struct WifiDisplaySource::PlaybackSession : public AHandler {
         kWhatSessionDead,
         kWhatBinaryData,
         kWhatSessionEstablished,
+        kWhatSessionDestroyed,
     };
 
 protected:
@@ -91,6 +92,7 @@ private:
 #endif
         kWhatMediaPullerNotify,
         kWhatConverterNotify,
+        kWhatTrackNotify,
         kWhatUpdateSurface,
         kWhatFinishPlay,
     };
diff --git a/media/libstagefright/wifi-display/source/RepeaterSource.cpp b/media/libstagefright/wifi-display/source/RepeaterSource.cpp
index 56e8860..483d29c 100644
--- a/media/libstagefright/wifi-display/source/RepeaterSource.cpp
+++ b/media/libstagefright/wifi-display/source/RepeaterSource.cpp
@@ -50,6 +50,8 @@ status_t RepeaterSource::start(MetaData *params) {
 }
 
 status_t RepeaterSource::stop() {
+    ALOGV("stopping");
+
     if (mLooper != NULL) {
         mLooper->stop();
         mLooper.clear();
@@ -57,7 +59,17 @@ status_t RepeaterSource::stop() {
         mReflector.clear();
     }
 
-    return mSource->stop();
+    if (mBuffer != NULL) {
+        ALOGV("releasing mbuf %p", mBuffer);
+        mBuffer->release();
+        mBuffer = NULL;
+    }
+
+    status_t err = mSource->stop();
+
+    ALOGV("stopped");
+
+    return err;
 }
 
 sp<MetaData> RepeaterSource::getFormat() {
@@ -117,6 +129,8 @@ void RepeaterSource::onMessageReceived(const sp<AMessage> &msg) {
             MediaBuffer *buffer;
             status_t err = mSource->read(&buffer);
 
+            ALOGV("read mbuf %p", buffer);
+
             Mutex::Autolock autoLock(mLock);
             if (mBuffer != NULL) {
                 mBuffer->release();
diff --git a/media/libstagefright/wifi-display/source/WifiDisplaySource.cpp b/media/libstagefright/wifi-display/source/WifiDisplaySource.cpp
index 16c0f35..8091cc4 100644
--- a/media/libstagefright/wifi-display/source/WifiDisplaySource.cpp
+++ b/media/libstagefright/wifi-display/source/WifiDisplaySource.cpp
@@ -293,6 +293,8 @@ void WifiDisplaySource::onMessageReceived(const sp<AMessage> &msg) {
                             mClientInfo.mPlaybackSession->height(),
                             0 /* flags */);
                 }
+            } else if (what == PlaybackSession::kWhatSessionDestroyed) {
+                disconnectClient2();
             } else {
                 CHECK_EQ(what, PlaybackSession::kWhatBinaryData);
 
@@ -1117,6 +1119,12 @@ status_t WifiDisplaySource::onTeardownRequest(
 void WifiDisplaySource::finishStop() {
     ALOGV("finishStop");
 
+    disconnectClientAsync();
+}
+
+void WifiDisplaySource::finishStopAfterDisconnectingClient() {
+    ALOGV("finishStopAfterDisconnectingClient");
+
 #if REQUIRE_HDCP
     if (mHDCP != NULL) {
         ALOGI("Initiating HDCP shutdown.");
@@ -1137,14 +1145,12 @@ void WifiDisplaySource::finishStop2() {
     mHDCP.clear();
 #endif
 
-    disconnectClient();
-
     if (mSessionID != 0) {
         mNetSession->destroySession(mSessionID);
         mSessionID = 0;
     }
 
-    ALOGV("finishStop2 completed.");
+    ALOGI("We're stopped.");
 
     status_t err = OK;
 
@@ -1264,14 +1270,26 @@ sp<WifiDisplaySource::PlaybackSession> WifiDisplaySource::findPlaybackSession(
     return mClientInfo.mPlaybackSession;
 }
 
-void WifiDisplaySource::disconnectClient() {
+void WifiDisplaySource::disconnectClientAsync() {
+    ALOGV("disconnectClient");
+
+    if (mClientInfo.mPlaybackSession == NULL) {
+        disconnectClient2();
+        return;
+    }
+
+    if (mClientInfo.mPlaybackSession != NULL) {
+        ALOGV("Destroying PlaybackSession");
+        mClientInfo.mPlaybackSession->destroyAsync();
+    }
+}
+
+void WifiDisplaySource::disconnectClient2() {
+    ALOGV("disconnectClient2");
+
     if (mClientInfo.mPlaybackSession != NULL) {
-        sp<PlaybackSession> playbackSession = mClientInfo.mPlaybackSession;
+        looper()->unregisterHandler(mClientInfo.mPlaybackSession->id());
         mClientInfo.mPlaybackSession.clear();
-
-        ALOGI("Destroying PlaybackSession");
-        playbackSession->destroy();
-        looper()->unregisterHandler(playbackSession->id());
     }
 
     if (mClientSessionID != 0) {
@@ -1280,6 +1298,8 @@ void WifiDisplaySource::disconnectClient() {
     }
 
     mClient->onDisplayDisconnected();
+
+    finishStopAfterDisconnectingClient();
 }
 
 #if REQUIRE_HDCP
diff --git a/media/libstagefright/wifi-display/source/WifiDisplaySource.h b/media/libstagefright/wifi-display/source/WifiDisplaySource.h
index 77b15f8..ade623a 100644
--- a/media/libstagefright/wifi-display/source/WifiDisplaySource.h
+++ b/media/libstagefright/wifi-display/source/WifiDisplaySource.h
@@ -200,11 +200,10 @@ private:
     sp<PlaybackSession> findPlaybackSession(
             const sp<ParsedMessage> &data, int32_t *playbackSessionID) const;
 
-    // Disconnects the current client and shuts down its playback session
-    // (if any).
-    void disconnectClient();
-
     void finishStop();
+    void disconnectClientAsync();
+    void disconnectClient2();
+    void finishStopAfterDisconnectingClient();
     void finishStop2();
 
     DISALLOW_EVIL_CONSTRUCTORS(WifiDisplaySource);