Transcoder: Improve video transcoding with audio passthrough performance

MediaSampleWriter was stalling the video transcoding when processing
chunks of audio samples at a time. Since the underlying muxer
buffers ~1 sec worth of data and interleaves tracks anyway the
sample writer now prioritizes pulling samples from the track that is
farthest behind while making sure the tracks are kept reasonably
in sync, i.e. max ~1 second divergence.

This change also raises the priority of the video transcoding thread
since that will be the bottleneck when audio is pased through.

Test: Unit tests
Fixes: 160268606
Change-Id: I004583b2a31a57882ea543072be321f9f1347508
diff --git a/media/libmediatranscoding/transcoder/MediaSampleQueue.cpp b/media/libmediatranscoding/transcoder/MediaSampleQueue.cpp
index 691ee1c..b085c98 100644
--- a/media/libmediatranscoding/transcoder/MediaSampleQueue.cpp
+++ b/media/libmediatranscoding/transcoder/MediaSampleQueue.cpp
@@ -47,6 +47,11 @@
     return mAborted;
 }
 
+bool MediaSampleQueue::isEmpty() {
+    std::scoped_lock<std::mutex> lock(mMutex);
+    return mSampleQueue.empty();
+}
+
 void MediaSampleQueue::abort() {
     std::scoped_lock<std::mutex> lock(mMutex);
     // Clear the queue and notify consumers.
diff --git a/media/libmediatranscoding/transcoder/MediaSampleWriter.cpp b/media/libmediatranscoding/transcoder/MediaSampleWriter.cpp
index 3676d73..bb0da88 100644
--- a/media/libmediatranscoding/transcoder/MediaSampleWriter.cpp
+++ b/media/libmediatranscoding/transcoder/MediaSampleWriter.cpp
@@ -127,18 +127,16 @@
         durationUs = 0;
     }
 
