Migrate buffers during surface change

Migrate graphic buffers during surface change in order to avoid
BufferQueue handling complexity later on.

Test: Manually using Chrome and google photo app
Bug: 132302078
Bug: 130862880
Change-Id: Ifb348b5d6a8f5a89dcc10a9f0be075057a5d3a6d
diff --git a/media/codec2/hidl/1.0/utils/ClientBlockHelper.cpp b/media/codec2/hidl/1.0/utils/ClientBlockHelper.cpp
index 1786c69..50790bc 100644
--- a/media/codec2/hidl/1.0/utils/ClientBlockHelper.cpp
+++ b/media/codec2/hidl/1.0/utils/ClientBlockHelper.cpp
@@ -82,21 +82,14 @@
 
 template <typename BlockProcessor>
 void forEachBlock(const std::list<std::unique_ptr<C2Work>>& workList,
-                  BlockProcessor process,
-                  bool processInput, bool processOutput) {
+                  BlockProcessor process) {
     for (const std::unique_ptr<C2Work>& work : workList) {
         if (!work) {
             continue;
         }
-        if (processInput) {
-            forEachBlock(work->input, process);
-        }
-        if (processOutput) {
-            for (const std::unique_ptr<C2Worklet>& worklet : work->worklets) {
-                if (worklet) {
-                    forEachBlock(worklet->output,
-                                 process);
-                }
+        for (const std::unique_ptr<C2Worklet>& worklet : work->worklets) {
+            if (worklet) {
+                forEachBlock(worklet->output, process);
             }
         }
     }
@@ -109,8 +102,6 @@
             new B2HGraphicBufferProducer(igbp);
 }
 
-} // unnamed namespace
-
 status_t attachToBufferQueue(const C2ConstGraphicBlock& block,
                              const sp<IGraphicBufferProducer>& igbp,
                              uint32_t generation,
@@ -154,73 +145,221 @@
             _C2BlockFactory::GetGraphicBlockPoolData(block),
             generation, bqId, bqSlot);
 }
+} // unnamed namespace
 
