Transcoder: Add TrackTranscoder benchmarking.

Adds a benchmarking test that can run TrackTranscoders
without getting samples directly from an extractor and
without writing the transcoded samples to a muxer.

Test: Running the benchmark.
Fixes: 152091444

Change-Id: Ia723f96b9033b1ec9b4502a1242e154bcdbfe9f5
diff --git a/media/libmediatranscoding/transcoder/benchmark/Android.bp b/media/libmediatranscoding/transcoder/benchmark/Android.bp
index b755206..c9a7b57 100644
--- a/media/libmediatranscoding/transcoder/benchmark/Android.bp
+++ b/media/libmediatranscoding/transcoder/benchmark/Android.bp
@@ -1,14 +1,23 @@
+cc_defaults {
+    name: "benchmarkdefaults",
+    shared_libs: ["libmediatranscoder", "libmediandk", "libbase"],
+    static_libs: ["libgoogle-benchmark"],
+}
+
 cc_test {
     name: "MediaTranscoderBenchmark",
     srcs: ["MediaTranscoderBenchmark.cpp"],
-    shared_libs: ["libmediatranscoder", "libmediandk"],
-    static_libs: ["libgoogle-benchmark"],
+    defaults: ["benchmarkdefaults"],
 }
 
 cc_test {
     name: "MediaSampleReaderBenchmark",
     srcs: ["MediaSampleReaderBenchmark.cpp"],
-    shared_libs: ["libmediatranscoder", "libmediandk", "libbase"],
-    static_libs: ["libgoogle-benchmark"],
+    defaults: ["benchmarkdefaults"],
 }
 
