Send audio-presentation metadata in access units

Audio presentation metadata is attached to the access units so that
each media buffer read() retreives it's associated audio
presentations information.

Bug: 120846068
Test: atest CtsMediaTestCases:MediaExtractorTest#testGetAudioPresentations
      repeat 3..5 times
Change-Id: I47f4d0dd959537fe98aa757bfa762ff02aa86a35
diff --git a/media/extractors/mpeg2/MPEG2PSExtractor.cpp b/media/extractors/mpeg2/MPEG2PSExtractor.cpp
index 554d252..b99fde9 100644
--- a/media/extractors/mpeg2/MPEG2PSExtractor.cpp
+++ b/media/extractors/mpeg2/MPEG2PSExtractor.cpp
@@ -734,6 +734,9 @@
     if (inMeta.findData(kKeySEI, &bufType, &bufData, &bufSize)) {
         AMediaFormat_setBuffer(outMeta, AMEDIAFORMAT_KEY_SEI, bufData, bufSize);
     }
+    if (inMeta.findData(kKeyAudioPresentationInfo, &bufType, &bufData, &bufSize)) {
+        AMediaFormat_setBuffer(outMeta, AMEDIAFORMAT_KEY_AUDIO_PRESENTATION_INFO, bufData, bufSize);
+    }
     mbuf->release();
     return AMEDIA_OK;
 }
diff --git a/media/extractors/mpeg2/MPEG2TSExtractor.cpp b/media/extractors/mpeg2/MPEG2TSExtractor.cpp
index 3bb2af7..4ed6098 100644
--- a/media/extractors/mpeg2/MPEG2TSExtractor.cpp
+++ b/media/extractors/mpeg2/MPEG2TSExtractor.cpp
@@ -224,6 +224,9 @@
     if (inMeta.findData(kKeySEI, &bufType, &bufData, &bufSize)) {
         AMediaFormat_setBuffer(outMeta, AMEDIAFORMAT_KEY_SEI, bufData, bufSize);
     }
+    if (inMeta.findData(kKeyAudioPresentationInfo, &bufType, &bufData, &bufSize)) {
+        AMediaFormat_setBuffer(outMeta, AMEDIAFORMAT_KEY_AUDIO_PRESENTATION_INFO, bufData, bufSize);
+    }
     mbuf->release();
     return AMEDIA_OK;
 }
diff --git a/media/libstagefright/MediaTrack.cpp b/media/libstagefright/MediaTrack.cpp
index 1c1be30..4669cc9 100644
--- a/media/libstagefright/MediaTrack.cpp
+++ b/media/libstagefright/MediaTrack.cpp
@@ -172,6 +172,10 @@
             meta.setData(kKeySEI,
                     MetaDataBase::Type::TYPE_NONE, valbuf->data(), valbuf->size());
         }
+        if (format->mFormat->findBuffer("audio-presentation-info", &valbuf)) {
+            meta.setData(kKeyAudioPresentationInfo,
+                    MetaDataBase::Type::TYPE_NONE, valbuf->data(), valbuf->size());
+        }
     } else {
         *buffer = nullptr;
     }
diff --git a/media/libstagefright/NuMediaExtractor.cpp b/media/libstagefright/NuMediaExtractor.cpp
index f5178dd..779372d 100644
--- a/media/libstagefright/NuMediaExtractor.cpp
+++ b/media/libstagefright/NuMediaExtractor.cpp
@@ -795,30 +795,34 @@
 }
 
 // Return OK if we have received an audio presentation info.