-bool holdBufferQueueBlock(const C2ConstGraphicBlock& block,
-                            const sp<IGraphicBufferProducer>& igbp,
-                            uint64_t bqId,
-                            uint32_t generation) {
-    std::shared_ptr<_C2BlockPoolData> data =
-            _C2BlockFactory::GetGraphicBlockPoolData(block);
-    if (!data) {
-        return false;
-    }
+class OutputBufferQueue::Impl {
+    std::mutex mMutex;
+    sp<IGraphicBufferProducer> mIgbp;
+    uint32_t mGeneration;
+    uint64_t mBqId;
+    std::shared_ptr<int> mOwner;
+    // To migrate existing buffers
+    sp<GraphicBuffer> mBuffers[BufferQueueDefs::NUM_BUFFER_SLOTS]; // find a better way
+    std::weak_ptr<_C2BlockPoolData>
+                    mPoolDatas[BufferQueueDefs::NUM_BUFFER_SLOTS];
 
-    uint32_t oldGeneration;
-    uint64_t oldId;
-    int32_t oldSlot;
-    // If the block is not bufferqueue-based, do nothing.
-    if (!_C2BlockFactory::GetBufferQueueData(
-            data, &oldGeneration, &oldId, &oldSlot) ||
-            (oldId == 0)) {
-        return false;
-    }
+public:
+    Impl(): mGeneration(0), mBqId(0) {}
 
-    // If the block's bqId is the same as the desired bqId, just hold.
-    if ((oldId == bqId) && (oldGeneration == generation)) {
-        LOG(VERBOSE) << "holdBufferQueueBlock -- import without attaching:"
-                     << " bqId " << oldId
-                     << ", bqSlot " << oldSlot
-                     << ", generation " << generation
-                     << ".";
-        _C2BlockFactory::HoldBlockFromBufferQueue(data, getHgbp(igbp));
+    bool configure(const sp<IGraphicBufferProducer>& igbp,
+                   uint32_t generation,
+                   uint64_t bqId) {
+        size_t tryNum = 0;
+        size_t success = 0;
+        sp<GraphicBuffer> buffers[BufferQueueDefs::NUM_BUFFER_SLOTS];
+        std::weak_ptr<_C2BlockPoolData>
+                poolDatas[BufferQueueDefs::NUM_BUFFER_SLOTS];
+        {
+            std::scoped_lock<std::mutex> l(mMutex);
+            if (generation == mGeneration) {
+                return false;
+            }
+            mIgbp = igbp;
+            mGeneration = generation;
+            mBqId = bqId;
+            mOwner = std::make_shared<int>(0);
+            for (int i = 0; i < BufferQueueDefs::NUM_BUFFER_SLOTS; ++i) {
+                if (mBqId == 0 || !mBuffers[i]) {
+                    continue;
+                }
+                std::shared_ptr<_C2BlockPoolData> data = mPoolDatas[i].lock();
+                if (!data ||
+                    !_C2BlockFactory::BeginAttachBlockToBufferQueue(data)) {
+                    continue;
+                }
+                ++tryNum;
+                int bqSlot;
+                mBuffers[i]->setGenerationNumber(generation);
+                status_t result = igbp->attachBuffer(&bqSlot, mBuffers[i]);
+                if (result != OK) {
+                    continue;
+                }
+                bool attach =
+                        _C2BlockFactory::EndAttachBlockToBufferQueue(
+                                data, mOwner, getHgbp(mIgbp),
+                                generation, bqId, bqSlot);
+                if (!attach) {
+                    igbp->cancelBuffer(bqSlot, Fence::NO_FENCE);
+                    continue;
+                }
+                buffers[bqSlot] = mBuffers[i];
+                poolDatas[bqSlot] = data;
+                ++success;
+            }
+            for (int i = 0; i < BufferQueueDefs::NUM_BUFFER_SLOTS; ++i) {
+                mBuffers[i] = buffers[i];
+                mPoolDatas[i] = poolDatas[i];
+            }
+        }
+        ALOGD("remote graphic buffer migration %zu/%zu", success, tryNum);
         return true;
     }
 
-    // Otherwise, attach to the given igbp, which must not be null.
-    if (!igbp) {
+    bool registerBuffer(const C2ConstGraphicBlock& block) {
+        std::shared_ptr<_C2BlockPoolData> data =
+                _C2BlockFactory::GetGraphicBlockPoolData(block);
+        if (!data) {
+            return false;
+        }
+        std::scoped_lock<std::mutex> l(mMutex);
+
+        if (!mIgbp) {
+            return false;
+        }
+
+        uint32_t oldGeneration;
+        uint64_t oldId;
+        int32_t oldSlot;
+        // If the block is not bufferqueue-based, do nothing.
+        if (!_C2BlockFactory::GetBufferQueueData(
+                data, &oldGeneration, &oldId, &oldSlot) || (oldId == 0)) {
+            return false;
+        }
+        // If the block's bqId is the same as the desired bqId, just hold.
+        if ((oldId == mBqId) && (oldGeneration == mGeneration)) {
+            LOG(VERBOSE) << "holdBufferQueueBlock -- import without attaching:"
+                         << " bqId " << oldId
+                         << ", bqSlot " << oldSlot
+                         << ", generation " << mGeneration
+                         << ".";
+            _C2BlockFactory::HoldBlockFromBufferQueue(data, mOwner, getHgbp(mIgbp));
+            mPoolDatas[oldSlot] = data;
+            mBuffers[oldSlot] = createGraphicBuffer(block);
+            mBuffers[oldSlot]->setGenerationNumber(mGeneration);
+            return true;
+        }
+        int32_t d = (int32_t) mGeneration - (int32_t) oldGeneration;
+        LOG(WARNING) << "receiving stale buffer: generation "
+                     << mGeneration << " , diff " << d  << " : slot "
+                     << oldSlot;
         return false;
     }
 
-    int32_t bqSlot;
-    status_t result = attachToBufferQueue(block, igbp, generation, &bqSlot);
+    status_t outputBuffer(
+            const C2ConstGraphicBlock& block,
+            const BnGraphicBufferProducer::QueueBufferInput& input,
+            BnGraphicBufferProducer::QueueBufferOutput* output) {
+        uint32_t generation;
+        uint64_t bqId;
+        int32_t bqSlot;
+        bool display = displayBufferQueueBlock(block);
+        if (!getBufferQueueAssignment(block, &generation, &bqId, &bqSlot) ||
+            bqId == 0) {
+            // Block not from bufferqueue -- it must be attached before queuing.
 
-    if (result != OK) {
-        LOG(ERROR) << "holdBufferQueueBlock -- fail to attach:"
-                   << " target bqId " << bqId
-                   << ", generation " << generation
-                   << ".";
-        return false;
+            mMutex.lock();
+            sp<IGraphicBufferProducer> outputIgbp = mIgbp;
+            uint32_t outputGeneration = mGeneration;
+            mMutex.unlock();
+
+            status_t status = attachToBufferQueue(
+                    block, outputIgbp, outputGeneration, &bqSlot);
+            if (status != OK) {
+                LOG(WARNING) << "outputBuffer -- attaching failed.";
+                return INVALID_OPERATION;
+            }
+
+            status = outputIgbp->queueBuffer(static_cast<int>(bqSlot),
+                                         input, output);
+            if (status != OK) {
+                LOG(ERROR) << "outputBuffer -- queueBuffer() failed "
+                           "on non-bufferqueue-based block. "
+                           "Error = " << status << ".";
+                return status;
+            }
+            return OK;
+        }
+
+        mMutex.lock();
+        sp<IGraphicBufferProducer> outputIgbp = mIgbp;
+        uint32_t outputGeneration = mGeneration;
+        uint64_t outputBqId = mBqId;
+        mMutex.unlock();
+
+        if (!outputIgbp) {
+            LOG(VERBOSE) << "outputBuffer -- output surface is null.";
+            return NO_INIT;
+        }
+
+        if (!display) {
+            LOG(WARNING) << "outputBuffer -- cannot display "
+                         "bufferqueue-based block to the bufferqueue.";
+            return UNKNOWN_ERROR;
+        }
+        if (bqId != outputBqId || generation != outputGeneration) {
+            int32_t diff = (int32_t) outputGeneration - (int32_t) generation;
+            LOG(WARNING) << "outputBuffer -- buffers from old generation to "
+                         << outputGeneration << " , diff: " << diff
+                         << " , slot: " << bqSlot;
+            return DEAD_OBJECT;
+        }
+
+        status_t status = outputIgbp->queueBuffer(static_cast<int>(bqSlot),
+                                              input, output);
+        if (status != OK) {
+            LOG(ERROR) << "outputBuffer -- queueBuffer() failed "
+                       "on bufferqueue-based block. "
+                       "Error = " << status << ".";
+            return status;
+        }
+        return OK;
     }
 
-    LOG(VERBOSE) << "holdBufferQueueBlock -- attached:"
-                 << " bqId " << bqId
-                 << ", bqSlot " << bqSlot
-                 << ", generation " << generation
-                 << ".";
-    _C2BlockFactory::AssignBlockToBufferQueue(
-            data, getHgbp(igbp), generation, bqId, bqSlot, true);
-    return true;
+    Impl *getPtr() {
+        return this;
+    }
+
+    ~Impl() {}
+};
+
+OutputBufferQueue::OutputBufferQueue(): mImpl(new Impl()) {}
+
+OutputBufferQueue::~OutputBufferQueue() {}
+
+bool OutputBufferQueue::configure(const sp<IGraphicBufferProducer>& igbp,
+                                  uint32_t generation,
+                                  uint64_t bqId) {
+    return mImpl && mImpl->configure(igbp, generation, bqId);
 }
 
-void holdBufferQueueBlocks(const std::list<std::unique_ptr<C2Work>>& workList,
-                           const sp<IGraphicBufferProducer>& igbp,
-                           uint64_t bqId,
-                           uint32_t generation,
-                           bool forInput) {
+status_t OutputBufferQueue::outputBuffer(
+    const C2ConstGraphicBlock& block,
+    const BnGraphicBufferProducer::QueueBufferInput& input,
+    BnGraphicBufferProducer::QueueBufferOutput* output) {
+    if (mImpl) {
+        return mImpl->outputBuffer(block, input, output);
+    }
+    return DEAD_OBJECT;
+}
+
+void OutputBufferQueue::holdBufferQueueBlocks(
+        const std::list<std::unique_ptr<C2Work>>& workList) {
+    if (!mImpl) {
+        return;
+    }
     forEachBlock(workList,
-                 std::bind(holdBufferQueueBlock,
-                           std::placeholders::_1, igbp, bqId, generation),
-                 forInput, !forInput);
+                 std::bind(&OutputBufferQueue::Impl::registerBuffer,
+                           mImpl->getPtr(), std::placeholders::_1));
 }
 
 }  // namespace utils
diff --git a/media/codec2/hidl/1.0/utils/Component.cpp b/media/codec2/hidl/1.0/utils/Component.cpp
index 5897dce..a9f20a4 100644
--- a/media/codec2/hidl/1.0/utils/Component.cpp
+++ b/media/codec2/hidl/1.0/utils/Component.cpp
@@ -107,19 +107,22 @@
             WorkBundle workBundle;
 
             sp<Component> strongComponent = mComponent.promote();
+            beginTransferBufferQueueBlocks(c2workItems, true);
             if (!objcpy(&workBundle, c2workItems, strongComponent ?
                     &strongComponent->mBufferPoolSender : nullptr)) {
                 LOG(ERROR) << "Component::Listener::onWorkDone_nb -- "
                            << "received corrupted work items.";
+                endTransferBufferQueueBlocks(c2workItems, false, true);
                 return;
             }
             Return<void> transStatus = listener->onWorkDone(workBundle);
             if (!transStatus.isOk()) {
                 LOG(ERROR) << "Component::Listener::onWorkDone_nb -- "
                            << "transaction failed.";
+                endTransferBufferQueueBlocks(c2workItems, false, true);
                 return;
             }
-            yieldBufferQueueBlocks(c2workItems, true);
+            endTransferBufferQueueBlocks(c2workItems, true, true);
         }
     }
 
@@ -254,13 +257,14 @@
 
     WorkBundle flushedWorkBundle;
     Status res = static_cast<Status>(c2res);
+    beginTransferBufferQueueBlocks(c2flushedWorks, true);
     if (c2res == C2_OK) {
         if (!objcpy(&flushedWorkBundle, c2flushedWorks, &mBufferPoolSender)) {
             res = Status::CORRUPTED;
         }
     }
     _hidl_cb(res, flushedWorkBundle);
-    yieldBufferQueueBlocks(c2flushedWorks, true);
+    endTransferBufferQueueBlocks(c2flushedWorks, true, true);
     return Void();
 }
 
