Minor tweaks to DirectRenderer and Converter

Converter now supports automatic prepending of SPS/PPS to IDR frames (h264)
as well as using the encoder in "surface-input" mode.
The new features are all opt-in and should not affect existing clients.

Change-Id: I543cf1d31ba068c1a01ab4e6814ac8d817b63faa

5 files changed, 209 insertions, 60 deletions

diff --git a/media/libstagefright/wifi-display/sink/DirectRenderer.cpp b/media/libstagefright/wifi-display/sink/DirectRenderer.cpp
index 15f9c88..cdb2267 100644
--- a/media/libstagefright/wifi-display/sink/DirectRenderer.cpp
+++ b/media/libstagefright/wifi-display/sink/DirectRenderer.cpp
@@ -29,9 +29,8 @@
 #include <media/stagefright/foundation/AMessage.h>
 #include <media/stagefright/foundation/hexdump.h>
 #include <media/stagefright/MediaCodec.h>
+#include <media/stagefright/MediaDefs.h>
 #include <media/stagefright/MediaErrors.h>
-#include <media/stagefright/MetaData.h>
-#include <media/stagefright/Utils.h>
 
 namespace android {
 
@@ -488,12 +487,38 @@ void DirectRenderer::onMessageReceived(const sp<AMessage> &msg) {
             break;
         }
 
+        case kWhatQueueAccessUnit:
+            onQueueAccessUnit(msg);
+            break;
+
+        case kWhatSetFormat:
+            onSetFormat(msg);
+            break;
+
         default:
             TRESPASS();
     }
 }
 
 void DirectRenderer::setFormat(size_t trackIndex, const sp<AMessage> &format) {
+    sp<AMessage> msg = new AMessage(kWhatSetFormat, id());
+    msg->setSize("trackIndex", trackIndex);
+    msg->setMessage("format", format);
+    msg->post();
+}
+
+void DirectRenderer::onSetFormat(const sp<AMessage> &msg) {
+    size_t trackIndex;
+    CHECK(msg->findSize("trackIndex", &trackIndex));
+
+    sp<AMessage> format;
+    CHECK(msg->findMessage("format", &format));
+
+    internalSetFormat(trackIndex, format);
+}
+
+void DirectRenderer::internalSetFormat(
+        size_t trackIndex, const sp<AMessage> &format) {
     CHECK_LT(trackIndex, 2u);
 
     CHECK(mDecoderContext[trackIndex] == NULL);
@@ -517,18 +542,21 @@ void DirectRenderer::setFormat(size_t trackIndex, const sp<AMessage> &format) {
 
 void DirectRenderer::queueAccessUnit(
         size_t trackIndex, const sp<ABuffer> &accessUnit) {
-    CHECK_LT(trackIndex, 2u);
+    sp<AMessage> msg = new AMessage(kWhatQueueAccessUnit, id());
+    msg->setSize("trackIndex", trackIndex);
+    msg->setBuffer("accessUnit", accessUnit);
+    msg->post();
+}
 
-    if (mDecoderContext[trackIndex] == NULL) {
-        CHECK_EQ(trackIndex, 0u);
+void DirectRenderer::onQueueAccessUnit(const sp<AMessage> &msg) {
+    size_t trackIndex;
+    CHECK(msg->findSize("trackIndex", &trackIndex));
 
-        sp<AMessage> format = new AMessage;
-        format->setString("mime", "video/avc");
-        format->setInt32("width", 640);
-        format->setInt32("height", 360);
+    sp<ABuffer> accessUnit;
+    CHECK(msg->findBuffer("accessUnit", &accessUnit));
 
-        setFormat(trackIndex, format);
-    }
+    CHECK_LT(trackIndex, 2u);
+    CHECK(mDecoderContext[trackIndex] != NULL);
 
     mDecoderContext[trackIndex]->queueInputBuffer(accessUnit);
 }
diff --git a/media/libstagefright/wifi-display/sink/DirectRenderer.h b/media/libstagefright/wifi-display/sink/DirectRenderer.h
index 1e7dc34..07c2170 100644
--- a/media/libstagefright/wifi-display/sink/DirectRenderer.h
+++ b/media/libstagefright/wifi-display/sink/DirectRenderer.h
@@ -43,6 +43,8 @@ private:
     enum {
         kWhatDecoderNotify,
         kWhatRenderVideo,
+        kWhatQueueAccessUnit,
+        kWhatSetFormat,
     };
 
     struct OutputInfo {
@@ -72,6 +74,11 @@ private:
     void scheduleVideoRenderIfNecessary();
     void onRenderVideo();
 
+    void onSetFormat(const sp<AMessage> &msg);
+    void onQueueAccessUnit(const sp<AMessage> &msg);
+
+    void internalSetFormat(size_t trackIndex, const sp<AMessage> &format);
+
     DISALLOW_EVIL_CONSTRUCTORS(DirectRenderer);
 };
 
diff --git a/media/libstagefright/wifi-display/source/Converter.cpp b/media/libstagefright/wifi-display/source/Converter.cpp
index 0214520..6f23854 100644
--- a/media/libstagefright/wifi-display/source/Converter.cpp
+++ b/media/libstagefright/wifi-display/source/Converter.cpp
@@ -21,6 +21,7 @@
 #include "Converter.h"
 
 #include "MediaPuller.h"
+#include "include/avc_utils.h"
 
 #include <cutils/properties.h>
 #include <gui/Surface.h>
@@ -33,6 +34,8 @@
 #include <media/stagefright/MediaDefs.h>
 #include <media/stagefright/MediaErrors.h>
 
+#include <arpa/inet.h>
+
 #include <OMX_Video.h>
 
 namespace android {
@@ -40,12 +43,14 @@ namespace android {
 Converter::Converter(
         const sp<AMessage> &notify,
         const sp<ALooper> &codecLooper,
-        const sp<AMessage> &outputFormat)
-    : mInitCheck(NO_INIT),
-      mNotify(notify),
+        const sp<AMessage> &outputFormat,
+        uint32_t flags)
+    : mNotify(notify),
       mCodecLooper(codecLooper),
       mOutputFormat(outputFormat),
+      mFlags(flags),
       mIsVideo(false),
+      mIsH264(false),
       mIsPCMAudio(false),
       mNeedToManuallyPrependSPSPPS(false),
       mDoMoreWorkPending(false)
@@ -55,21 +60,18 @@ Converter::Converter(
 #endif
       ,mPrevVideoBitrate(-1)
       ,mNumFramesToDrop(0)
+      ,mEncodingSuspended(false)
     {
     AString mime;
     CHECK(mOutputFormat->findString("mime", &mime));
 
     if (!strncasecmp("video/", mime.c_str(), 6)) {
         mIsVideo = true;
+
+        mIsH264 = !strcasecmp(mime.c_str(), MEDIA_MIMETYPE_VIDEO_AVC);
     } else if (!strcasecmp(MEDIA_MIMETYPE_AUDIO_RAW, mime.c_str())) {
         mIsPCMAudio = true;
     }
-
-    mInitCheck = initEncoder();
-
-    if (mInitCheck != OK) {
-        releaseEncoder();
-    }
 }
 
 static void ReleaseMediaBufferReference(const sp<ABuffer> &accessUnit) {
@@ -118,8 +120,19 @@ void Converter::shutdownAsync() {
     (new AMessage(kWhatShutdown, id()))->post();
 }
 
-status_t Converter::initCheck() const {
-    return mInitCheck;
+status_t Converter::init() {
+    status_t err = initEncoder();
+
+    if (err != OK) {
+        releaseEncoder();
+    }
+
+    return err;
+}
+
+sp<IGraphicBufferProducer> Converter::getGraphicBufferProducer() {
+    CHECK(mFlags & FLAG_USE_SURFACE_INPUT);
+    return mGraphicBufferProducer;
 }
 
 size_t Converter::getInputBufferCount() const {
@@ -244,6 +257,16 @@ status_t Converter::initEncoder() {
         return err;
     }
 
+    if (mFlags & FLAG_USE_SURFACE_INPUT) {
+        CHECK(mIsVideo);
+
+        err = mEncoder->createInputSurface(&mGraphicBufferProducer);
+
+        if (err != OK) {
+            return err;
+        }
+    }
+
     err = mEncoder->start();
 
     if (err != OK) {
@@ -256,7 +279,17 @@ status_t Converter::initEncoder() {
         return err;
     }
 
-    return mEncoder->getOutputBuffers(&mEncoderOutputBuffers);
+    err = mEncoder->getOutputBuffers(&mEncoderOutputBuffers);
+
+    if (err != OK) {
+        return err;
+    }
+
+    if (mFlags & FLAG_USE_SURFACE_INPUT) {
+        scheduleDoMoreWork();
+    }
+
+    return OK;
 }
 
 void Converter::notifyError(status_t err) {
@@ -312,9 +345,12 @@ void Converter::onMessageReceived(const sp<AMessage> &msg) {
                 sp<ABuffer> accessUnit;
                 CHECK(msg->findBuffer("accessUnit", &accessUnit));
 
-                if (mIsVideo && mNumFramesToDrop) {
-                    --mNumFramesToDrop;
-                    ALOGI("dropping frame.");
+                if (mNumFramesToDrop > 0 || mEncodingSuspended) {
+                    if (mNumFramesToDrop > 0) {
+                        --mNumFramesToDrop;
+                        ALOGI("dropping frame.");
+                    }
+
                     ReleaseMediaBufferReference(accessUnit);
                     break;
                 }
@@ -396,7 +432,7 @@ void Converter::onMessageReceived(const sp<AMessage> &msg) {
             }
 
             if (mIsVideo) {
-                ALOGI("requesting IDR frame");
+                ALOGV("requesting IDR frame");
                 mEncoder->requestIDRFrame();
             }
             break;
@@ -411,6 +447,10 @@ void Converter::onMessageReceived(const sp<AMessage> &msg) {
             AString mime;
             CHECK(mOutputFormat->findString("mime", &mime));
             ALOGI("encoder (%s) shut down.", mime.c_str());
+
+            sp<AMessage> notify = mNotify->dup();
+            notify->setInt32("what", kWhatShutdownCompleted);
+            notify->post();
             break;
         }
 
@@ -431,6 +471,21 @@ void Converter::onMessageReceived(const sp<AMessage> &msg) {
             break;
         }
 
+        case kWhatSuspendEncoding:
+        {
+            int32_t suspend;
+            CHECK(msg->findInt32("suspend", &suspend));
+
+            mEncodingSuspended = suspend;
+
+            if (mFlags & FLAG_USE_SURFACE_INPUT) {
+                sp<AMessage> params = new AMessage;
+                params->setInt32("drop-input-frames",suspend);
+                mEncoder->setParameters(params);
+            }
+            break;
+        }
+
         default:
             TRESPASS();
     }
@@ -616,22 +671,39 @@ status_t Converter::feedEncoderInputBuffers() {
     return OK;
 }
 
+sp<ABuffer> Converter::prependCSD(const sp<ABuffer> &accessUnit) const {
+    CHECK(mCSD0 != NULL);
+
+    sp<ABuffer> dup = new ABuffer(accessUnit->size() + mCSD0->size());
+    memcpy(dup->data(), mCSD0->data(), mCSD0->size());
+    memcpy(dup->data() + mCSD0->size(), accessUnit->data(), accessUnit->size());
+
+    int64_t timeUs;
+    CHECK(accessUnit->meta()->findInt64("timeUs", &timeUs));
+
+    dup->meta()->setInt64("timeUs", timeUs);
+
+    return dup;
+}
+
 status_t Converter::doMoreWork() {
     status_t err;
 
-    for (;;) {
-        size_t bufferIndex;
-        err = mEncoder->dequeueInputBuffer(&bufferIndex);
+    if (!(mFlags & FLAG_USE_SURFACE_INPUT)) {
+        for (;;) {
+            size_t bufferIndex;
+            err = mEncoder->dequeueInputBuffer(&bufferIndex);
 
-        if (err != OK) {
-            break;
+            if (err != OK) {
+                break;
+            }
+
+            mAvailEncoderInputIndices.push_back(bufferIndex);
         }
 
-        mAvailEncoderInputIndices.push_back(bufferIndex);
+        feedEncoderInputBuffers();
     }
 
-    feedEncoderInputBuffers();
-
     for (;;) {
         size_t bufferIndex;
         size_t offset;
@@ -705,9 +777,19 @@ status_t Converter::doMoreWork() {
 
             if (flags & MediaCodec::BUFFER_FLAG_CODECCONFIG) {
                 if (!handle) {
+                    if (mIsH264) {
+                        mCSD0 = buffer;
+                    }
                     mOutputFormat->setBuffer("csd-0", buffer);
                 }
             } else {
+                if (mNeedToManuallyPrependSPSPPS
+                        && mIsH264
+                        && (mFlags & FLAG_PREPEND_CSD_IF_NECESSARY)
+                        && IsIDR(buffer)) {
+                    buffer = prependCSD(buffer);
+                }
+
                 sp<AMessage> notify = mNotify->dup();
                 notify->setInt32("what", kWhatAccessUnit);
                 notify->setBuffer("accessUnit", buffer);
@@ -732,9 +814,18 @@ void Converter::requestIDRFrame() {
 }
 
 void Converter::dropAFrame() {
+    // Unsupported in surface input mode.
+    CHECK(!(mFlags & FLAG_USE_SURFACE_INPUT));
+
     (new AMessage(kWhatDropAFrame, id()))->post();
 }
 
+void Converter::suspendEncoding(bool suspend) {
+    sp<AMessage> msg = new AMessage(kWhatSuspendEncoding, id());
+    msg->setInt32("suspend", suspend);
+    msg->post();
+}
+
 int32_t Converter::getVideoBitrate() const {
     return mPrevVideoBitrate;
 }
diff --git a/media/libstagefright/wifi-display/source/Converter.h b/media/libstagefright/wifi-display/source/Converter.h
index 76c8b19..5876e07 100644
--- a/media/libstagefright/wifi-display/source/Converter.h
+++ b/media/libstagefright/wifi-display/source/Converter.h
@@ -18,13 +18,12 @@
 
 #define CONVERTER_H_
 
-#include "WifiDisplaySource.h"
-
 #include <media/stagefright/foundation/AHandler.h>
 
 namespace android {
 
 struct ABuffer;
+struct IGraphicBufferProducer;
 struct MediaCodec;
 
 #define ENABLE_SILENCE_DETECTION        0
@@ -33,11 +32,25 @@ struct MediaCodec;
 // media access unit of a different format.
 // Right now this'll convert raw video into H.264 and raw audio into AAC.
 struct Converter : public AHandler {
+    enum {
+        kWhatAccessUnit,
+        kWhatEOS,
+        kWhatError,
+        kWhatShutdownCompleted,
+    };
+
+    enum FlagBits {
+        FLAG_USE_SURFACE_INPUT          = 1,
+        FLAG_PREPEND_CSD_IF_NECESSARY   = 2,
+    };
     Converter(const sp<AMessage> &notify,
               const sp<ALooper> &codecLooper,
-              const sp<AMessage> &outputFormat);
+              const sp<AMessage> &outputFormat,
+              uint32_t flags = 0);
 
-    status_t initCheck() const;
+    status_t init();
+
+    sp<IGraphicBufferProducer> getGraphicBufferProducer();
 
     size_t getInputBufferCount() const;
 
@@ -50,22 +63,7 @@ struct Converter : public AHandler {
     void requestIDRFrame();
 
     void dropAFrame();
-
-    enum {
-        kWhatAccessUnit,
-        kWhatEOS,
-        kWhatError,
-    };
-
-    enum {
-        kWhatDoMoreWork,
-        kWhatRequestIDRFrame,
-        kWhatShutdown,
-        kWhatMediaPullerNotify,
-        kWhatEncoderActivity,
-        kWhatDropAFrame,
-        kWhatReleaseOutputBuffer,
-    };
+    void suspendEncoding(bool suspend);
 
     void shutdownAsync();
 
@@ -74,22 +72,40 @@ struct Converter : public AHandler {
 
     static int32_t GetInt32Property(const char *propName, int32_t defaultValue);
 
+    enum {
+        // MUST not conflict with private enums below.
+        kWhatMediaPullerNotify = 'pulN',
+    };
+
 protected:
     virtual ~Converter();
     virtual void onMessageReceived(const sp<AMessage> &msg);
 
 private:
-    status_t mInitCheck;
+    enum {
+        kWhatDoMoreWork,
+        kWhatRequestIDRFrame,
+        kWhatSuspendEncoding,
+        kWhatShutdown,
+        kWhatEncoderActivity,
+        kWhatDropAFrame,
+        kWhatReleaseOutputBuffer,
+    };
+
     sp<AMessage> mNotify;
     sp<ALooper> mCodecLooper;
     sp<AMessage> mOutputFormat;
+    uint32_t mFlags;
     bool mIsVideo;
+    bool mIsH264;
     bool mIsPCMAudio;
     bool mNeedToManuallyPrependSPSPPS;
 
     sp<MediaCodec> mEncoder;
     sp<AMessage> mEncoderActivityNotify;
 
+    sp<IGraphicBufferProducer> mGraphicBufferProducer;
+
     Vector<sp<ABuffer> > mEncoderInputBuffers;
     Vector<sp<ABuffer> > mEncoderOutputBuffers;
 
@@ -97,6 +113,8 @@ private:
 
     List<sp<ABuffer> > mInputBufferQueue;
 
+    sp<ABuffer> mCSD0;
+
     bool mDoMoreWorkPending;
 
 #if ENABLE_SILENCE_DETECTION
@@ -109,6 +127,7 @@ private:
     int32_t mPrevVideoBitrate;
 
     int32_t mNumFramesToDrop;
+    bool mEncodingSuspended;
 
     status_t initEncoder();
     void releaseEncoder();
@@ -127,6 +146,8 @@ private:
 
     static bool IsSilence(const sp<ABuffer> &accessUnit);
 
+    sp<ABuffer> prependCSD(const sp<ABuffer> &accessUnit) const;
+
     DISALLOW_EVIL_CONSTRUCTORS(Converter);
 };
 
diff --git a/media/libstagefright/wifi-display/source/PlaybackSession.cpp b/media/libstagefright/wifi-display/source/PlaybackSession.cpp
index a15fbac..0aa4ee5 100644
--- a/media/libstagefright/wifi-display/source/PlaybackSession.cpp
+++ b/media/libstagefright/wifi-display/source/PlaybackSession.cpp
@@ -521,7 +521,7 @@ void WifiDisplaySource::PlaybackSession::onMessageReceived(
                 if (mTracks.isEmpty()) {
                     ALOGI("Reached EOS");
                 }
-            } else {
+            } else if (what != Converter::kWhatShutdownCompleted) {
                 CHECK_EQ(what, Converter::kWhatError);
 
                 status_t err;
@@ -957,14 +957,16 @@ status_t WifiDisplaySource::PlaybackSession::addSource(
 
     sp<Converter> converter = new Converter(notify, codecLooper, format);
 
-    err = converter->initCheck();
+    looper()->registerHandler(converter);
+
+    err = converter->init();
     if (err != OK) {
         ALOGE("%s converter returned err %d", isVideo ? "video" : "audio", err);
+
+        looper()->unregisterHandler(converter->id());
         return err;
     }
 
-    looper()->registerHandler(converter);
-
     notify = new AMessage(Converter::kWhatMediaPullerNotify, converter->id());
     notify->setSize("trackIndex", trackIndex);