+// Return ERROR_END_OF_STREAM if no tracks are available.
 // Return ERROR_UNSUPPORTED if the track has no audio presentation.
 // Return INVALID_OPERATION if audio presentation metadata version does not match.
 status_t NuMediaExtractor::getAudioPresentations(
-        size_t trackIndex, AudioPresentationCollection *presentations) const {
+        size_t trackIndex, AudioPresentationCollection *presentations) {
     Mutex::Autolock autoLock(mLock);
-
-    if (mImpl == NULL) {
-        return -EINVAL;
+    ssize_t minIndex = fetchAllTrackSamples();
+    if (minIndex < 0) {
+        return ERROR_END_OF_STREAM;
     }
+    for (size_t i = 0; i < mSelectedTracks.size(); ++i) {
+        TrackInfo *info = &mSelectedTracks.editItemAt(i);
 
-    if (trackIndex >= mImpl->countTracks()) {
-        return -ERANGE;
+        if (info->mTrackIndex == trackIndex) {
+            sp<MetaData> meta = new MetaData(info->mSamples.begin()->mBuffer->meta_data());
+
+            uint32_t type;
+            const void *data;
+            size_t size;
+            if (meta != NULL && meta->findData(kKeyAudioPresentationInfo, &type, &data, &size)) {
+                std::istringstream inStream(std::string(static_cast<const char*>(data), size));
+                return deserializeAudioPresentations(&inStream, presentations);
+            }
+            ALOGV("Track %zu does not contain any audio presentation", trackIndex);
+            return ERROR_UNSUPPORTED;
+        }
     }
-
-    sp<MetaData> meta = mImpl->getTrackMetaData(trackIndex);
-
-    uint32_t type;
-    const void *data;
-    size_t size;
-    if (meta != NULL && meta->findData(kKeyAudioPresentationInfo, &type, &data, &size)) {
-        std::istringstream inStream(std::string(static_cast<const char*>(data), size));
-        return deserializeAudioPresentations(&inStream, presentations);
-    }
-    ALOGE("Source does not contain any audio presentation");
+    ALOGV("Source does not contain any audio presentation");
     return ERROR_UNSUPPORTED;
 }
 
diff --git a/media/libstagefright/include/media/stagefright/NuMediaExtractor.h b/media/libstagefright/include/media/stagefright/NuMediaExtractor.h
index 8dc8d38..f34f9b6 100644
--- a/media/libstagefright/include/media/stagefright/NuMediaExtractor.h
+++ b/media/libstagefright/include/media/stagefright/NuMediaExtractor.h
@@ -96,8 +96,7 @@
 
     bool getCachedDuration(int64_t *durationUs, bool *eos) const;
 
-    status_t getAudioPresentations(size_t trackIdx,
-            AudioPresentationCollection *presentations) const;
+    status_t getAudioPresentations(size_t trackIdx, AudioPresentationCollection *presentations);
 
 protected:
     virtual ~NuMediaExtractor();
diff --git a/media/libstagefright/mpeg2ts/ATSParser.cpp b/media/libstagefright/mpeg2ts/ATSParser.cpp
index e9baa1a..7ce9dea 100644
--- a/media/libstagefright/mpeg2ts/ATSParser.cpp
+++ b/media/libstagefright/mpeg2ts/ATSParser.cpp
@@ -217,6 +217,9 @@
     sp<IDescrambler> mDescrambler;
     AudioPresentationCollection mAudioPresentations;
 
+    // Send audio presentations along with access units.
+    void addAudioPresentations(const sp<ABuffer> &buffer);
+
     // Flush accumulated payload if necessary --- i.e. at EOS or at the start of
     // another payload. event is set if the flushed payload is PES with a sync
     // frame.
@@ -1708,6 +1711,13 @@
     return err;
 }
 
+void ATSParser::Stream::addAudioPresentations(const sp<ABuffer> &buffer) {
+    std::ostringstream outStream(std::ios::out);
+    serializeAudioPresentations(mAudioPresentations, &outStream);
+    sp<ABuffer> ap = ABuffer::CreateAsCopy(outStream.str().data(), outStream.str().size());
+    buffer->meta()->setBuffer("audio-presentation-info", ap);
+}
+
 void ATSParser::Stream::onPayloadData(
         unsigned PTS_DTS_flags, uint64_t PTS, uint64_t /* DTS */,
         unsigned PES_scrambling_control,
@@ -1758,8 +1768,10 @@
                     }
                 }
                 mSource = new AnotherPacketSource(meta);
+                if (mAudioPresentations.size() > 0) {
+                    addAudioPresentations(accessUnit);
+                }
                 mSource->queueAccessUnit(accessUnit);