diff --git a/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/ClientBlockHelper.h b/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/ClientBlockHelper.h
index be429ac..0a2298c 100644
--- a/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/ClientBlockHelper.h
+++ b/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/ClientBlockHelper.h
@@ -28,62 +28,42 @@
 namespace V1_0 {
 namespace utils {
 
-
 // BufferQueue-Based Block Operations
 // ==================================
 
-// Create a GraphicBuffer object from a graphic block and attach it to an
-// IGraphicBufferProducer.
-status_t attachToBufferQueue(const C2ConstGraphicBlock& block,
-                             const sp<IGraphicBufferProducer>& igbp,
-                             uint32_t generation,
-                             int32_t* bqSlot);
+// Manage BufferQueue and graphic blocks for both component and codec.
+// Manage graphic blocks ownership consistently during surface change.
+struct OutputBufferQueue {
 
-// Return false if block does not come from a bufferqueue-based blockpool.
-// Otherwise, extract generation, bqId and bqSlot and return true.
-bool getBufferQueueAssignment(const C2ConstGraphicBlock& block,
-                              uint32_t* generation,
-                              uint64_t* bqId,
-                              int32_t* bqSlot);
+    OutputBufferQueue();
 
-// Assign the given block to a bufferqueue so that when the block is destroyed,
-// cancelBuffer() will be called.
-//
-// If the block does not come from a bufferqueue-based blockpool, this function
-// returns false.
-//
-// If the block already has a bufferqueue assignment that matches the given one,
-// the function returns true.
-//
-// If the block already has a bufferqueue assignment that does not match the
-// given one, the block will be reassigned to the given bufferqueue. This
-// will call attachBuffer() on the given igbp. The function then returns true on
-// success or false on any failure during the operation.
-//
-// Note: This function should be called after detachBuffer() or dequeueBuffer()
-// is called manually.
-bool holdBufferQueueBlock(const C2ConstGraphicBlock& block,
-                          const sp<IGraphicBufferProducer>& igbp,
-                          uint64_t bqId,
-                          uint32_t generation);
+    ~OutputBufferQueue();
 
-// Call holdBufferQueueBlock() on input or output blocks in the given workList.
-// Since the bufferqueue assignment for input and output buffers can be
-// different, this function takes forInput to determine whether the given
-// bufferqueue is for input buffers or output buffers. (The default value of
-// forInput is false.)
-//
-// In the (rare) case that both input and output buffers are bufferqueue-based,
-// this function must be called twice, once for the input buffers and once for
-// the output buffers.
-//
-// Note: This function should be called after WorkBundle has been received from
-// another process.
-void holdBufferQueueBlocks(const std::list<std::unique_ptr<C2Work>>& workList,
-                           const sp<IGraphicBufferProducer>& igbp,
-                           uint64_t bqId,
-                           uint32_t generation,
-                           bool forInput = false);
+    // Configure a new surface to render graphic blocks.
+    // Graphic blocks from older surface will be migrated to new surface.
+    bool configure(const sp<IGraphicBufferProducer>& igbp,
+                   uint32_t generation,
+                   uint64_t bqId);
+
+    // Render a graphic block to current surface.
+    status_t outputBuffer(
+            const C2ConstGraphicBlock& block,
+            const BnGraphicBufferProducer::QueueBufferInput& input,
+            BnGraphicBufferProducer::QueueBufferOutput* output);
+
+    // Call holdBufferQueueBlock() on output blocks in the given workList.
+    // The OutputBufferQueue will take the ownership of output blocks.
+    //
+    // Note: This function should be called after WorkBundle has been received
+    // from another process.
+    void holdBufferQueueBlocks(
+            const std::list<std::unique_ptr<C2Work>>& workList);
+
+private:
+
+    class Impl;
+    std::unique_ptr<Impl> mImpl;
+};
 
 }  // namespace utils
 }  // namespace V1_0