-    const char* mime = nullptr;
-    const bool isVideo = AMediaFormat_getString(trackFormat.get(), AMEDIAFORMAT_KEY_MIME, &mime) &&
-                         (strncmp(mime, "video/", 6) == 0);
-
-    mTracks.emplace_back(sampleQueue, static_cast<size_t>(trackIndex), durationUs, isVideo);
+    mAllTracks.push_back(std::make_unique<TrackRecord>(sampleQueue, static_cast<size_t>(trackIndex),
+                                                       durationUs));
+    mSortedTracks.insert(mAllTracks.back().get());
     return true;
 }
 
 bool MediaSampleWriter::start() {
     std::scoped_lock lock(mStateMutex);
 
-    if (mTracks.size() == 0) {
+    if (mAllTracks.size() == 0) {
         LOG(ERROR) << "No tracks to write.";
         return false;
     } else if (mState != INITIALIZED) {
@@ -165,8 +163,8 @@
     }
 
     // Stop the sources, and wait for thread to join.
-    for (auto& track : mTracks) {
-        track.mSampleQueue->abort();
+    for (auto& track : mAllTracks) {
+        track->mSampleQueue->abort();
     }
     mThread.join();
     mState = STOPPED;
@@ -193,76 +191,102 @@
     return writeStatus != AMEDIA_OK ? writeStatus : muxerStatus;
 }
 
+std::multiset<MediaSampleWriter::TrackRecord*>::iterator MediaSampleWriter::getNextOutputTrack() {
+    // Find the first track that has samples ready in its queue AND is not more than
+    // mMaxTrackDivergenceUs ahead of the slowest track. If no such track exists then return the
+    // slowest track and let the writer wait for samples to become ready. Note that mSortedTracks is
+    // sorted by each track's previous sample timestamp in ascending order.
+    auto slowestTrack = mSortedTracks.begin();
+    if (slowestTrack == mSortedTracks.end() || !(*slowestTrack)->mSampleQueue->isEmpty()) {
+        return slowestTrack;
+    }
+
+    const int64_t slowestTimeUs = (*slowestTrack)->mPrevSampleTimeUs;
+    int64_t divergenceUs;
+
+    for (auto it = std::next(slowestTrack); it != mSortedTracks.end(); ++it) {
+        // If the current track has diverged then the rest will have too, so we can stop the search.
+        // If not and it has samples ready then return it, otherwise keep looking.
+        if (__builtin_sub_overflow((*it)->mPrevSampleTimeUs, slowestTimeUs, &divergenceUs) ||
+            divergenceUs >= mMaxTrackDivergenceUs) {
+            break;
+        } else if (!(*it)->mSampleQueue->isEmpty()) {
+            return it;
+        }
+    }
+
+    // No track with pending samples within acceptable time interval was found, so let the writer
+    // wait for the slowest track to produce a new sample.
+    return slowestTrack;
+}
+
 media_status_t MediaSampleWriter::runWriterLoop() {
     AMediaCodecBufferInfo bufferInfo;
-    uint32_t segmentEndTimeUs = mTrackSegmentLengthUs;
-    bool samplesLeft = true;
     int32_t lastProgressUpdate = 0;
 
     // Set the "primary" track that will be used to determine progress to the track with longest
     // duration.
     int primaryTrackIndex = -1;
     int64_t longestDurationUs = 0;
-    for (int trackIndex = 0; trackIndex < mTracks.size(); ++trackIndex) {
-        if (mTracks[trackIndex].mDurationUs > longestDurationUs) {
-            primaryTrackIndex = trackIndex;
-            longestDurationUs = mTracks[trackIndex].mDurationUs;
+    for (auto& track : mAllTracks) {
+        if (track->mDurationUs > longestDurationUs) {
+            primaryTrackIndex = track->mTrackIndex;
+            longestDurationUs = track->mDurationUs;
         }
     }
 
-    while (samplesLeft) {
-        samplesLeft = false;
-        for (auto& track : mTracks) {
-            if (track.mReachedEos) continue;
+    while (true) {
+        auto outputTrackIter = getNextOutputTrack();
 
-            std::shared_ptr<MediaSample> sample;
-            do {
-                if (track.mSampleQueue->dequeue(&sample)) {
-                    // Track queue was aborted.
-                    return AMEDIA_ERROR_UNKNOWN;  // TODO(lnilsson): Custom error code.
-                } else if (sample->info.flags & SAMPLE_FLAG_END_OF_STREAM) {
-                    // Track reached end of stream.
-                    track.mReachedEos = true;
-
-                    // Preserve source track duration by setting the appropriate timestamp on the
-                    // empty End-Of-Stream sample.
-                    if (track.mDurationUs > 0 && track.mFirstSampleTimeSet) {
-                        sample->info.presentationTimeUs =
-                                track.mDurationUs + track.mFirstSampleTimeUs;
-                    }
-                } else {
-                    samplesLeft = true;
-                }
-
-                track.mPrevSampleTimeUs = sample->info.presentationTimeUs;
-                if (!track.mFirstSampleTimeSet) {
-                    // Record the first sample's timestamp in order to translate duration to EOS
-                    // time for tracks that does not start at 0.
-                    track.mFirstSampleTimeUs = sample->info.presentationTimeUs;
-                    track.mFirstSampleTimeSet = true;
-                }
-
-                bufferInfo.offset = sample->dataOffset;
-                bufferInfo.size = sample->info.size;
-                bufferInfo.flags = sample->info.flags;
-                bufferInfo.presentationTimeUs = sample->info.presentationTimeUs;
-
-                media_status_t status =
-                        mMuxer->writeSampleData(track.mTrackIndex, sample->buffer, &bufferInfo);
-                if (status != AMEDIA_OK) {
-                    LOG(ERROR) << "writeSampleData returned " << status;
-                    return status;
-                }
-
-            } while (sample->info.presentationTimeUs < segmentEndTimeUs && !track.mReachedEos);
+        // Exit if all tracks have reached end of stream.
+        if (outputTrackIter == mSortedTracks.end()) {
+            break;
         }
 
-        // TODO(lnilsson): Add option to toggle progress reporting on/off.
-        if (primaryTrackIndex >= 0) {
-            const TrackRecord& track = mTracks[primaryTrackIndex];
+        // Remove the track from the set, update it, and then reinsert it to keep the set in order.
+        TrackRecord* track = *outputTrackIter;
+        mSortedTracks.erase(outputTrackIter);
 
-            const int64_t elapsed = track.mPrevSampleTimeUs - track.mFirstSampleTimeUs;
-            int32_t progress = (elapsed * 100) / track.mDurationUs;
+        std::shared_ptr<MediaSample> sample;
+        if (track->mSampleQueue->dequeue(&sample)) {
+            // Track queue was aborted.
+            return AMEDIA_ERROR_UNKNOWN;  // TODO(lnilsson): Custom error code.
+        } else if (sample->info.flags & SAMPLE_FLAG_END_OF_STREAM) {
+            // Track reached end of stream.
+            track->mReachedEos = true;
+
+            // Preserve source track duration by setting the appropriate timestamp on the
+            // empty End-Of-Stream sample.
+            if (track->mDurationUs > 0 && track->mFirstSampleTimeSet) {
+                sample->info.presentationTimeUs = track->mDurationUs + track->mFirstSampleTimeUs;
+            }
+        }
+
+        track->mPrevSampleTimeUs = sample->info.presentationTimeUs;
+        if (!track->mFirstSampleTimeSet) {
+            // Record the first sample's timestamp in order to translate duration to EOS
+            // time for tracks that does not start at 0.
+            track->mFirstSampleTimeUs = sample->info.presentationTimeUs;
+            track->mFirstSampleTimeSet = true;
+        }
+
+        bufferInfo.offset = sample->dataOffset;
+        bufferInfo.size = sample->info.size;
+        bufferInfo.flags = sample->info.flags;
+        bufferInfo.presentationTimeUs = sample->info.presentationTimeUs;
+
+        media_status_t status =
+                mMuxer->writeSampleData(track->mTrackIndex, sample->buffer, &bufferInfo);
+        if (status != AMEDIA_OK) {
+            LOG(ERROR) << "writeSampleData returned " << status;
+            return status;
+        }
+        sample.reset();
+
+        // TODO(lnilsson): Add option to toggle progress reporting on/off.
+        if (track->mTrackIndex == primaryTrackIndex) {
+            const int64_t elapsed = track->mPrevSampleTimeUs - track->mFirstSampleTimeUs;
+            int32_t progress = (elapsed * 100) / track->mDurationUs;
             progress = std::clamp(progress, 0, 100);
 
             if (progress > lastProgressUpdate) {
@@ -273,7 +297,9 @@
             }
         }
 
-        segmentEndTimeUs += mTrackSegmentLengthUs;
+        if (!track->mReachedEos) {
+            mSortedTracks.insert(track);
+        }
     }
 
     return AMEDIA_OK;
diff --git a/media/libmediatranscoding/transcoder/VideoTrackTranscoder.cpp b/media/libmediatranscoding/transcoder/VideoTrackTranscoder.cpp
index 65dcad3..d2a7154 100644
--- a/media/libmediatranscoding/transcoder/VideoTrackTranscoder.cpp
+++ b/media/libmediatranscoding/transcoder/VideoTrackTranscoder.cpp
@@ -19,6 +19,7 @@
 
 #include <android-base/logging.h>
 #include <media/VideoTrackTranscoder.h>
+#include <utils/AndroidThreads.h>
 
 namespace android {
 
@@ -437,6 +438,8 @@
 }
 
 media_status_t VideoTrackTranscoder::runTranscodeLoop() {
+    androidSetThreadPriority(0 /* tid (0 = current) */, ANDROID_PRIORITY_VIDEO);
+
     // Push start decoder and encoder as two messages, so that these are subject to the
     // stop request as well. If the job is cancelled (or paused) immediately after start,
     // we don't need to waste time start then stop the codecs.
diff --git a/media/libmediatranscoding/transcoder/benchmark/MediaTranscoderBenchmark.cpp b/media/libmediatranscoding/transcoder/benchmark/MediaTranscoderBenchmark.cpp
index a680ea3..b31b675 100644
--- a/media/libmediatranscoding/transcoder/benchmark/MediaTranscoderBenchmark.cpp
+++ b/media/libmediatranscoding/transcoder/benchmark/MediaTranscoderBenchmark.cpp
@@ -161,6 +161,7 @@
         }
 
         if (!callbacks->waitForTranscodingFinished()) {
+            transcoder->cancel();
             state.SkipWithError("Transcoder timed out");
             goto exit;
         }
diff --git a/media/libmediatranscoding/transcoder/include/media/MediaSampleQueue.h b/media/libmediatranscoding/transcoder/include/media/MediaSampleQueue.h
index dc22423..c6cf1a4 100644
--- a/media/libmediatranscoding/transcoder/include/media/MediaSampleQueue.h
+++ b/media/libmediatranscoding/transcoder/include/media/MediaSampleQueue.h
@@ -50,6 +50,12 @@
     bool dequeue(std::shared_ptr<MediaSample>* sample /* nonnull */);
 
     /**
+     * Checks if the queue currently holds any media samples.
+     * @return True if the queue is empty or has been aborted. False otherwise.
+     */
+    bool isEmpty();
+
+    /**
      * Aborts the queue operation. This clears the queue and notifies waiting consumers. After the
      * has been aborted it is not possible to enqueue more samples, and dequeue will return null.
      */
diff --git a/media/libmediatranscoding/transcoder/include/media/MediaSampleWriter.h b/media/libmediatranscoding/transcoder/include/media/MediaSampleWriter.h
index 92ddc2f..d4b1fcf 100644
--- a/media/libmediatranscoding/transcoder/include/media/MediaSampleWriter.h
+++ b/media/libmediatranscoding/transcoder/include/media/MediaSampleWriter.h
@@ -26,6 +26,7 @@
 #include <functional>
 #include <memory>
 #include <mutex>
+#include <set>
 #include <thread>
 
 namespace android {
@@ -61,15 +62,18 @@
 };
 
 /**
- * MediaSampleWriter writes samples in interleaved segments of a configurable duration.
- * Each track have its own MediaSampleQueue from which samples are dequeued by the sample writer in
- * output order. The dequeued samples are written to an instance of the writer's muxer interface.
- * The default muxer interface implementation is based directly on AMediaMuxer.
+ * MediaSampleWriter writes samples to a muxer while keeping its input sources synchronized. Each
+ * source track have its own MediaSampleQueue from which samples are dequeued by the sample writer
+ * and written to the muxer. The sample writer always prioritizes dequeueing samples from the source
+ * track that is farthest behind by comparing sample timestamps. If the slowest track does not have
+ * any samples pending the writer moves on to the next track but never allows tracks to diverge more
+ * than a configurable duration of time. The default muxer interface implementation is based
+ * directly on AMediaMuxer.
  */
 class MediaSampleWriter {
 public:
-    /** The default segment length. */
-    static constexpr uint32_t kDefaultTrackSegmentLengthUs = 1 * 1000 * 1000;  // 1 sec.
+    /** The default maximum track divergence in microseconds. */
+    static constexpr uint32_t kDefaultMaxTrackDivergenceUs = 1 * 1000 * 1000;  // 1 second.
 
     /** Callback interface. */
     class CallbackInterface {
@@ -87,14 +91,14 @@
     };
 
     /**
-     * Constructor with custom segment length.
-     * @param trackSegmentLengthUs The segment length to use for this MediaSampleWriter.
+     * Constructor with custom maximum track divergence.
+     * @param maxTrackDivergenceUs The maximum track divergence in microseconds.
      */
-    MediaSampleWriter(uint32_t trackSegmentLengthUs)
-          : mTrackSegmentLengthUs(trackSegmentLengthUs), mMuxer(nullptr), mState(UNINITIALIZED){};
+    MediaSampleWriter(uint32_t maxTrackDivergenceUs)
+          : mMaxTrackDivergenceUs(maxTrackDivergenceUs), mMuxer(nullptr), mState(UNINITIALIZED){};
 
-    /** Constructor using the default segment length. */
-    MediaSampleWriter() : MediaSampleWriter(kDefaultTrackSegmentLengthUs){};
+    /** Constructor using the default maximum track divergence. */
+    MediaSampleWriter() : MediaSampleWriter(kDefaultMaxTrackDivergenceUs){};
 
     /** Destructor. */
     ~MediaSampleWriter();
@@ -147,20 +151,16 @@
     bool stop();
 
 private:
-    media_status_t writeSamples();
-    media_status_t runWriterLoop();
-
     struct TrackRecord {
         TrackRecord(const std::shared_ptr<MediaSampleQueue>& sampleQueue, size_t trackIndex,
-                    int64_t durationUs, bool isVideo)
+                    int64_t durationUs)
               : mSampleQueue(sampleQueue),
                 mTrackIndex(trackIndex),
                 mDurationUs(durationUs),
                 mFirstSampleTimeUs(0),
-                mPrevSampleTimeUs(0),
+                mPrevSampleTimeUs(INT64_MIN),
                 mFirstSampleTimeSet(false),
-                mReachedEos(false),
-                mIsVideo(isVideo) {}
+                mReachedEos(false) {}
 
         std::shared_ptr<MediaSampleQueue> mSampleQueue;
         const size_t mTrackIndex;
@@ -169,13 +169,19 @@
         int64_t mPrevSampleTimeUs;
         bool mFirstSampleTimeSet;
         bool mReachedEos;
-        bool mIsVideo;
+
+        struct compare {
+            bool operator()(const TrackRecord* lhs, const TrackRecord* rhs) const {
+                return lhs->mPrevSampleTimeUs < rhs->mPrevSampleTimeUs;
+            }
+        };
     };
 
-    const uint32_t mTrackSegmentLengthUs;
+    const uint32_t mMaxTrackDivergenceUs;
     std::weak_ptr<CallbackInterface> mCallbacks;
     std::shared_ptr<MediaSampleWriterMuxerInterface> mMuxer;
-    std::vector<TrackRecord> mTracks;
+    std::vector<std::unique_ptr<TrackRecord>> mAllTracks;
+    std::multiset<TrackRecord*, TrackRecord::compare> mSortedTracks;
     std::thread mThread;
 
     std::mutex mStateMutex;
@@ -185,6 +191,10 @@
         STARTED,
         STOPPED,
     } mState GUARDED_BY(mStateMutex);
+
+    media_status_t writeSamples();
+    media_status_t runWriterLoop();
+    std::multiset<TrackRecord*>::iterator getNextOutputTrack();
 };
 
 }  // namespace android
diff --git a/media/libmediatranscoding/transcoder/tests/MediaSampleQueueTests.cpp b/media/libmediatranscoding/transcoder/tests/MediaSampleQueueTests.cpp
index 2046ca0..6357e4d 100644
--- a/media/libmediatranscoding/transcoder/tests/MediaSampleQueueTests.cpp
+++ b/media/libmediatranscoding/transcoder/tests/MediaSampleQueueTests.cpp
@@ -46,10 +46,12 @@
 
     static constexpr int kNumSamples = 4;
     MediaSampleQueue sampleQueue;
+    EXPECT_TRUE(sampleQueue.isEmpty());
 
     // Enqueue loop.
     for (int i = 0; i < kNumSamples; ++i) {
         sampleQueue.enqueue(newSample(i));
+        EXPECT_FALSE(sampleQueue.isEmpty());
     }
 
     // Dequeue loop.
@@ -60,6 +62,7 @@
         EXPECT_EQ(sample->bufferId, i);
         EXPECT_FALSE(aborted);
     }
+    EXPECT_TRUE(sampleQueue.isEmpty());
 }
 
 TEST_F(MediaSampleQueueTests, TestInterleavedDequeueOrder) {
@@ -71,12 +74,14 @@
     // Enqueue and dequeue.
     for (int i = 0; i < kNumSamples; ++i) {
         sampleQueue.enqueue(newSample(i));
+        EXPECT_FALSE(sampleQueue.isEmpty());
 
         std::shared_ptr<MediaSample> sample;
         bool aborted = sampleQueue.dequeue(&sample);
         EXPECT_NE(sample, nullptr);
         EXPECT_EQ(sample->bufferId, i);
         EXPECT_FALSE(aborted);
+        EXPECT_TRUE(sampleQueue.isEmpty());
     }
 }
 
@@ -98,6 +103,7 @@
     EXPECT_NE(sample, nullptr);
     EXPECT_EQ(sample->bufferId, 1);
     EXPECT_FALSE(aborted);
+    EXPECT_TRUE(sampleQueue.isEmpty());
 
     enqueueThread.join();
 }
@@ -160,7 +166,9 @@
         EXPECT_FALSE(bufferReleased[i]);
     }
 
+    EXPECT_FALSE(sampleQueue.isEmpty());
     sampleQueue.abort();
+    EXPECT_TRUE(sampleQueue.isEmpty());
 
     for (int i = 0; i < kNumSamples; ++i) {
         EXPECT_TRUE(bufferReleased[i]);
diff --git a/media/libmediatranscoding/transcoder/tests/MediaSampleWriterTests.cpp b/media/libmediatranscoding/transcoder/tests/MediaSampleWriterTests.cpp
index c82ec28..c6e35c4 100644
--- a/media/libmediatranscoding/transcoder/tests/MediaSampleWriterTests.cpp
+++ b/media/libmediatranscoding/transcoder/tests/MediaSampleWriterTests.cpp
@@ -78,7 +78,22 @@
         return {.type = Event::WriteSample, .trackIndex = trackIndex, .data = data, .info = *info};
     }
 
-    const Event& popEvent() {
+    static Event WriteSampleWithPts(size_t trackIndex, int64_t pts) {
+        return {.type = Event::WriteSample, .trackIndex = trackIndex, .info = {0, 0, pts, 0}};
+    }
+
+    void pushEvent(const Event& e) {
+        std::unique_lock<std::mutex> lock(mMutex);
+        mEventQueue.push_back(e);
+        mCondition.notify_one();
+    }
+
+    const Event& popEvent(bool wait = false) {
+        std::unique_lock<std::mutex> lock(mMutex);
+        while (wait && mEventQueue.empty()) {
+            mCondition.wait_for(lock, std::chrono::milliseconds(200));
+        }
+
         if (mEventQueue.empty()) {
             mPoppedEvent = NoEvent;
         } else {
@@ -92,6 +107,8 @@
     Event mPoppedEvent;
     std::list<Event> mEventQueue;
     ssize_t mTrackCount = 0;
+    std::mutex mMutex;
+    std::condition_variable mCondition;
 };
 
 bool operator==(const AMediaCodecBufferInfo& lhs, const AMediaCodecBufferInfo& rhs) {
@@ -250,6 +267,10 @@
         return newSample(ptsUs, 0, sampleId, sampleId, reinterpret_cast<const uint8_t*>(sampleId));
     }
 
+    static std::shared_ptr<MediaSample> newSampleWithPtsOnly(int64_t ptsUs) {
+        return newSample(ptsUs, 0, 0, 0, nullptr);
+    }
+
     void SetUp() override {
         LOG(DEBUG) << "MediaSampleWriterTests set up";
         mTestMuxer = std::make_shared<TestMuxer>();
@@ -345,10 +366,9 @@
 }
 
 TEST_F(MediaSampleWriterTests, TestProgressUpdate) {
-    static constexpr uint32_t kSegmentLengthUs = 1;
     const TestMediaSource& mediaSource = getMediaSource();
 
-    MediaSampleWriter writer{kSegmentLengthUs};
+    MediaSampleWriter writer{};
     EXPECT_TRUE(writer.init(mTestMuxer, mTestCallbacks));
 
     std::shared_ptr<AMediaFormat> videoFormat =
@@ -370,9 +390,7 @@
 }
 
 TEST_F(MediaSampleWriterTests, TestInterleaving) {
-    static constexpr uint32_t kSegmentLength = MediaSampleWriter::kDefaultTrackSegmentLengthUs;
-
-    MediaSampleWriter writer{kSegmentLength};
+    MediaSampleWriter writer{};
     EXPECT_TRUE(writer.init(mTestMuxer, mTestCallbacks));
 
     // Use two tracks for this test.
@@ -398,18 +416,19 @@
     };
 
     addSampleToTrackWithPts(0, 0);
-    addSampleToTrackWithPts(0, kSegmentLength / 2);
-    addSampleToTrackWithPts(0, kSegmentLength);  // Track 0 reached 1st segment end
+    addSampleToTrackWithPts(1, 4);
 
-    addSampleToTrackWithPts(1, 0);
-    addSampleToTrackWithPts(1, kSegmentLength);  // Track 1 reached 1st segment end
+    addSampleToTrackWithPts(0, 1);
+    addSampleToTrackWithPts(0, 2);
+    addSampleToTrackWithPts(0, 3);
+    addSampleToTrackWithPts(0, 10);
 
-    addSampleToTrackWithPts(0, kSegmentLength * 2);  // Track 0 reached 2nd segment end
+    addSampleToTrackWithPts(1, 5);
+    addSampleToTrackWithPts(1, 6);
+    addSampleToTrackWithPts(1, 11);
 
-    addSampleToTrackWithPts(1, kSegmentLength + 1);
-    addSampleToTrackWithPts(1, kSegmentLength * 2);  // Track 1 reached 2nd segment end
-
-    addSampleToTrackWithPts(0, kSegmentLength * 2 + 1);
+    addSampleToTrackWithPts(0, 12);
+    addSampleToTrackWithPts(1, 13);
 
     for (int trackIndex = 0; trackIndex < kNumTracks; ++trackIndex) {
         sampleQueues[trackIndex]->enqueue(newSampleEos());
@@ -443,7 +462,10 @@
         int64_t duration = 0;
         AMediaFormat_getInt64(trackFormat.get(), AMEDIAFORMAT_KEY_DURATION, &duration);
 
-        const AMediaCodecBufferInfo info = {0, 0, duration, AMEDIACODEC_BUFFER_FLAG_END_OF_STREAM};
+        // EOS timestamp = first sample timestamp + duration.
+        const int64_t endTime = duration + (trackIndex == 1 ? 4 : 0);
+        const AMediaCodecBufferInfo info = {0, 0, endTime, AMEDIACODEC_BUFFER_FLAG_END_OF_STREAM};
+
         EXPECT_EQ(mTestMuxer->popEvent(), TestMuxer::WriteSample(trackIndex, nullptr, &info));
     }
 
@@ -452,6 +474,124 @@
     EXPECT_TRUE(mTestCallbacks->hasFinished());
 }
 
+TEST_F(MediaSampleWriterTests, TestMaxDivergence) {
+    static constexpr uint32_t kMaxDivergenceUs = 10;
+
+    MediaSampleWriter writer{kMaxDivergenceUs};
+    EXPECT_TRUE(writer.init(mTestMuxer, mTestCallbacks));
+
+    // Use two tracks for this test.
+    static constexpr int kNumTracks = 2;
+    std::shared_ptr<MediaSampleQueue> sampleQueues[kNumTracks];
+    const TestMediaSource& mediaSource = getMediaSource();
+
+    for (int trackIdx = 0; trackIdx < kNumTracks; ++trackIdx) {
+        sampleQueues[trackIdx] = std::make_shared<MediaSampleQueue>();
+
+        auto trackFormat = mediaSource.mTrackFormats[trackIdx % mediaSource.mTrackCount];
+        EXPECT_TRUE(writer.addTrack(sampleQueues[trackIdx], trackFormat));
+        EXPECT_EQ(mTestMuxer->popEvent(), TestMuxer::AddTrack(trackFormat.get()));
+    }
+
+    ASSERT_TRUE(writer.start());
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::Start());
+
+    // The first samples of each track can be written in any order since the writer does not have
+    // any previous timestamps to compare.
+    sampleQueues[0]->enqueue(newSampleWithPtsOnly(0));
+    sampleQueues[1]->enqueue(newSampleWithPtsOnly(1));
+    mTestMuxer->popEvent(true);
+    mTestMuxer->popEvent(true);
+
+    // The writer will now be waiting on track 0 since it has the lowest previous timestamp.
+    sampleQueues[0]->enqueue(newSampleWithPtsOnly(kMaxDivergenceUs + 1));
+    sampleQueues[0]->enqueue(newSampleWithPtsOnly(kMaxDivergenceUs + 2));
+
+    // The writer should dequeue the first sample above but not the second since track 0 now is too
+    // far ahead. Instead it should wait for track 1.
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(0, kMaxDivergenceUs + 1));
+
+    // Enqueue a sample from track 1 that puts it within acceptable divergence range again. The
+    // writer should dequeue that sample and then go back to track 0 since track 1 is empty.
+    sampleQueues[1]->enqueue(newSampleWithPtsOnly(kMaxDivergenceUs));
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(1, kMaxDivergenceUs));
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(0, kMaxDivergenceUs + 2));
+
+    // Both tracks are now empty so the writer should wait for track 1 which is farthest behind.
+    sampleQueues[1]->enqueue(newSampleWithPtsOnly(kMaxDivergenceUs + 3));
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(1, kMaxDivergenceUs + 3));
+
+    for (int trackIndex = 0; trackIndex < kNumTracks; ++trackIndex) {
+        sampleQueues[trackIndex]->enqueue(newSampleEos());
+    }
+
+    // Wait for writer to complete.
+    mTestCallbacks->waitForWritingFinished();
+
+    // Verify EOS samples.
+    for (int trackIndex = 0; trackIndex < kNumTracks; ++trackIndex) {
+        auto trackFormat = mediaSource.mTrackFormats[trackIndex % mediaSource.mTrackCount];
+        int64_t duration = 0;
+        AMediaFormat_getInt64(trackFormat.get(), AMEDIAFORMAT_KEY_DURATION, &duration);
+
+        // EOS timestamp = first sample timestamp + duration.
+        const int64_t endTime = duration + (trackIndex == 1 ? 1 : 0);
+        const AMediaCodecBufferInfo info = {0, 0, endTime, AMEDIACODEC_BUFFER_FLAG_END_OF_STREAM};
+        EXPECT_EQ(mTestMuxer->popEvent(), TestMuxer::WriteSample(trackIndex, nullptr, &info));
+    }
+
+    EXPECT_EQ(mTestMuxer->popEvent(), TestMuxer::Stop());
+    EXPECT_TRUE(writer.stop());
+    EXPECT_TRUE(mTestCallbacks->hasFinished());
+}
+
+TEST_F(MediaSampleWriterTests, TestTimestampDivergenceOverflow) {
+    auto testCallbacks = std::make_shared<TestCallbacks>(false /* expectSuccess */);
+    MediaSampleWriter writer{};
+    EXPECT_TRUE(writer.init(mTestMuxer, testCallbacks));
+
+    // Use two tracks for this test.
+    static constexpr int kNumTracks = 2;
+    std::shared_ptr<MediaSampleQueue> sampleQueues[kNumTracks];
+    const TestMediaSource& mediaSource = getMediaSource();
+
+    for (int trackIdx = 0; trackIdx < kNumTracks; ++trackIdx) {
+        sampleQueues[trackIdx] = std::make_shared<MediaSampleQueue>();
+
+        auto trackFormat = mediaSource.mTrackFormats[trackIdx % mediaSource.mTrackCount];
+        EXPECT_TRUE(writer.addTrack(sampleQueues[trackIdx], trackFormat));
+        EXPECT_EQ(mTestMuxer->popEvent(), TestMuxer::AddTrack(trackFormat.get()));
+    }
+
+    // Prime track 0 with lower end of INT64 range, and track 1 with positive timestamps making the
+    // difference larger than INT64_MAX.
+    sampleQueues[0]->enqueue(newSampleWithPtsOnly(INT64_MIN + 1));
+    sampleQueues[1]->enqueue(newSampleWithPtsOnly(1000));
+    sampleQueues[1]->enqueue(newSampleWithPtsOnly(1001));
+
+    ASSERT_TRUE(writer.start());
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::Start());
+
+    // The first sample of each track can be pulled in any order.
+    mTestMuxer->popEvent(true);
+    mTestMuxer->popEvent(true);
+
+    // Wait to make sure the writer compares track 0 empty against track 1 non-empty. The writer
+    // should handle the large timestamp differences and chose to wait for track 0 even though
+    // track 1 has a sample ready.
+    std::this_thread::sleep_for(std::chrono::milliseconds(20));
+
+    sampleQueues[0]->enqueue(newSampleWithPtsOnly(INT64_MIN + 2));
+    sampleQueues[0]->enqueue(newSampleWithPtsOnly(1000));  // <-- Close the gap between the tracks.
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(0, INT64_MIN + 2));
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(0, 1000));
+    EXPECT_EQ(mTestMuxer->popEvent(true), TestMuxer::WriteSampleWithPts(1, 1001));
+
+    EXPECT_TRUE(writer.stop());
+    EXPECT_EQ(mTestMuxer->popEvent(), TestMuxer::Stop());
+    EXPECT_TRUE(testCallbacks->hasFinished());
+}
+
 TEST_F(MediaSampleWriterTests, TestAbortInputQueue) {
     MediaSampleWriter writer{};
     std::shared_ptr<TestCallbacks> callbacks =