+cc_test {
+    name: "MediaTrackTranscoderBenchmark",
+    srcs: ["MediaTrackTranscoderBenchmark.cpp"],
+    defaults: ["benchmarkdefaults"],
+}
diff --git a/media/libmediatranscoding/transcoder/benchmark/MediaTrackTranscoderBenchmark.cpp b/media/libmediatranscoding/transcoder/benchmark/MediaTrackTranscoderBenchmark.cpp
new file mode 100644
index 0000000..e15186c
--- /dev/null
+++ b/media/libmediatranscoding/transcoder/benchmark/MediaTrackTranscoderBenchmark.cpp
@@ -0,0 +1,399 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Native media track transcoder benchmark tests.
+ *
+ * How to run the benchmark:
+ *
+ * 1. Download the media assets from http://go/transcodingbenchmark and push the directory
+ *    ("TranscodingBenchmark") to /data/local/tmp.
+ *
+ * 2. Compile the benchmark and sync to device:
+ *      $ mm -j72 && adb sync
+ *
+ * 3. Run:
+ *      $ adb shell /data/nativetest64/MediaTrackTranscoderBenchmark/MediaTrackTranscoderBenchmark
+ */
+
+// #define LOG_NDEBUG 0
+#define LOG_TAG "MediaTrackTranscoderBenchmark"
+
+#include <android-base/logging.h>
+#include <benchmark/benchmark.h>
+#include <fcntl.h>
+#include <media/MediaSampleReader.h>
+#include <media/MediaSampleReaderNDK.h>
+#include <media/MediaTrackTranscoder.h>
+#include <media/MediaTrackTranscoderCallback.h>
+#include <media/NdkCommon.h>
+#include <media/PassthroughTrackTranscoder.h>
+#include <media/VideoTrackTranscoder.h>
+
+using namespace android;
+
+typedef enum {
+    kVideo,
+    kAudio,
+} MediaType;
+
+class TrackTranscoderCallbacks : public MediaTrackTranscoderCallback {
+public:
+    virtual void onTrackFormatAvailable(const MediaTrackTranscoder* transcoder __unused) override {}
+
+    virtual void onTrackFinished(const MediaTrackTranscoder* transcoder __unused) override {
+        std::unique_lock lock(mMutex);
+        mFinished = true;
+        mCondition.notify_all();
+    }
+
+    virtual void onTrackError(const MediaTrackTranscoder* transcoder __unused,
+                              media_status_t status) override {
+        std::unique_lock lock(mMutex);
+        mFinished = true;
+        mStatus = status;
+        mCondition.notify_all();
+    }
+
+    void waitForTranscodingFinished() {
+        std::unique_lock lock(mMutex);
+        while (!mFinished) {
+            mCondition.wait(lock);
+        }
+    }
+
+    media_status_t mStatus = AMEDIA_OK;
+
+private:
+    std::mutex mMutex;
+    std::condition_variable mCondition;
+    bool mFinished = false;
+};
+
+/**
+ * MockSampleReader holds a ringbuffer of the first samples in the provided source track. Samples
+ * are returned to the caller from the ringbuffer in a round-robin fashion with increasing
+ * timestamps. The number of samples returned before EOS matches the number of frames in the source
+ * track.
+ */
+class MockSampleReader : public MediaSampleReader {
+public:
+    static std::shared_ptr<MediaSampleReader> createFromFd(int fd, size_t offset, size_t size) {
+        AMediaExtractor* extractor = AMediaExtractor_new();
+        media_status_t status = AMediaExtractor_setDataSourceFd(extractor, fd, offset, size);
+        if (status != AMEDIA_OK) return nullptr;
+
+        auto sampleReader = std::shared_ptr<MockSampleReader>(new MockSampleReader(extractor));
+        return sampleReader;
+    }
+
+    AMediaFormat* getFileFormat() override { return AMediaExtractor_getFileFormat(mExtractor); }
+
+    size_t getTrackCount() const override { return AMediaExtractor_getTrackCount(mExtractor); }
+
+    AMediaFormat* getTrackFormat(int trackIndex) override {
+        return AMediaExtractor_getTrackFormat(mExtractor, trackIndex);
+    }
+
+    media_status_t selectTrack(int trackIndex) override {
+        if (mSelectedTrack >= 0) return AMEDIA_ERROR_UNSUPPORTED;
+        mSelectedTrack = trackIndex;
+
+        media_status_t status = AMediaExtractor_selectTrack(mExtractor, trackIndex);
+        if (status != AMEDIA_OK) return status;
+
+        // Get the sample count.
+        AMediaFormat* format = getTrackFormat(trackIndex);
+        const bool haveSampleCount =
+                AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_FRAME_COUNT, &mSampleCount);
+        AMediaFormat_delete(format);
+
+        if (!haveSampleCount) {
+            LOG(ERROR) << "No sample count in track format.";
+            return AMEDIA_ERROR_UNSUPPORTED;
+        }
+
+        // Buffer samples.
+        const int32_t targetBufferCount = 60;
+        std::unique_ptr<uint8_t[]> buffer;
+        MediaSampleInfo info;
+        while (true) {
+            info.presentationTimeUs = AMediaExtractor_getSampleTime(mExtractor);
+            info.flags = AMediaExtractor_getSampleFlags(mExtractor);
+            info.size = AMediaExtractor_getSampleSize(mExtractor);
+
+            // Finish buffering after either reading all the samples in the track or after
+            // completing the GOP satisfying the target count.
+            if (mSamples.size() == mSampleCount ||
+                (mSamples.size() >= targetBufferCount && info.flags & SAMPLE_FLAG_SYNC_SAMPLE)) {
+                break;
+            }
+
+            buffer.reset(new uint8_t[info.size]);
+
+            ssize_t bytesRead = AMediaExtractor_readSampleData(mExtractor, buffer.get(), info.size);
+            if (bytesRead != info.size) {
+                return AMEDIA_ERROR_UNKNOWN;
+            }
+
+            mSamples.emplace_back(std::move(buffer), info);
+
+            AMediaExtractor_advance(mExtractor);
+        }
+
+        mFirstPtsUs = mSamples[0].second.presentationTimeUs;
+        mPtsDiff = mSamples[1].second.presentationTimeUs - mSamples[0].second.presentationTimeUs;
+
+        return AMEDIA_OK;
+    }
+
+    media_status_t setEnforceSequentialAccess(bool enforce __unused) override { return AMEDIA_OK; }
+
+    media_status_t getEstimatedBitrateForTrack(int trackIndex __unused,
+                                               int32_t* bitrate __unused) override {
+        return AMEDIA_ERROR_UNSUPPORTED;
+    }
+
+    media_status_t getSampleInfoForTrack(int trackIndex, MediaSampleInfo* info) override {
+        if (trackIndex != mSelectedTrack) return AMEDIA_ERROR_INVALID_PARAMETER;
+
+        if (mCurrentSampleIndex >= mSampleCount) {
+            info->presentationTimeUs = 0;
+            info->size = 0;
+            info->flags = SAMPLE_FLAG_END_OF_STREAM;
+            return AMEDIA_ERROR_END_OF_STREAM;
+        }
+
+        *info = mSamples[mCurrentSampleIndex % mSamples.size()].second;
+        info->presentationTimeUs = mFirstPtsUs + mCurrentSampleIndex * mPtsDiff;
+        return AMEDIA_OK;
+    }
+
+    media_status_t readSampleDataForTrack(int trackIndex, uint8_t* buffer,
+                                          size_t bufferSize) override {
+        if (trackIndex != mSelectedTrack) return AMEDIA_ERROR_INVALID_PARAMETER;
+
+        if (mCurrentSampleIndex >= mSampleCount) return AMEDIA_ERROR_END_OF_STREAM;
+
+        auto& p = mSamples[mCurrentSampleIndex % mSamples.size()];
+
+        if (bufferSize < p.second.size) return AMEDIA_ERROR_INVALID_PARAMETER;
+        memcpy(buffer, p.first.get(), bufferSize);
+
+        advanceTrack(trackIndex);
+        return AMEDIA_OK;
+    }
+
+    void advanceTrack(int trackIndex) {
+        if (trackIndex != mSelectedTrack) return;
+        ++mCurrentSampleIndex;
+    }
+
+    virtual ~MockSampleReader() override { AMediaExtractor_delete(mExtractor); }
+
+private:
+    MockSampleReader(AMediaExtractor* extractor) : mExtractor(extractor) {}
+    AMediaExtractor* mExtractor = nullptr;
+    int32_t mSampleCount = 0;
+    std::vector<std::pair<std::unique_ptr<uint8_t[]>, MediaSampleInfo>> mSamples;
+    int mSelectedTrack = -1;
+    int32_t mCurrentSampleIndex = 0;
+    int64_t mFirstPtsUs = 0;
+    int64_t mPtsDiff = 0;
+};
+
+static std::shared_ptr<AMediaFormat> GetDefaultTrackFormat(MediaType mediaType,
+                                                           AMediaFormat* sourceFormat) {
+    // Default video config.
+    static constexpr int32_t kVideoBitRate = 20 * 1000 * 1000;  // 20 mbps
+    static constexpr float kVideoFrameRate = 30.0f;             // 30 fps
+
+    AMediaFormat* format = nullptr;
+
+    if (mediaType == kVideo) {
+        format = AMediaFormat_new();
+        AMediaFormat_copy(format, sourceFormat);
+        AMediaFormat_setString(format, AMEDIAFORMAT_KEY_MIME, AMEDIA_MIMETYPE_VIDEO_AVC);
+        AMediaFormat_setInt32(format, AMEDIAFORMAT_KEY_BIT_RATE, kVideoBitRate);
+        AMediaFormat_setFloat(format, AMEDIAFORMAT_KEY_FRAME_RATE, kVideoFrameRate);
+    }
+    // nothing for audio.
+
+    return std::shared_ptr<AMediaFormat>(format, &AMediaFormat_delete);
+}
+
+/** Gets a MediaSampleReader for the source file */
+static std::shared_ptr<MediaSampleReader> GetSampleReader(const std::string& srcFileName,
+                                                          bool mock) {
+    // Asset directory
+    static const std::string kAssetDirectory = "/data/local/tmp/TranscodingBenchmark/";
+
+    int srcFd = 0;
+    std::string srcPath = kAssetDirectory + srcFileName;
+
+    if ((srcFd = open(srcPath.c_str(), O_RDONLY)) < 0) {
+        return nullptr;
+    }
+
+    const size_t fileSize = lseek(srcFd, 0, SEEK_END);
+    lseek(srcFd, 0, SEEK_SET);
+
+    std::shared_ptr<MediaSampleReader> sampleReader;
+
+    if (mock) {
+        sampleReader = MockSampleReader::createFromFd(srcFd, 0 /* offset */, fileSize);
+    } else {
+        sampleReader = MediaSampleReaderNDK::createFromFd(srcFd, 0 /* offset */, fileSize);
+    }
+
+    if (srcFd > 0) close(srcFd);
+    return sampleReader;
+}
+
+/**
+ * Configures a MediaTrackTranscoder with an empty sample consumer so that the samples are returned
+ * to the transcoder immediately.
+ */
+static void ConfigureEmptySampleConsumer(const std::shared_ptr<MediaTrackTranscoder>& transcoder,
+                                         uint32_t& sampleCount) {
+    transcoder->setSampleConsumer([&sampleCount](const std::shared_ptr<MediaSample>& sample) {
+        if (!(sample->info.flags & SAMPLE_FLAG_CODEC_CONFIG) && sample->info.size > 0) {
+            ++sampleCount;
+        }
+    });
+}
+
+/**
+ * Configures a MediaTrackTranscoder with the provided MediaSampleReader, reading from the first
+ * track that matches the specified media type.
+ */
+static bool ConfigureSampleReader(const std::shared_ptr<MediaTrackTranscoder>& transcoder,
+                                  const std::shared_ptr<MediaSampleReader>& sampleReader,
+                                  MediaType mediaType) {
+    int srcTrackIndex = -1;
+    std::shared_ptr<AMediaFormat> srcTrackFormat = nullptr;
+
+    for (int trackIndex = 0; trackIndex < sampleReader->getTrackCount(); ++trackIndex) {
+        AMediaFormat* trackFormat = sampleReader->getTrackFormat(trackIndex);
+
+        const char* mime = nullptr;
+        AMediaFormat_getString(trackFormat, AMEDIAFORMAT_KEY_MIME, &mime);
+
+        if ((mediaType == kVideo && strncmp(mime, "video/", 6) == 0) ||
+            (mediaType == kAudio && strncmp(mime, "audio/", 6) == 0)) {
+            srcTrackIndex = trackIndex;
+            srcTrackFormat = std::shared_ptr<AMediaFormat>(trackFormat, &AMediaFormat_delete);
+            break;
+        }
+        AMediaFormat_delete(trackFormat);
+    }
+
+    if (srcTrackIndex == -1) {
+        LOG(ERROR) << "No matching source track found";
+        return false;
+    }
+
+    media_status_t status = sampleReader->selectTrack(srcTrackIndex);
+    if (status != AMEDIA_OK) {
+        LOG(ERROR) << "Unable to select track";
+        return false;
+    }
+
+    auto destinationFormat = GetDefaultTrackFormat(mediaType, srcTrackFormat.get());
+    status = transcoder->configure(sampleReader, srcTrackIndex, destinationFormat);
+    if (status != AMEDIA_OK) {
+        LOG(ERROR) << "transcoder configure returned " << status;
+        return false;
+    }
+
+    return true;
+}
+
+static void BenchmarkTranscoder(benchmark::State& state, const std::string& srcFileName,
+                                bool mockReader, MediaType mediaType) {
+    for (auto _ : state) {
+        std::shared_ptr<TrackTranscoderCallbacks> callbacks =
+                std::make_shared<TrackTranscoderCallbacks>();
+        std::shared_ptr<MediaTrackTranscoder> transcoder;
+
+        if (mediaType == kVideo) {
+            transcoder = VideoTrackTranscoder::create(callbacks);
+        } else {
+            transcoder = std::make_shared<PassthroughTrackTranscoder>(callbacks);
+        }
+
+        std::shared_ptr<MediaSampleReader> sampleReader = GetSampleReader(srcFileName, mockReader);
+        if (sampleReader == nullptr) {
+            state.SkipWithError("Unable to create sample reader");
+            return;
+        }
+
+        if (!ConfigureSampleReader(transcoder, sampleReader, mediaType)) {
+            state.SkipWithError("Unable to configure the transcoder");
+            return;
+        }
+
+        uint32_t sampleCount = 0;
+        ConfigureEmptySampleConsumer(transcoder, sampleCount);
+
+        if (!transcoder->start()) {
+            state.SkipWithError("Unable to start the transcoder");
+            return;
+        }
+
+        callbacks->waitForTranscodingFinished();
+        transcoder->stop();
+
+        if (callbacks->mStatus != AMEDIA_OK) {
+            state.SkipWithError("Transcoder failed with error");
+            return;
+        }
+
+        LOG(DEBUG) << "Number of samples received: " << sampleCount;
+        state.counters["FrameRate"] = benchmark::Counter(sampleCount, benchmark::Counter::kIsRate);
+    }
+}
+
+// Benchmark registration wrapper for transcoding.
+#define TRANSCODER_BENCHMARK(func) \
+    BENCHMARK(func)->UseRealTime()->MeasureProcessCPUTime()->Unit(benchmark::kMillisecond)
+
+static void BM_VideoTranscode_AVC2AVC_NoMuxer(benchmark::State& state) {
+    const char* srcFile = "video_1920x1080_3648frame_h264_22Mbps_30fps_aac.mp4";
+    BenchmarkTranscoder(state, srcFile, false /* mockReader */, kVideo);
+}
+
+static void BM_VideoTranscode_AVC2AVC_NoMuxer_NoExtractor(benchmark::State& state) {
+    const char* srcFile = "video_1920x1080_3648frame_h264_22Mbps_30fps_aac.mp4";
+    BenchmarkTranscoder(state, srcFile, true /* mockReader */, kVideo);
+}
+
+static void BM_VideoTranscode_HEVC2AVC_NoMuxer(benchmark::State& state) {
+    const char* srcFile = "video_1920x1080_3863frame_hevc_4Mbps_30fps_aac.mp4";
+    BenchmarkTranscoder(state, srcFile, false /* mockReader */, kVideo);
+}
+
+static void BM_VideoTranscode_HEVC2AVC_NoMuxer_NoExtractor(benchmark::State& state) {
+    const char* srcFile = "video_1920x1080_3863frame_hevc_4Mbps_30fps_aac.mp4";
+    BenchmarkTranscoder(state, srcFile, true /* mockReader */, kVideo);
+}
+
+TRANSCODER_BENCHMARK(BM_VideoTranscode_AVC2AVC_NoMuxer);
+TRANSCODER_BENCHMARK(BM_VideoTranscode_AVC2AVC_NoMuxer_NoExtractor);
+TRANSCODER_BENCHMARK(BM_VideoTranscode_HEVC2AVC_NoMuxer);
+TRANSCODER_BENCHMARK(BM_VideoTranscode_HEVC2AVC_NoMuxer_NoExtractor);
+
+BENCHMARK_MAIN();