diff --git a/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/types.h b/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/types.h
index 0ddfec5..a9928b3 100644
--- a/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/types.h
+++ b/media/codec2/hidl/1.0/utils/include/codec2/hidl/1.0/types.h
@@ -299,26 +299,46 @@
 // BufferQueue-Based Block Operations
 // ==================================
 
-// Disassociate the given block with its designated bufferqueue so that
-// cancelBuffer() will not be called when the block is destroyed. If the block
-// does not have a designated bufferqueue, the function returns false.
-// Otherwise, it returns true.
+// Call before transferring block to other processes.
 //
-// Note: This function should be called after attachBuffer() or queueBuffer() is
-// called manually.
-bool yieldBufferQueueBlock(const C2ConstGraphicBlock& block);
+// The given block is ready to transfer to other processes. This will guarantee
+// the given block data is not mutated by bufferqueue migration.
+bool beginTransferBufferQueueBlock(const C2ConstGraphicBlock& block);
 
-// Call yieldBufferQueueBlock() on blocks in the given workList. processInput
-// determines whether input blocks are yielded. processOutput works similarly on
-// output blocks. (The default value of processInput is false while the default
-// value of processOutput is true. This implies that in most cases, only output
-// buffers contain bufferqueue-based blocks.)
+// Call beginTransferBufferQueueBlock() on blocks in the given workList.
+// processInput determines whether input blocks are yielded. processOutput
+// works similarly on output blocks. (The default value of processInput is
+// false while the default value of processOutput is true. This implies that in
+// most cases, only output buffers contain bufferqueue-based blocks.)
+void beginTransferBufferQueueBlocks(
+        const std::list<std::unique_ptr<C2Work>>& workList,
+        bool processInput = false,
+        bool processOutput = true);
+
+// Call after transferring block is finished and make sure that
+// beginTransferBufferQueueBlock() is called before.
 //