-                mSource->convertAudioPresentationInfoToMetadata(mAudioPresentations);
                 ALOGV("onPayloadData: created AnotherPacketSource PID 0x%08x of type 0x%02x",
                         mElementaryPID, mStreamType);
             }
@@ -1771,8 +1783,10 @@
             if (mSource->getFormat() == NULL) {
                 mSource->setFormat(mQueue->getFormat());
             }
+            if (mAudioPresentations.size() > 0) {
+                addAudioPresentations(accessUnit);
+            }
             mSource->queueAccessUnit(accessUnit);
-            mSource->convertAudioPresentationInfoToMetadata(mAudioPresentations);
         }
 
         // Every access unit has a pesStartOffset queued in |mPesStartOffsets|.
diff --git a/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp b/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp
index e2c5031..3cdd228 100644
--- a/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp
+++ b/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp
@@ -223,6 +223,12 @@
                     kKeyMpegUserData, 0, mpegUserData->data(), mpegUserData->size());
         }
 
+        sp<ABuffer> ap;
+        if (buffer->meta()->findBuffer("audio-presentation-info", &ap) && ap != NULL) {
+            bufmeta.setData(
+                    kKeyAudioPresentationInfo, 0, ap->data(), ap->size());
+        }
+
         int32_t cryptoMode;
         if (buffer->meta()->findInt32("cryptoMode", &cryptoMode)) {
             int32_t cryptoKey;
@@ -697,23 +703,4 @@
     return firstMeta;
 }
 
-void AnotherPacketSource::convertAudioPresentationInfoToMetadata(
-        const AudioPresentationCollection& presentations) {
-    sp<MetaData> meta = getFormat();
-    if (meta == NULL) {
-        return;
-    }
-    if (presentations.empty()) {
-        // Clear audio presentation info in metadata.
-        Mutex::Autolock autoLock(mLock);
-        meta->remove(kKeyAudioPresentationInfo);
-    } else {
-        std::ostringstream outStream(std::ios::out);
-        serializeAudioPresentations(presentations, &outStream);
-        Mutex::Autolock autoLock(mLock);
-        meta->setData(kKeyAudioPresentationInfo, MetaData::TYPE_NONE,
-                outStream.str().data(), outStream.str().size());
-    }
-}
-
 }  // namespace android
diff --git a/media/libstagefright/mpeg2ts/AnotherPacketSource.h b/media/libstagefright/mpeg2ts/AnotherPacketSource.h
index 57a6c33..f4a6acb 100644
--- a/media/libstagefright/mpeg2ts/AnotherPacketSource.h
+++ b/media/libstagefright/mpeg2ts/AnotherPacketSource.h
@@ -85,8 +85,6 @@
     void trimBuffersAfterMeta(const sp<AMessage> &meta);
     sp<AMessage> trimBuffersBeforeMeta(const sp<AMessage> &meta);
 
-    void convertAudioPresentationInfoToMetadata(const AudioPresentationCollection &presentations);
-
 protected:
     virtual ~AnotherPacketSource();
 
diff --git a/media/ndk/NdkMediaExtractor.cpp b/media/ndk/NdkMediaExtractor.cpp
index f697bd1..8296598 100644
--- a/media/ndk/NdkMediaExtractor.cpp
+++ b/media/ndk/NdkMediaExtractor.cpp
@@ -448,6 +448,16 @@
         meta->setBuffer(AMEDIAFORMAT_KEY_MPEG_USER_DATA, mpegUserData);
     }
 
+    const void *audioPresentationsPointer;
+    size_t audioPresentationsLength;
+    if (sampleMeta->findData(
+            kKeyAudioPresentationInfo, &dataType,
+            &audioPresentationsPointer, &audioPresentationsLength)) {
+        sp<ABuffer> audioPresentationsData = ABuffer::CreateAsCopy(
+                audioPresentationsPointer, audioPresentationsLength);
+        meta->setBuffer(AMEDIAFORMAT_KEY_AUDIO_PRESENTATION_INFO, audioPresentationsData);
+    }
+
     return AMEDIA_OK;
 }