3 files changed, 74 insertions, 2 deletions

diff --git a/libs/gui/BLASTBufferQueue.cpp b/libs/gui/BLASTBufferQueue.cpp
index 94998e5ab7..c65eafa541 100644
--- a/libs/gui/BLASTBufferQueue.cpp
+++ b/libs/gui/BLASTBufferQueue.cpp
@@ -1266,6 +1266,11 @@ void BLASTBufferQueue::setTransactionHangCallback(
     mTransactionHangCallback = std::move(callback);
 }
 
+void BLASTBufferQueue::setApplyToken(sp<IBinder> applyToken) {
+    std::lock_guard _lock{mMutex};
+    mApplyToken = std::move(applyToken);
+}
+
 #if COM_ANDROID_GRAPHICS_LIBGUI_FLAGS(BUFFER_RELEASE_CHANNEL)
 
 BLASTBufferQueue::BufferReleaseReader::BufferReleaseReader(
diff --git a/libs/gui/include/gui/BLASTBufferQueue.h b/libs/gui/include/gui/BLASTBufferQueue.h
index 868dbd06ec..ba58a15e17 100644
--- a/libs/gui/include/gui/BLASTBufferQueue.h
+++ b/libs/gui/include/gui/BLASTBufferQueue.h
@@ -143,7 +143,7 @@ public:
      * indicates the reason for the hang.
      */
     void setTransactionHangCallback(std::function<void(const std::string&)> callback);
-
+    void setApplyToken(sp<IBinder>);
     virtual ~BLASTBufferQueue();
 
     void onFirstRef() override;
@@ -272,7 +272,7 @@ private:
 
     // Queues up transactions using this token in SurfaceFlinger. This prevents queued up
     // transactions from other parts of the client from blocking this transaction.
-    const sp<IBinder> mApplyToken GUARDED_BY(mMutex) = sp<BBinder>::make();
+    sp<IBinder> mApplyToken GUARDED_BY(mMutex) = sp<BBinder>::make();
 
     // Guards access to mDequeueTimestamps since we cannot hold to mMutex in onFrameDequeued or
     // we will deadlock.
diff --git a/libs/gui/tests/BLASTBufferQueue_test.cpp b/libs/gui/tests/BLASTBufferQueue_test.cpp
index eb2a61d151..53f4a36c42 100644
--- a/libs/gui/tests/BLASTBufferQueue_test.cpp
+++ b/libs/gui/tests/BLASTBufferQueue_test.cpp
@@ -186,6 +186,10 @@ public:
         mBlastBufferQueueAdapter->mergeWithNextTransaction(merge, frameNumber);
     }
 
+    void setApplyToken(sp<IBinder> applyToken) {
+        mBlastBufferQueueAdapter->setApplyToken(std::move(applyToken));
+    }
+
 private:
     sp<TestBLASTBufferQueue> mBlastBufferQueueAdapter;
 };
@@ -511,6 +515,69 @@ TEST_F(BLASTBufferQueueTest, TripleBuffering) {
     adapter.waitForCallbacks();
 }
 
+class WaitForCommittedCallback {
+public:
+    WaitForCommittedCallback() = default;
+    ~WaitForCommittedCallback() = default;
+
+    void wait() {
+        std::unique_lock lock(mMutex);
+        cv.wait(lock, [this] { return mCallbackReceived; });
+    }
+
+    void notify() {
+        std::unique_lock lock(mMutex);
+        mCallbackReceived = true;
+        cv.notify_one();
+        mCallbackReceivedTimeStamp = std::chrono::system_clock::now();
+    }
+    auto getCallback() {
+        return [this](void* /* unused context */, nsecs_t /* latchTime */,
+                      const sp<Fence>& /* presentFence */,
+                      const std::vector<SurfaceControlStats>& /* stats */) { notify(); };
+    }
+    std::chrono::time_point<std::chrono::system_clock> mCallbackReceivedTimeStamp;
+
+private:
+    std::mutex mMutex;
+    std::condition_variable cv;
+    bool mCallbackReceived = false;
+};
+
+TEST_F(BLASTBufferQueueTest, setApplyToken) {
+    sp<IBinder> applyToken = sp<BBinder>::make();
+    WaitForCommittedCallback firstTransaction;
+    WaitForCommittedCallback secondTransaction;
+    {
+        BLASTBufferQueueHelper adapter(mSurfaceControl, mDisplayWidth, mDisplayHeight);
+        adapter.setApplyToken(applyToken);
+        sp<IGraphicBufferProducer> igbProducer;
+        setUpProducer(adapter, igbProducer);
+
+        Transaction t;
+        t.addTransactionCommittedCallback(firstTransaction.getCallback(), nullptr);
+        adapter.mergeWithNextTransaction(&t, 1);
+        queueBuffer(igbProducer, 127, 127, 127,
+                    /*presentTimeDelay*/ std::chrono::nanoseconds(500ms).count());
+    }
+    {
+        BLASTBufferQueueHelper adapter(mSurfaceControl, mDisplayWidth, mDisplayHeight);
+        adapter.setApplyToken(applyToken);
+        sp<IGraphicBufferProducer> igbProducer;
+        setUpProducer(adapter, igbProducer);
+
+        Transaction t;
+        t.addTransactionCommittedCallback(secondTransaction.getCallback(), nullptr);
+        adapter.mergeWithNextTransaction(&t, 1);
+        queueBuffer(igbProducer, 127, 127, 127, /*presentTimeDelay*/ 0);
+    }
+
+    firstTransaction.wait();
+    secondTransaction.wait();
+    EXPECT_GT(secondTransaction.mCallbackReceivedTimeStamp,
+              firstTransaction.mCallbackReceivedTimeStamp);
+}
+
 TEST_F(BLASTBufferQueueTest, SetCrop_Item) {
     uint8_t r = 255;
     uint8_t g = 0;