-// Note: This function should be called after WorkBundle has been successfully
-// sent over the Treble boundary to another process.
-void yieldBufferQueueBlocks(const std::list<std::unique_ptr<C2Work>>& workList,
-                            bool processInput = false,
-                            bool processOutput = true);
+// The transfer of given block is finished. If transfer is successful the given
+// block is not owned by process anymore. Since transfer is finished the given
+// block data is OK to mutate by bufferqueue migration after this call.
+bool endTransferBufferQueueBlock(const C2ConstGraphicBlock& block,
+                                 bool transfer);
+
+// Call endTransferBufferQueueBlock() on blocks in the given workList.
+// processInput determines whether input blocks are yielded. processOutput
+// works similarly on output blocks. (The default value of processInput is
+// false while the default value of processOutput is true. This implies that in
+// most cases, only output buffers contain bufferqueue-based blocks.)
+void endTransferBufferQueueBlocks(
+        const std::list<std::unique_ptr<C2Work>>& workList,
+        bool transfer,
+        bool processInput = false,
+        bool processOutput = true);
+
+// The given block is ready to be rendered. the given block is not owned by
+// process anymore. If migration is in progress, this returns false in order
+// not to render.
+bool displayBufferQueueBlock(const C2ConstGraphicBlock& block);
 
 }  // namespace utils
 }  // namespace V1_0
