Merge "codec2: fix drain flow and EOS work reporting"
diff --git a/C2VDAComponent.cpp b/C2VDAComponent.cpp
index b212e16..c935596 100644
--- a/C2VDAComponent.cpp
+++ b/C2VDAComponent.cpp
@@ -482,6 +482,7 @@
         mThread("C2VDAComponentThread"),
         mVDAInitResult(VideoDecodeAcceleratorAdaptor::Result::ILLEGAL_STATE),
         mComponentState(ComponentState::UNINITIALIZED),
+        mDrainWithEOS(false),
         mColorFormat(0u),
         mLastOutputTimestamp(-1),
         mCodecProfile(media::VIDEO_CODEC_PROFILE_UNKNOWN),
@@ -556,11 +557,17 @@
 
 void C2VDAComponent::onQueueWork(std::unique_ptr<C2Work> work) {
     DCHECK(mTaskRunner->BelongsToCurrentThread());
-    ALOGV("onQueueWork");
+    ALOGV("onQueueWork: flags=0x%x, index=%llu, timestamp=%llu", work->input.flags,
+          work->input.ordinal.frameIndex.peekull(), work->input.ordinal.timestamp.peekull());
     EXPECT_RUNNING_OR_RETURN_ON_ERROR();
     // It is illegal for client to put new works while component is still flushing.
     CHECK_NE(mComponentState, ComponentState::FLUSHING);
-    mQueue.emplace(std::move(work));
+
+    uint32_t drainMode = NO_DRAIN;
+    if (work->input.flags & C2FrameData::FLAG_END_OF_STREAM) {
+        drainMode = DRAIN_COMPONENT_WITH_EOS;
+    }
+    mQueue.push({std::move(work), drainMode});
     // TODO(johnylin): set a maximum size of mQueue and check if mQueue is already full.
 
     mTaskRunner->PostTask(FROM_HERE,
@@ -584,24 +591,32 @@
     }
 
     // Dequeue a work from mQueue.
-    std::unique_ptr<C2Work> work(std::move(mQueue.front()));
+    std::unique_ptr<C2Work> work(std::move(mQueue.front().mWork));
+    auto drainMode = mQueue.front().mDrainMode;
     mQueue.pop();
 
-    // Send input buffer to VDA for decode.
-    // Use frameIndex as bitstreamId.
     CHECK_EQ(work->input.buffers.size(), 1u);
     C2ConstLinearBlock linearBlock = work->input.buffers.front()->data().linearBlocks().front();
+    // linearBlock.size() == 0 means this is a dummy work. No decode needed.
     if (linearBlock.size() > 0) {
+        // Send input buffer to VDA for decode.
+        // Use frameIndex as bitstreamId.
         int32_t bitstreamId = frameIndexToBitstreamId(work->input.ordinal.frameIndex);
         sendInputBufferToAccelerator(linearBlock, bitstreamId);
     }
 
-    if (work->input.flags & C2FrameData::FLAG_END_OF_STREAM) {
+    CHECK_EQ(work->worklets.size(), 1u);
+    work->worklets.front()->output.flags = static_cast<C2FrameData::flags_t>(0);
+    work->worklets.front()->output.buffers.clear();
+    work->worklets.front()->output.ordinal = work->input.ordinal;
+
+    if (drainMode != NO_DRAIN) {
         mVDAAdaptor->flush();
         mComponentState = ComponentState::DRAINING;
+        mDrainWithEOS = drainMode == DRAIN_COMPONENT_WITH_EOS;
     }
 
-    // Put work to mPendingWork.
+    // Put work to mPendingWorks.
     mPendingWorks.emplace_back(std::move(work));
 
     if (!mQueue.empty()) {
@@ -674,13 +689,9 @@
     info->mState = GraphicBlockInfo::State::OWNED_BY_CLIENT;
 
     // Attach output buffer to the work corresponded to bitstreamId.
-    CHECK_EQ(work->worklets.size(), 1u);
-    work->worklets.front()->output.buffers.clear();
     work->worklets.front()->output.buffers.emplace_back(std::make_shared<C2VDAGraphicBuffer>(
             info->mGraphicBlock, base::Bind(&C2VDAComponent::returnOutputBuffer,
                                             mWeakThisFactory.GetWeakPtr(), pictureBufferId)));
-    work->worklets.front()->output.ordinal = work->input.ordinal;
-    work->workletsProcessed = 1u;
 
     // TODO: this does not work for timestamps as they can wrap around
     int64_t currentTimestamp = base::checked_cast<int64_t>(work->input.ordinal.timestamp.peek());
@@ -690,26 +701,27 @@
     reportFinishedWorkIfAny();
 }
 
-void C2VDAComponent::onDrain() {
+void C2VDAComponent::onDrain(uint32_t drainMode) {
     DCHECK(mTaskRunner->BelongsToCurrentThread());
-    ALOGV("onDrain");
-    EXPECT_STATE_OR_RETURN_ON_ERROR(STARTED);
+    ALOGV("onDrain: mode = %u", drainMode);
+    EXPECT_RUNNING_OR_RETURN_ON_ERROR();
 
-    // Set input flag as C2FrameData::FLAG_END_OF_STREAM to the last queued work. If mQueue is
-    // empty, set to the last work in mPendingWorks and then signal flush immediately.
     if (!mQueue.empty()) {
-        mQueue.back()->input.flags = static_cast<C2FrameData::flags_t>(
-                mQueue.back()->input.flags | C2FrameData::FLAG_END_OF_STREAM);
-    } else if (!mPendingWorks.empty()) {
-        C2Work* work = getPendingWorkLastToFinish();
-        if (!work) {
-            reportError(C2_CORRUPTED);
-            return;
+        // Mark last queued work as "drain-till-here" by setting drainMode. Do not change drainMode
+        // if last work already has one.
+        if (mQueue.back().mDrainMode == NO_DRAIN) {
+            mQueue.back().mDrainMode = drainMode;
         }
-        mPendingWorks.back()->input.flags = static_cast<C2FrameData::flags_t>(
-                mPendingWorks.back()->input.flags | C2FrameData::FLAG_END_OF_STREAM);
-        mVDAAdaptor->flush();
-        mComponentState = ComponentState::DRAINING;
+    } else if (!mPendingWorks.empty()) {
+        // Neglect drain request if component is not in STARTED mode. Otherwise, enters DRAINING
+        // mode and signal VDA flush immediately.
+        if (mComponentState == ComponentState::STARTED) {
+            mVDAAdaptor->flush();
+            mComponentState = ComponentState::DRAINING;
+            mDrainWithEOS = drainMode == DRAIN_COMPONENT_WITH_EOS;
+        } else {
+            ALOGV("Neglect drain. Component in state: %d", mComponentState);
+        }
     } else {
         // Do nothing.
         ALOGV("No buffers in VDA, drain takes no effect.");
@@ -730,6 +742,13 @@
         return;
     }
 
+    if (mDrainWithEOS) {
+        // Return EOS work.
+        reportEOSWork();
+    }
+    // mPendingWorks must be empty after draining is finished.
+    CHECK(mPendingWorks.empty());
+
     // Last stream is finished. Reset the timestamp record.
     mLastOutputTimestamp = -1;
 
@@ -746,7 +765,7 @@
     mVDAAdaptor->reset();
     // Pop all works in mQueue and put into mPendingWorks.
     while (!mQueue.empty()) {
-        mPendingWorks.emplace_back(std::move(mQueue.front()));
+        mPendingWorks.emplace_back(std::move(mQueue.front().mWork));
         mQueue.pop();
     }
     mComponentState = ComponentState::FLUSHING;
@@ -760,7 +779,7 @@
     mVDAAdaptor->reset();
     // Pop all works in mQueue and put into mPendingWorks.
     while (!mQueue.empty()) {
-        mPendingWorks.emplace_back(std::move(mQueue.front()));
+        mPendingWorks.emplace_back(std::move(mQueue.front().mWork));
         mQueue.pop();
     }
 
@@ -856,20 +875,6 @@
     return workIter->get();
 }
 
-C2Work* C2VDAComponent::getPendingWorkLastToFinish() {
-    // Get the work with largest timestamp.
-    auto workIter = std::max_element(
-            mPendingWorks.begin(), mPendingWorks.end(), [](const auto& w1, const auto& w2) {
-                return w1->input.ordinal.timestamp < w2->input.ordinal.timestamp;
-            });
-
-    if (workIter == mPendingWorks.end()) {
-        ALOGE("Can't get last finished work from mPendingWork");
-        return nullptr;
-    }
-    return workIter->get();
-}
-
 C2VDAComponent::GraphicBlockInfo* C2VDAComponent::getGraphicBlockById(int32_t blockId) {
     if (blockId < 0 || blockId >= static_cast<int32_t>(mGraphicBlocks.size())) {
         ALOGE("getGraphicBlockById failed: id=%d", blockId);
@@ -1109,13 +1114,14 @@
 }
 
 c2_status_t C2VDAComponent::drain_nb(drain_mode_t mode) {
-    if (mode != DRAIN_COMPONENT_WITH_EOS) {
+    if (mode != DRAIN_COMPONENT_WITH_EOS && mode != DRAIN_COMPONENT_NO_EOS) {
         return C2_OMITTED;  // Tunneling is not supported by now
     }
     if (mState != State::RUNNING) {
         return C2_BAD_STATE;
     }
-    mTaskRunner->PostTask(FROM_HERE, base::Bind(&C2VDAComponent::onDrain, base::Unretained(this)));
+    mTaskRunner->PostTask(FROM_HERE, base::Bind(&C2VDAComponent::onDrain, base::Unretained(this),
+                                                static_cast<uint32_t>(mode)));
     return C2_OK;
 }
 
@@ -1250,10 +1256,12 @@
     // However, the timestamp is guaranteed to be monotonic increasing for buffers in display order.
     // That is, since VDA output is in display order, if we get a returned output with timestamp T,
     // it implies all works with timestamp <= T are done.
+    // EOS work will not be reported here. reportEOSWork() does it.
     auto iter = mPendingWorks.begin();
     while (iter != mPendingWorks.end()) {
         if (isWorkDone(iter->get())) {
             iter->get()->result = C2_OK;
+            iter->get()->workletsProcessed = static_cast<uint32_t>(iter->get()->worklets.size());
             finishedWorks.emplace_back(std::move(*iter));
             iter = mPendingWorks.erase(iter);
         } else {
@@ -1268,7 +1276,16 @@
 
 bool C2VDAComponent::isWorkDone(const C2Work* work) const {
     if (!work->input.buffers.empty()) {
-        return false;  // Input buffer is still owned by VDA.
+        // Input buffer is still owned by VDA.
+        // This condition could also recognize dummy EOS work since it won't get
+        // onInputBufferDone(), input.buffers won't be cleared until reportEOSWork().
+        return false;
+    }
+    if (mComponentState == ComponentState::DRAINING && mDrainWithEOS &&
+        mPendingWorks.size() == 1u) {
+        // If component is in DRAINING state and mDrainWithEOS is true. The last returned work
+        // should be marked EOS flag and returned by reportEOSWork() instead.
+        return false;
     }
     if (mLastOutputTimestamp < 0) {
         return false;  // No output buffer is returned yet.
@@ -1279,6 +1296,28 @@
     return true;  // Output buffer is returned, or it has no related output buffer.
 }
 
+void C2VDAComponent::reportEOSWork() {
+    ALOGV("reportEOSWork");
+    DCHECK(mTaskRunner->BelongsToCurrentThread());
+    // In this moment all works prior to EOS work should be done and returned to listener.
+    if (mPendingWorks.size() != 1u) {  // only EOS work left
+        ALOGE("It shouldn't have remaining works in mPendingWorks except EOS work.");
+        reportError(C2_CORRUPTED);
+        return;
+    }
+
+    std::unique_ptr<C2Work> eosWork(std::move(mPendingWorks.front()));
+    mPendingWorks.pop_front();
+    eosWork->input.buffers.clear();
+    eosWork->result = C2_OK;
+    eosWork->workletsProcessed = static_cast<uint32_t>(eosWork->worklets.size());
+    eosWork->worklets.front()->output.flags = C2FrameData::FLAG_END_OF_STREAM;
+
+    std::list<std::unique_ptr<C2Work>> finishedWorks;
+    finishedWorks.emplace_back(std::move(eosWork));
+    mListener->onWorkDone_nb(shared_from_this(), std::move(finishedWorks));
+}
+
 void C2VDAComponent::reportAbandonedWorks() {
     DCHECK(mTaskRunner->BelongsToCurrentThread());
     std::list<std::unique_ptr<C2Work>> abandonedWorks;
diff --git a/cmds/codec2.cpp b/cmds/codec2.cpp
index d5d1250..155c5aa 100644
--- a/cmds/codec2.cpp
+++ b/cmds/codec2.cpp
@@ -271,7 +271,8 @@
                 mProcessedWork.pop_front();
             }
 
-            if (work->workletsProcessed > 0) {
+            CHECK_EQ(work->worklets.size(), 1u);
+            if (work->worklets.front()->output.buffers.size() == 1u) {
                 int slot;
                 sp<Fence> fence;
                 std::shared_ptr<C2Buffer> output = work->worklets.front()->output.buffers[0];
@@ -313,12 +314,13 @@
 #endif
             }
 
+            bool eos = work->worklets.front()->output.flags & C2FrameData::FLAG_END_OF_STREAM;
             // input buffers should be cleared in component side.
             CHECK(work->input.buffers.empty());
             work->worklets.clear();
             work->workletsProcessed = 0;
 
-            if (work->input.flags & C2FrameData::FLAG_END_OF_STREAM) {
+            if (eos) {
                 running.store(false);  // stop the thread
             }
 
diff --git a/include/C2VDAComponent.h b/include/C2VDAComponent.h
index 9744259..c7c6eb3 100644
--- a/include/C2VDAComponent.h
+++ b/include/C2VDAComponent.h
@@ -179,6 +179,17 @@
         kDpbOutputBufferExtraCount = 3,  // Use the same number as ACodec.
     };
 
+    // This constant is used to tell apart from drain_mode_t enumerations in C2Component.h, which
+    // means no drain request.
+    // Note: this value must be different than all enumerations in drain_mode_t.
+    static constexpr uint32_t NO_DRAIN = ~0u;
+
+    // Internal struct for work queue.
+    struct WorkEntry {
+        std::unique_ptr<C2Work> mWork;
+        uint32_t mDrainMode = NO_DRAIN;
+    };
+
     // Internal struct to keep the information of a specific graphic block.
     struct GraphicBlockInfo {
         enum class State {
@@ -224,7 +235,7 @@
     void onDequeueWork();
     void onInputBufferDone(int32_t bitstreamId);
     void onOutputBufferDone(int32_t pictureBufferId, int32_t bitstreamId);
-    void onDrain();
+    void onDrain(uint32_t drainMode);
     void onDrainDone();
     void onFlush();
     void onStop(base::WaitableEvent* done);
@@ -245,8 +256,6 @@
     GraphicBlockInfo* getGraphicBlockById(int32_t blockId);
     // Helper function to get the specified work in mPendingWorks by bitstream id.
     C2Work* getPendingWorkByBitstreamId(int32_t bitstreamId);
-    // Helper function to get the work which is last to finish in mPendingWorks.
-    C2Work* getPendingWorkLastToFinish();
     // Try to apply the output format change.
     void tryChangeOutputFormat();
     // Allocate output buffers (graphic blocks) from block allocator.
@@ -256,6 +265,8 @@
 
     // Check for finished works in mPendingWorks. If any, make onWorkDone call to listener.
     void reportFinishedWorkIfAny();
+    // Make onWorkDone call to listener for reporting EOS work in mPendingWorks.
+    void reportEOSWork();
     // Abandon all works in mPendingWorks.
     void reportAbandonedWorks();
     // Make onError call to listener for reporting errors.
@@ -284,11 +295,14 @@
     base::WaitableEvent* mStopDoneEvent;
     // The state machine on component thread.
     ComponentState mComponentState;
+    // The indicator of drain mode (true for draining with EOS). This should be always set along
+    // with component going to DRAINING state, and only regarded under DRAINING state.
+    bool mDrainWithEOS;
     // The vector of storing allocated output graphic block information.
     std::vector<GraphicBlockInfo> mGraphicBlocks;
-    // The work queue. Works are queued from component API queue_nb and dequeued by the decode
-    // process of component.
-    std::queue<std::unique_ptr<C2Work>> mQueue;
+    // The work queue. Works are queued along with drain mode from component API queue_nb and
+    // dequeued by the decode process of component.
+    std::queue<WorkEntry> mQueue;
     // Store all pending works. The dequeued works are placed here until they are finished and then
     // sent out by onWorkDone call to listener.
     std::deque<std::unique_ptr<C2Work>> mPendingWorks;
diff --git a/tests/C2VDAComponent_test.cpp b/tests/C2VDAComponent_test.cpp
index 29041f0..2e345b3 100644
--- a/tests/C2VDAComponent_test.cpp
+++ b/tests/C2VDAComponent_test.cpp
@@ -139,6 +139,12 @@
           : C2Buffer({block->share(block->offset(), block->size(), C2Fence())}) {}
 };
 
+class C2VDADummyLinearBuffer : public C2Buffer {
+public:
+    explicit C2VDADummyLinearBuffer(const std::shared_ptr<C2LinearBlock>& block)
+          : C2Buffer({block->share(0, 0, C2Fence())}) {}
+};
+
 class Listener;
 
 class C2VDAComponentTest : public ::testing::Test {
@@ -396,13 +402,16 @@
 // - Sanity check. If this is true, decoded content sanity check is enabled. Test will compute the
 //   MD5Sum for output frame data for a play-though iteration (not flushed), and compare to golden
 //   MD5Sums which should be stored in the file |video_filename|.md5
+// - Use dummy EOS work. If this is true, test will queue a dummy work with end-of-stream flag in
+//   the end of all input works. On the contrary, test will call drain_nb() to component.
 class C2VDAComponentParamTest
       : public C2VDAComponentTest,
-        public ::testing::WithParamInterface<std::tuple<int, uint32_t, bool>> {
+        public ::testing::WithParamInterface<std::tuple<int, uint32_t, bool, bool>> {
 protected:
     int mFlushAfterWorkIndex;
     uint32_t mNumberOfPlaythrough;
     bool mSanityCheck;
+    bool mUseDummyEOSWork;
 };
 
 TEST_P(C2VDAComponentParamTest, SimpleDecodeTest) {
@@ -420,14 +429,18 @@
     }
 
     mSanityCheck = std::get<2>(GetParam());
+    mUseDummyEOSWork = std::get<3>(GetParam());
 
     // Reset counters and determine the expected answers for all iterations.
     mOutputFrameCounts.resize(mNumberOfPlaythrough, 0);
     mFinishedWorkCounts.resize(mNumberOfPlaythrough, 0);
     mMD5Strings.resize(mNumberOfPlaythrough);
     std::vector<int> expectedOutputFrameCounts(mNumberOfPlaythrough, mTestVideoFile->mNumFrames);
-    std::vector<int> expectedFinishedWorkCounts(mNumberOfPlaythrough,
-                                                mTestVideoFile->mNumFragments);
+    auto expectedWorkCount = mTestVideoFile->mNumFragments;
+    if (mUseDummyEOSWork) {
+        expectedWorkCount += 1;  // plus one dummy EOS work
+    }
+    std::vector<int> expectedFinishedWorkCounts(mNumberOfPlaythrough, expectedWorkCount);
     if (mFlushAfterWorkIndex >= 0) {
         // First iteration performs the mid-stream flushing.
         expectedOutputFrameCounts[0] = mFlushAfterWorkIndex + 1;
@@ -473,13 +486,13 @@
                 mProcessedWork.pop_front();
             }
             mFinishedWorkCounts[iteration]++;
-            ALOGV("Output: frame index: %llu result: %d outputs: %zu",
+            ALOGV("Output: frame index: %llu result: %d flags: 0x%x buffers: %zu",
                   work->input.ordinal.frameIndex.peekull(), work->result,
+                  work->worklets.front()->output.flags,
                   work->worklets.front()->output.buffers.size());
 
-            if (work->workletsProcessed == 1u) {
-                ASSERT_EQ(work->worklets.size(), 1u);
-                ASSERT_EQ(work->worklets.front()->output.buffers.size(), 1u);
+            ASSERT_EQ(work->worklets.size(), 1u);
+            if (work->worklets.front()->output.buffers.size() == 1u) {
                 std::shared_ptr<C2Buffer> output = work->worklets.front()->output.buffers[0];
                 C2ConstGraphicBlock graphicBlock = output->data().graphicBlocks().front();
                 ASSERT_EQ(mTestVideoFile->mWidth, static_cast<int>(graphicBlock.width()));
@@ -507,12 +520,14 @@
                 mOutputFrameCounts[iteration]++;
             }
 
+            bool iteration_end =
+                    work->worklets.front()->output.flags & C2FrameData::FLAG_END_OF_STREAM;
+
             // input buffers should be cleared in component side.
             ASSERT_TRUE(work->input.buffers.empty());
             work->worklets.clear();
             work->workletsProcessed = 0;
 
-            bool iteration_end = work->input.flags & C2FrameData::FLAG_END_OF_STREAM;
             if (iteration == 0 && work->input.ordinal.frameIndex.peeku() ==
                                           static_cast<uint64_t>(mFlushAfterWorkIndex)) {
                 ULock l(mFlushDoneLock);
@@ -564,6 +579,7 @@
             int64_t timestamp = 0u;
             MediaBuffer* buffer = nullptr;
             sp<ABuffer> csd;
+            bool queueDummyEOSWork = false;
             if (!csds.empty()) {
                 csd = std::move(csds.front());
                 csds.pop_front();
@@ -572,14 +588,23 @@
             } else {
                 if (mTestVideoFile->mData->read(&buffer) != OK) {
                     ASSERT_TRUE(buffer == nullptr);
-                    ALOGV("Meet end of stream. Now drain the component.");
-                    ASSERT_EQ(component->drain_nb(C2Component::DRAIN_COMPONENT_WITH_EOS), C2_OK);
-                    break;
+                    if (mUseDummyEOSWork) {
+                        ALOGV("Meet end of stream. Put a dummy EOS work.");
+                        queueDummyEOSWork = true;
+                    } else {
+                        ALOGV("Meet end of stream. Now drain the component.");
+                        ASSERT_EQ(component->drain_nb(C2Component::DRAIN_COMPONENT_WITH_EOS),
+                                  C2_OK);
+                        break;
+                    }
+                    // TODO(johnylin): add test with drain with DRAIN_COMPONENT_NO_EOS when we know
+                    //                 the actual use case of it.
+                } else {
+                    sp<MetaData> meta = buffer->meta_data();
+                    ASSERT_TRUE(meta->findInt64(kKeyTime, &timestamp));
+                    size = buffer->size();
+                    data = buffer->data();
                 }
-                sp<MetaData> meta = buffer->meta_data();
-                ASSERT_TRUE(meta->findInt64(kKeyTime, &timestamp));
-                size = buffer->size();
-                data = buffer->data();
             }
 
             std::unique_ptr<C2Work> work;
@@ -592,25 +617,38 @@
                     mQueueCondition.wait_for(l, 100ms);
                 }
             }
-            work->input.flags = static_cast<C2FrameData::flags_t>(0);
-            work->input.ordinal.timestamp = static_cast<uint64_t>(timestamp);
+
             work->input.ordinal.frameIndex = static_cast<uint64_t>(numWorks);
-
-            // Allocate input buffer.
-            std::shared_ptr<C2LinearBlock> block;
-            mLinearBlockPool->fetchLinearBlock(
-                    size, {C2MemoryUsage::CPU_READ, C2MemoryUsage::CPU_WRITE}, &block);
-            C2WriteView view = block->map().get();
-            ASSERT_EQ(view.error(), C2_OK);
-            memcpy(view.base(), data, size);
-
             work->input.buffers.clear();
-            work->input.buffers.emplace_back(new C2VDALinearBuffer(std::move(block)));
+
+            std::shared_ptr<C2LinearBlock> block;
+            if (queueDummyEOSWork) {
+                work->input.flags = C2FrameData::FLAG_END_OF_STREAM;
+                work->input.ordinal.timestamp = 0;  // timestamp is invalid for dummy EOS work
+
+                // Create a dummy input buffer by allocating minimal size of buffer from block pool.
+                mLinearBlockPool->fetchLinearBlock(
+                        1, {C2MemoryUsage::CPU_READ, C2MemoryUsage::CPU_WRITE}, &block);
+                work->input.buffers.emplace_back(new C2VDADummyLinearBuffer(std::move(block)));
+                ALOGV("Input: (Dummy EOS) id: %llu", work->input.ordinal.frameIndex.peekull());
+            } else {
+                work->input.flags = static_cast<C2FrameData::flags_t>(0);
+                work->input.ordinal.timestamp = static_cast<uint64_t>(timestamp);
+
+                // Allocate an input buffer with data size.
+                mLinearBlockPool->fetchLinearBlock(
+                        size, {C2MemoryUsage::CPU_READ, C2MemoryUsage::CPU_WRITE}, &block);
+                C2WriteView view = block->map().get();
+                ASSERT_EQ(view.error(), C2_OK);
+                memcpy(view.base(), data, size);
+                work->input.buffers.emplace_back(new C2VDALinearBuffer(std::move(block)));
+                ALOGV("Input: bitstream id: %llu timestamp: %llu size: %zu",
+                      work->input.ordinal.frameIndex.peekull(),
+                      work->input.ordinal.timestamp.peekull(), size);
+            }
+
             work->worklets.clear();
             work->worklets.emplace_back(new C2Worklet);
-            ALOGV("Input: bitstream id: %llu timestamp: %llu size: %zu",
-                  work->input.ordinal.frameIndex.peekull(), work->input.ordinal.timestamp.peekull(),
-                  size);
 
             std::list<std::unique_ptr<C2Work>> items;
             items.push_back(std::move(work));
@@ -631,6 +669,10 @@
                           C2_OK);
                 break;
             }
+
+            if (queueDummyEOSWork) {
+                break;
+            }
         }
 
         if (iteration == 0 && mFlushAfterWorkIndex >= 0) {
@@ -645,7 +687,7 @@
             ALOGV("Got flush done signal");
             EXPECT_EQ(numWorks, mFlushAfterWorkIndex + 1);
         } else {
-            EXPECT_EQ(numWorks, mTestVideoFile->mNumFragments);
+            EXPECT_EQ(numWorks, expectedWorkCount);
         }
         ASSERT_EQ(mTestVideoFile->mData->stop(), OK);
     }
@@ -678,36 +720,41 @@
     }
 }
 
-// Play input video once.
+// Play input video once, end by draining.
 INSTANTIATE_TEST_CASE_P(SinglePlaythroughTest, C2VDAComponentParamTest,
                         ::testing::Values(std::make_tuple(static_cast<int>(FlushPoint::NO_FLUSH),
-                                                          1u, false)));
+                                                          1u, false, false)));
+// Play input video once, end by dummy EOS work.
+INSTANTIATE_TEST_CASE_P(DummyEOSWorkTest, C2VDAComponentParamTest,
+                        ::testing::Values(std::make_tuple(static_cast<int>(FlushPoint::NO_FLUSH),
+                                                          1u, false, true)));
 
 // Play 5 times of input video, and check sanity by MD5Sum.
 INSTANTIATE_TEST_CASE_P(MultiplePlaythroughSanityTest, C2VDAComponentParamTest,
                         ::testing::Values(std::make_tuple(static_cast<int>(FlushPoint::NO_FLUSH),
-                                                          5u, true)));
+                                                          5u, true, false)));
 
 // Test mid-stream flush then play once entirely.
 INSTANTIATE_TEST_CASE_P(FlushPlaythroughTest, C2VDAComponentParamTest,
-                        ::testing::Values(std::make_tuple(40, 1u, true)));
+                        ::testing::Values(std::make_tuple(40, 1u, true, false)));
 
 // Test mid-stream flush then stop.
 INSTANTIATE_TEST_CASE_P(FlushStopTest, C2VDAComponentParamTest,
                         ::testing::Values(std::make_tuple(
-                                static_cast<int>(FlushPoint::MID_STREAM_FLUSH), 0u, false)));
+                                static_cast<int>(FlushPoint::MID_STREAM_FLUSH), 0u, false, false)));
 
 // Test early flush (after a few works) then stop.
 INSTANTIATE_TEST_CASE_P(EarlyFlushStopTest, C2VDAComponentParamTest,
-                        ::testing::Values(std::make_tuple(0, 0u, false),
-                                          std::make_tuple(1, 0u, false),
-                                          std::make_tuple(2, 0u, false),
-                                          std::make_tuple(3, 0u, false)));
+                        ::testing::Values(std::make_tuple(0, 0u, false, false),
+                                          std::make_tuple(1, 0u, false, false),
+                                          std::make_tuple(2, 0u, false, false),
+                                          std::make_tuple(3, 0u, false, false)));
 
 // Test end-of-stream flush then stop.
-INSTANTIATE_TEST_CASE_P(EndOfStreamFlushStopTest, C2VDAComponentParamTest,
-                        ::testing::Values(std::make_tuple(
-                                static_cast<int>(FlushPoint::END_OF_STREAM_FLUSH), 0u, false)));
+INSTANTIATE_TEST_CASE_P(
+        EndOfStreamFlushStopTest, C2VDAComponentParamTest,
+        ::testing::Values(std::make_tuple(static_cast<int>(FlushPoint::END_OF_STREAM_FLUSH), 0u,
+                                          false, false)));
 
 }  // namespace android