diff --git a/media/codec2/hidl/1.0/utils/types.cpp b/media/codec2/hidl/1.0/utils/types.cpp
index 3fd2f92..07dbf67 100644
--- a/media/codec2/hidl/1.0/utils/types.cpp
+++ b/media/codec2/hidl/1.0/utils/types.cpp
@@ -737,7 +737,15 @@
                     bufferPoolSender, baseBlocks, baseBlockIndices);
         }
     case _C2BlockPoolData::TYPE_BUFFERQUEUE:
-        // Do the same thing as a NATIVE block.
+        uint32_t gen;
+        uint64_t bqId;
+        int32_t bqSlot;
+        // Update handle if migration happened.
+        if (_C2BlockFactory::GetBufferQueueData(
+                blockPoolData, &gen, &bqId, &bqSlot)) {
+            android::MigrateNativeCodec2GrallocHandle(
+                    const_cast<native_handle_t*>(handle), gen, bqId, bqSlot);
+        }
         return _addBaseBlock(
                 index, handle,
                 baseBlocks, baseBlockIndices);
@@ -1772,20 +1780,53 @@
 
 } // unnamed namespace
 
-bool yieldBufferQueueBlock(const C2ConstGraphicBlock& block) {
+bool beginTransferBufferQueueBlock(const C2ConstGraphicBlock& block) {
     std::shared_ptr<_C2BlockPoolData> data =
             _C2BlockFactory::GetGraphicBlockPoolData(block);
     if (data && _C2BlockFactory::GetBufferQueueData(data)) {
-        _C2BlockFactory::YieldBlockToBufferQueue(data);
+        _C2BlockFactory::BeginTransferBlockToClient(data);
         return true;
     }
     return false;
 }
 
-void yieldBufferQueueBlocks(
+void beginTransferBufferQueueBlocks(
         const std::list<std::unique_ptr<C2Work>>& workList,
         bool processInput, bool processOutput) {
-    forEachBlock(workList, yieldBufferQueueBlock, processInput, processOutput);
+    forEachBlock(workList, beginTransferBufferQueueBlock,
+                 processInput, processOutput);
+}
+
+bool endTransferBufferQueueBlock(
+        const C2ConstGraphicBlock& block,
+        bool transfer) {
+    std::shared_ptr<_C2BlockPoolData> data =
+            _C2BlockFactory::GetGraphicBlockPoolData(block);
+    if (data && _C2BlockFactory::GetBufferQueueData(data)) {
+        _C2BlockFactory::EndTransferBlockToClient(data, transfer);
+        return true;
+    }
+    return false;
+}
+
+void endTransferBufferQueueBlocks(
+        const std::list<std::unique_ptr<C2Work>>& workList,
+        bool transfer,
+        bool processInput, bool processOutput) {
+    forEachBlock(workList,
+                 std::bind(endTransferBufferQueueBlock,
+                           std::placeholders::_1, transfer),
+                 processInput, processOutput);
+}
+
+bool displayBufferQueueBlock(const C2ConstGraphicBlock& block) {
+    std::shared_ptr<_C2BlockPoolData> data =
+            _C2BlockFactory::GetGraphicBlockPoolData(block);
+    if (data && _C2BlockFactory::GetBufferQueueData(data)) {
+        _C2BlockFactory::DisplayBlockToBufferQueue(data);
+        return true;
+    }
+    return false;
 }
 
 }  // namespace utils
diff --git a/media/codec2/hidl/client/client.cpp b/media/codec2/hidl/client/client.cpp
index 6aca4a3..2b417a6 100644
--- a/media/codec2/hidl/client/client.cpp
+++ b/media/codec2/hidl/client/client.cpp
@@ -1080,15 +1080,7 @@
 void Codec2Client::Component::handleOnWorkDone(
         const std::list<std::unique_ptr<C2Work>> &workItems) {
     // Output bufferqueue-based blocks' lifetime management
-    mOutputBufferQueueMutex.lock();
-    sp<IGraphicBufferProducer> igbp = mOutputIgbp;
-    uint64_t bqId = mOutputBqId;
-    uint32_t generation = mOutputGeneration;
-    mOutputBufferQueueMutex.unlock();
-
-    if (igbp) {
-        holdBufferQueueBlocks(workItems, igbp, bqId, generation);
-    }
+    mOutputBufferQueue.holdBufferQueueBlocks(workItems);
 }
 
 c2_status_t Codec2Client::Component::queue(
@@ -1151,15 +1143,7 @@
     }
 
     // Output bufferqueue-based blocks' lifetime management
-    mOutputBufferQueueMutex.lock();
-    sp<IGraphicBufferProducer> igbp = mOutputIgbp;
-    uint64_t bqId = mOutputBqId;
-    uint32_t generation = mOutputGeneration;
-    mOutputBufferQueueMutex.unlock();
-
-    if (igbp) {
-        holdBufferQueueBlocks(*flushedWork, igbp, bqId, generation);
-    }
+    mOutputBufferQueue.holdBufferQueueBlocks(*flushedWork);
 
     return status;
 }
@@ -1239,15 +1223,31 @@
         C2BlockPool::local_id_t blockPoolId,
         const sp<IGraphicBufferProducer>& surface,
         uint32_t generation) {
-    sp<HGraphicBufferProducer2> igbp =
-            surface->getHalInterface<HGraphicBufferProducer2>();
+    uint64_t bqId = 0;
+    sp<IGraphicBufferProducer> nullIgbp;
+    sp<HGraphicBufferProducer2> nullHgbp;
 
-    if (!igbp) {
+    sp<HGraphicBufferProducer2> igbp = surface ?
+            surface->getHalInterface<HGraphicBufferProducer2>() : nullHgbp;
+    if (surface && !igbp) {
         igbp = new B2HGraphicBufferProducer2(surface);
     }
 
+    if (!surface) {
+        mOutputBufferQueue.configure(nullIgbp, generation, 0);
+    } else if (surface->getUniqueId(&bqId) != OK) {
+        LOG(ERROR) << "setOutputSurface -- "
+                   "cannot obtain bufferqueue id.";
+        bqId = 0;
+        mOutputBufferQueue.configure(nullIgbp, generation, 0);
+    } else {
+        mOutputBufferQueue.configure(surface, generation, bqId);
+    }
+    ALOGD("generation remote change %u", generation);
+
     Return<Status> transStatus = mBase->setOutputSurface(
-            static_cast<uint64_t>(blockPoolId), igbp);
+            static_cast<uint64_t>(blockPoolId),
+            bqId == 0 ? nullHgbp : igbp);
     if (!transStatus.isOk()) {
         LOG(ERROR) << "setOutputSurface -- transaction failed.";
         return C2_TRANSACTION_FAILED;
@@ -1256,18 +1256,6 @@
             static_cast<c2_status_t>(static_cast<Status>(transStatus));
     if (status != C2_OK) {
         LOG(DEBUG) << "setOutputSurface -- call failed: " << status << ".";
-    } else {
-        std::lock_guard<std::mutex> lock(mOutputBufferQueueMutex);
-        if (mOutputIgbp != surface) {
-            mOutputIgbp = surface;
-            if (!surface) {
-                mOutputBqId = 0;
-            } else if (surface->getUniqueId(&mOutputBqId) != OK) {
-                LOG(ERROR) << "setOutputSurface -- "
-                              "cannot obtain bufferqueue id.";
-            }
-        }
-        mOutputGeneration = generation;
     }
     return status;
 }
@@ -1276,74 +1264,7 @@
         const C2ConstGraphicBlock& block,
         const QueueBufferInput& input,
         QueueBufferOutput* output) {
-    uint32_t generation;
-    uint64_t bqId;
-    int32_t bqSlot;
-    if (!getBufferQueueAssignment(block, &generation, &bqId, &bqSlot) ||
-            bqId == 0) {
-        // Block not from bufferqueue -- it must be attached before queuing.
-
-        mOutputBufferQueueMutex.lock();
-        sp<IGraphicBufferProducer> outputIgbp = mOutputIgbp;
-        uint32_t outputGeneration = mOutputGeneration;
-        mOutputBufferQueueMutex.unlock();
-
-        status_t status = attachToBufferQueue(block,
-                                               outputIgbp,
-                                               outputGeneration,
-                                               &bqSlot);
-        if (status != OK) {
-            LOG(WARNING) << "queueToOutputSurface -- attaching failed.";
-            return INVALID_OPERATION;
-        }
-
-        status = outputIgbp->queueBuffer(static_cast<int>(bqSlot),
-                                         input, output);
-        if (status != OK) {
-            LOG(ERROR) << "queueToOutputSurface -- queueBuffer() failed "
-                          "on non-bufferqueue-based block. "
-                          "Error = " << status << ".";
-            return status;
-        }
-        return OK;
-    }
-
-    mOutputBufferQueueMutex.lock();
-    sp<IGraphicBufferProducer> outputIgbp = mOutputIgbp;
-    uint64_t outputBqId = mOutputBqId;
-    uint32_t outputGeneration = mOutputGeneration;
-    mOutputBufferQueueMutex.unlock();
-
-    if (!outputIgbp) {
-        LOG(VERBOSE) << "queueToOutputSurface -- output surface is null.";
-        return NO_INIT;
-    }
-
-    if (bqId != outputBqId || generation != outputGeneration) {
-        if (!holdBufferQueueBlock(block, mOutputIgbp, mOutputBqId, mOutputGeneration)) {
-            LOG(ERROR) << "queueToOutputSurface -- migration failed.";
-            return DEAD_OBJECT;
-        }
-        if (!getBufferQueueAssignment(block, &generation, &bqId, &bqSlot)) {
-            LOG(ERROR) << "queueToOutputSurface -- corrupted bufferqueue assignment.";
-            return UNKNOWN_ERROR;
-        }
-    }
-
-    status_t status = outputIgbp->queueBuffer(static_cast<int>(bqSlot),
-                                              input, output);
-    if (status != OK) {
-        LOG(DEBUG) << "queueToOutputSurface -- queueBuffer() failed "
-                      "on bufferqueue-based block. "
-                      "Error = " << status << ".";
-        return status;
-    }
-    if (!yieldBufferQueueBlock(block)) {
-        LOG(DEBUG) << "queueToOutputSurface -- cannot yield "
-                      "bufferqueue-based block to the bufferqueue.";
-        return UNKNOWN_ERROR;
-    }
-    return OK;
+    return mOutputBufferQueue.outputBuffer(block, input, output);
 }
 
 c2_status_t Codec2Client::Component::connectToInputSurface(
diff --git a/media/codec2/hidl/client/include/codec2/hidl/client.h b/media/codec2/hidl/client/include/codec2/hidl/client.h
index 03db515..b8a7fb5 100644
--- a/media/codec2/hidl/client/include/codec2/hidl/client.h
+++ b/media/codec2/hidl/client/include/codec2/hidl/client.h
@@ -369,10 +369,8 @@
     ::android::hardware::media::c2::V1_0::utils::DefaultBufferPoolSender
             mBufferPoolSender;
 
-    std::mutex mOutputBufferQueueMutex;
-    sp<IGraphicBufferProducer> mOutputIgbp;
-    uint64_t mOutputBqId;
-    uint32_t mOutputGeneration;
+    ::android::hardware::media::c2::V1_0::utils::OutputBufferQueue
+            mOutputBufferQueue;
 
     static c2_status_t setDeathListener(
             const std::shared_ptr<Component>& component,