Native VK Swapchain: Fix the lifetime of DisplayVk::ColorBufferInfo

The release of ColorBuffer itself and the release of the id of the
ColorBuffer doesn't happen at the same time. More than one ColorBuffers
may share the same ID at the same time - one is in use, and others are
wait to be freed. Therefore, if we try to use ColorBuffer id to reference
a ColorBuffer inside DisplayVk we may encounter a collision. Therefore,
instead of the id, the client of DisplayVk now creates a
DisplayVk::DisplayBufferInfo object and have the lifetime of that object
via shared_ptr. Inside DisplayVk, when post is called a weak_ptr is used
an compared if current composition is the same as the previous one.

* Rename DisplayVk::ColorBufferInfo to DisplayVk::DisplayBufferInfo
* Remove the releaseDisplayImport interface and its present in
  ColorBuffer destructor. Now the lifetime of a DisplayBufferInfo is
  entirely controlled by the object created via
  DisplayVk::createDisplayBuffer, and is released automatically when
  ColorBufferRef is released.

Test: run the emulator with host vulkan native swapchain, and it won't crash now

Change-Id: If3da93b059bcccfb1bbce40134d4d0789676a00b
diff --git a/stream-servers/ColorBuffer.cpp b/stream-servers/ColorBuffer.cpp
index f12f2da..24d17ae 100644
--- a/stream-servers/ColorBuffer.cpp
+++ b/stream-servers/ColorBuffer.cpp
@@ -299,8 +299,6 @@
 ColorBuffer::~ColorBuffer() {
     RecursiveScopedHelperContext context(m_helper);
 
-    m_helper->releaseDisplayImport(getHndl());
-
     if (m_blitEGLImage) {
         s_egl.eglDestroyImageKHR(m_display, m_blitEGLImage);
     }
@@ -935,15 +933,14 @@
 
 bool ColorBuffer::importMemory(
 #ifdef _WIN32
-        void* handle,
+    void* handle,
 #else
-        int handle,
+    int handle,
 #endif
-        uint64_t size,
-        bool dedicated,
-        bool linearTiling,
-        bool vulkanOnly) {
+    uint64_t size, bool dedicated, bool linearTiling, bool vulkanOnly,
+    std::shared_ptr<DisplayVk::DisplayBufferInfo> displayBufferVk) {
     RecursiveScopedHelperContext context(m_helper);
+    m_displayBufferVk = std::move(displayBufferVk);
     s_gles2.glCreateMemoryObjectsEXT(1, &m_memoryObject);
     if (dedicated) {
         static const GLint DEDICATED_FLAG = GL_TRUE;
diff --git a/stream-servers/ColorBuffer.h b/stream-servers/ColorBuffer.h
index 6249b77..9d4eb91 100644
--- a/stream-servers/ColorBuffer.h
+++ b/stream-servers/ColorBuffer.h
@@ -22,12 +22,13 @@
 #include <GLES3/gl3.h>
 #include "base/Stream.h"
 // #include "android/skin/rect.h"
-#include "snapshot/LazySnapshotObj.h"
+#include <memory>
+
+#include "DisplayVk.h"
 #include "FrameworkFormats.h"
 #include "Hwc2.h"
 #include "RenderContext.h"
-
-#include <memory>
+#include "snapshot/LazySnapshotObj.h"
 
 class TextureDraw;
 class TextureResize;
@@ -74,7 +75,6 @@
         virtual void teardownContext() = 0;
         virtual TextureDraw* getTextureDraw() const = 0;
         virtual bool isBound() const = 0;
-        virtual void releaseDisplayImport(HandleType) = 0;
     };
 
     // Helper class to use a ColorBuffer::Helper context.
@@ -251,16 +251,19 @@
     void postLayer(ComposeLayer* l, int frameWidth, int frameHeight);
     GLuint getTexture();
 
+    const std::shared_ptr<DisplayVk::DisplayBufferInfo>& getDisplayBufferVk()
+        const {
+        return m_displayBufferVk;
+    };
+
     bool importMemory(
 #ifdef _WIN32
         void* handle,
 #else
         int handle,
 #endif
-        uint64_t size,
-        bool dedicated,
-        bool linearTiling,
-        bool vulkanOnly);
+        uint64_t size, bool dedicated, bool linearTiling, bool vulkanOnly,
+        std::shared_ptr<DisplayVk::DisplayBufferInfo> displayBufferVk);
     void setInUse(bool inUse);
     bool isInUse() const { return m_inUse; }
 
@@ -320,6 +323,9 @@
     GLuint m_buf = 0;
     uint32_t m_displayId = 0;
     bool m_BRSwizzle = false;
+    // Won't share with others so that m_displayBufferVk lives shorter than this
+    // ColorBuffer.
+    std::shared_ptr<DisplayVk::DisplayBufferInfo> m_displayBufferVk;
 };
 
 typedef std::shared_ptr<ColorBuffer> ColorBufferPtr;
diff --git a/stream-servers/DisplayVk.cpp b/stream-servers/DisplayVk.cpp
index f10a53b..78e5046 100644
--- a/stream-servers/DisplayVk.cpp
+++ b/stream-servers/DisplayVk.cpp
@@ -1,11 +1,11 @@
 #include "DisplayVk.h"
 
-#include "ErrorLog.h"
-#include "NativeSubWindow.h"
-
 #include <glm/glm.hpp>
 #include <glm/gtx/matrix_transform_2d.hpp>
 
+#include "ErrorLog.h"
+#include "NativeSubWindow.h"
+
 DisplayVk::DisplayVk(const goldfish_vk::VulkanDispatch &vk,
                      VkPhysicalDevice vkPhysicalDevice,
                      uint32_t swapChainQueueFamilyIndex,
@@ -109,55 +109,21 @@
     m_surfaceState = std::move(surfaceState);
 }
 
-void DisplayVk::importVkImage(HandleType id, VkImage image, VkFormat format,
-                              uint32_t width, uint32_t height) {
-    if (m_colorBuffers.find(id) != m_colorBuffers.end()) {
-        ERR("%s(%s:%d): DisplayVk can't import VkImage with duplicate id = "
-            "%" PRIu64 ".\n",
-            __FUNCTION__, __FILE__, static_cast<int>(__LINE__),
-            static_cast<uint64_t>(id));
-        ::abort();
-    }
-    ColorBufferInfo cbInfo = {};
-    cbInfo.m_width = width;
-    cbInfo.m_height = height;
-    cbInfo.m_vkFormat = format;
-
-    VkImageViewCreateInfo imageViewCi = {};
-    imageViewCi.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
-    imageViewCi.image = image;
-    imageViewCi.viewType = VK_IMAGE_VIEW_TYPE_2D;
-    imageViewCi.format = format;
-    imageViewCi.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
-    imageViewCi.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
-    imageViewCi.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
-    imageViewCi.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
-    imageViewCi.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
-    imageViewCi.subresourceRange.baseMipLevel = 0;
-    imageViewCi.subresourceRange.levelCount = 1;
-    imageViewCi.subresourceRange.baseArrayLayer = 0;
-    imageViewCi.subresourceRange.layerCount = 1;
-    VK_CHECK(m_vk.vkCreateImageView(m_vkDevice, &imageViewCi, nullptr,
-                                    &cbInfo.m_vkImageView));
-    m_colorBuffers.emplace(id, cbInfo);
+std::shared_ptr<DisplayVk::DisplayBufferInfo> DisplayVk::createDisplayBuffer(
+    VkImage image, VkFormat format, uint32_t width, uint32_t height) {
+    return std::shared_ptr<DisplayBufferInfo>(
+        new DisplayBufferInfo(m_vk, m_vkDevice, width, height, format, image));
 }
 
-void DisplayVk::releaseImport(HandleType id) {
-    auto it = m_colorBuffers.find(id);
-    if (it == m_colorBuffers.end()) {
-        return;
-    }
-    m_vk.vkDestroyImageView(m_vkDevice, it->second.m_vkImageView, nullptr);
-    m_colorBuffers.erase(it);
-}
-
-void DisplayVk::post(HandleType id) {
+void DisplayVk::post(
+    const std::shared_ptr<DisplayBufferInfo> &displayBufferPtr) {
     if (m_swapChainStateVk == nullptr || m_compositorVk == nullptr) {
         ERR("%s(%s:%d): Haven't bound to a surface, can't post color buffer.\n",
             __FUNCTION__, __FILE__, static_cast<int>(__LINE__));
         return;
     }
-    const auto &surfaceState = *m_surfaceState;
+    auto &surfaceState = *m_surfaceState;
+    const auto &db = *displayBufferPtr;
 
     VK_CHECK(m_vk.vkWaitForFences(m_vkDevice, 1, &m_frameDrawCompleteFence,
                                   VK_TRUE, UINT64_MAX));
@@ -165,26 +131,26 @@
     VK_CHECK(m_vk.vkAcquireNextImageKHR(
         m_vkDevice, m_swapChainStateVk->getSwapChain(), UINT64_MAX,
         m_imageReadySem, VK_NULL_HANDLE, &imageIndex));
-    if (!surfaceState.m_prevColorBuffer.has_value() ||
-        surfaceState.m_prevColorBuffer.value() != id) {
-        const auto &cb = m_colorBuffers.at(id);
-        if (!canComposite(cb.m_vkFormat)) {
-            ERR("%s(%s:%d): Can't composite the ColorBuffer(id = %" PRIu64
+    auto maybePrevDisplayBuffer = surfaceState.m_prevDisplayBuffer.lock();
+    if (!maybePrevDisplayBuffer || maybePrevDisplayBuffer != displayBufferPtr) {
+        if (!canComposite(db.m_vkFormat)) {
+            ERR("%s(%s:%d): Can't composite the DisplayBuffer(0x%" PRIxPTR
                 "). The image(VkFormat = %" PRIu64 ") can be sampled from.\n",
                 __FUNCTION__, __FILE__, static_cast<int>(__LINE__),
-                static_cast<uint64_t>(id),
-                static_cast<uint64_t>(cb.m_vkFormat));
+                reinterpret_cast<uintptr_t>(displayBufferPtr.get()),
+                static_cast<uint64_t>(db.m_vkFormat));
             return;
         }
         auto composition = std::make_unique<Composition>(
-            cb.m_vkImageView, m_compositionVkSampler, cb.m_width, cb.m_height);
+            db.m_vkImageView, m_compositionVkSampler, db.m_width, db.m_height);
         composition->m_transform =
             glm::scale(glm::mat3(1.0),
                        glm::vec2(static_cast<float>(surfaceState.m_width) /
-                                     static_cast<float>(cb.m_width),
+                                     static_cast<float>(db.m_width),
                                  static_cast<float>(surfaceState.m_height) /
-                                     static_cast<float>(cb.m_height)));
+                                     static_cast<float>(db.m_height)));
         m_compositorVk->setComposition(imageIndex, std::move(composition));
+        surfaceState.m_prevDisplayBuffer = displayBufferPtr;
     }
 
     auto cmdBuff = m_compositorVk->getCommandBuffer(imageIndex);
@@ -223,8 +189,41 @@
         return it->second;
     }
     VkFormatProperties formatProps = {};
-    m_vk.vkGetPhysicalDeviceFormatProperties(m_vkPhysicalDevice, format, &formatProps);
-    bool res = formatProps.optimalTilingFeatures & VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT;
+    m_vk.vkGetPhysicalDeviceFormatProperties(m_vkPhysicalDevice, format,
+                                             &formatProps);
+    bool res =
+        formatProps.optimalTilingFeatures & VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT;
     m_canComposite.emplace(format, res);
     return res;
+}
+
+DisplayVk::DisplayBufferInfo::DisplayBufferInfo(
+    const goldfish_vk::VulkanDispatch &vk, VkDevice vkDevice, uint32_t width,
+    uint32_t height, VkFormat format, VkImage image)
+    : m_vk(vk),
+      m_vkDevice(vkDevice),
+      m_width(width),
+      m_height(height),
+      m_vkFormat(format),
+      m_vkImageView(VK_NULL_HANDLE) {
+    VkImageViewCreateInfo imageViewCi = {};
+    imageViewCi.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
+    imageViewCi.image = image;
+    imageViewCi.viewType = VK_IMAGE_VIEW_TYPE_2D;
+    imageViewCi.format = format;
+    imageViewCi.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
+    imageViewCi.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
+    imageViewCi.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
+    imageViewCi.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
+    imageViewCi.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+    imageViewCi.subresourceRange.baseMipLevel = 0;
+    imageViewCi.subresourceRange.levelCount = 1;
+    imageViewCi.subresourceRange.baseArrayLayer = 0;
+    imageViewCi.subresourceRange.layerCount = 1;
+    VK_CHECK(m_vk.vkCreateImageView(m_vkDevice, &imageViewCi, nullptr,
+                                    &m_vkImageView));
+}
+
+DisplayVk::DisplayBufferInfo::~DisplayBufferInfo() {
+    m_vk.vkDestroyImageView(m_vkDevice, m_vkImageView, nullptr);
 }
\ No newline at end of file
diff --git a/stream-servers/DisplayVk.h b/stream-servers/DisplayVk.h
index 05dae47..88e86c2 100644
--- a/stream-servers/DisplayVk.h
+++ b/stream-servers/DisplayVk.h
@@ -3,10 +3,8 @@
 
 #include <functional>
 #include <memory>
-#include <optional>
 #include <unordered_map>
 
-#include "ColorBuffer.h"
 #include "CompositorVk.h"
 #include "RenderContext.h"
 #include "SwapChainStateVk.h"
@@ -17,17 +15,39 @@
 
 class DisplayVk {
    public:
+    class DisplayBufferInfo {
+       public:
+        ~DisplayBufferInfo();
+
+       private:
+        DisplayBufferInfo(const goldfish_vk::VulkanDispatch &, VkDevice,
+                          uint32_t width, uint32_t height, VkFormat, VkImage);
+
+        const goldfish_vk::VulkanDispatch &m_vk;
+        VkDevice m_vkDevice;
+        uint32_t m_width;
+        uint32_t m_height;
+        VkFormat m_vkFormat;
+
+        VkImageView m_vkImageView;
+
+        friend class DisplayVk;
+    };
     DisplayVk(const goldfish_vk::VulkanDispatch &, VkPhysicalDevice,
               uint32_t swapChainQueueFamilyIndex,
               uint32_t compositorQueueFamilyIndex, VkDevice,
               VkQueue compositorVkQueue, VkQueue swapChainVkQueue);
     ~DisplayVk();
     void bindToSurface(VkSurfaceKHR, uint32_t width, uint32_t height);
-    void importVkImage(HandleType id, VkImage, VkFormat, uint32_t width,
-                       uint32_t height);
-    void releaseImport(HandleType id);
-    // To display a colorbuffer, first import it by calling importVkImage
-    void post(HandleType id);
+    // The caller is responsible to make sure the VkImage lives longer than the
+    // DisplayBufferInfo created here. However, given that DisplayBufferInfo
+    // lives in a shared_ptr, the potential lifetime of DisplayBufferInfo is
+    // aligned to DisplayVk when DisplayVk::m_surfaceState::m_prevDisplayBuffer
+    // is locked and upgraded to a shared_ptr in DisplayVk::post.
+    std::shared_ptr<DisplayBufferInfo> createDisplayBuffer(VkImage, VkFormat,
+                                                           uint32_t width,
+                                                           uint32_t height);
+    void post(const std::shared_ptr<DisplayBufferInfo> &);
 
    private:
     bool canComposite(VkFormat);
@@ -45,23 +65,13 @@
     VkSemaphore m_imageReadySem;
     VkSemaphore m_frameDrawCompleteSem;
 
-    struct ColorBufferInfo {
-        uint32_t m_width;
-        uint32_t m_height;
-        VkFormat m_vkFormat;
-        VkImageView m_vkImageView;
-    };
-    // This map won't automatically relaase all the created VkImageView. The
-    // client is responsible for calling releaseImport on all imported ids
-    // before this object is released.
-    std::unordered_map<HandleType, ColorBufferInfo> m_colorBuffers;
-
     std::unique_ptr<SwapChainStateVk> m_swapChainStateVk;
     std::unique_ptr<CompositorVk> m_compositorVk;
     struct SurfaceState {
         uint32_t m_width = 0;
         uint32_t m_height = 0;
-        std::optional<HandleType> m_prevColorBuffer = std::nullopt;
+        std::weak_ptr<DisplayBufferInfo> m_prevDisplayBuffer =
+            std::weak_ptr<DisplayBufferInfo>();
     };
     std::unique_ptr<SurfaceState> m_surfaceState;
     std::unordered_map<VkFormat, bool> m_canComposite;
diff --git a/stream-servers/FrameBuffer.cpp b/stream-servers/FrameBuffer.cpp
index 06b384d..d410cd5 100644
--- a/stream-servers/FrameBuffer.cpp
+++ b/stream-servers/FrameBuffer.cpp
@@ -59,7 +59,7 @@
 // to a FrameBuffer instance.
 class ColorBufferHelper : public ColorBuffer::Helper {
    public:
-    ColorBufferHelper(FrameBuffer* fb) : mFb(fb), mDisplayVk(nullptr) {}
+    ColorBufferHelper(FrameBuffer* fb) : mFb(fb) {}
 
     virtual bool setupContext() override {
         mIsBound = mFb->bind_locked();
@@ -77,17 +77,8 @@
 
     virtual bool isBound() const override { return mIsBound; }
 
-    virtual void releaseDisplayImport(HandleType id) override {
-        if (mDisplayVk) {
-            mDisplayVk->releaseImport(id);
-        }
-    }
-
-    void setDisplayVk(DisplayVk* displayVk) { mDisplayVk = displayVk; }
-
    private:
     FrameBuffer* mFb;
-    DisplayVk* mDisplayVk;
     bool mIsBound = false;
 };
 
@@ -377,8 +368,6 @@
                 *dispatch, emu->physdev, emu->queueFamilyIndex,
                 emu->queueFamilyIndex, emu->device, emu->queue, emu->queue);
             fb->m_vkInstance = emu->instance;
-            static_cast<ColorBufferHelper*>(fb->m_colorBufferHelper)
-                ->setDisplayVk(fb->m_displayVk.get());
         }
     }
 
@@ -703,12 +692,14 @@
     }
 
     auto& cb = *c->second.cb;
+    std::shared_ptr<DisplayVk::DisplayBufferInfo> db = nullptr;
     if (m_displayVk != nullptr) {
-        m_displayVk->importVkImage(colorBufferHandle, image, format,
-                                   static_cast<uint32_t>(cb.getWidth()),
-                                   static_cast<uint32_t>(cb.getHeight()));
+        db = m_displayVk->createDisplayBuffer(
+            image, format, static_cast<uint32_t>(cb.getWidth()),
+            static_cast<uint32_t>(cb.getHeight()));
     }
-    return cb.importMemory(handle, size, dedicated, linearTiling, vulkanOnly);
+    return cb.importMemory(handle, size, dedicated, linearTiling, vulkanOnly,
+                           std::move(db));
 }
 
 void FrameBuffer::setColorBufferInUse(
@@ -1232,6 +1223,7 @@
         // emugl::emugl_crash_reporter(
         //     "FATAL: color buffer with handle %u already exists",
         //     handle);
+        ::abort();
     }
 
     createColorBufferWithHandleLocked(
@@ -1270,7 +1262,8 @@
         // Explicitly set refcount to 1 to avoid the colorbuffer being added to
         // m_colorBufferDelayedCloseList in FrameBuffer::onLoad().
         if (m_refCountPipeEnabled) {
-            m_colorbuffers[handle] = { std::move(cb), 1, false, 0 };
+            m_colorbuffers.try_emplace(
+                handle, ColorBufferRef{std::move(cb), 1, false, 0});
         } else {
             // Android master default api level is 1000
             int apiLevel = 1000;
@@ -1278,7 +1271,9 @@
             // pre-O and post-O use different color buffer memory management
             // logic
             if (apiLevel > 0 && apiLevel < 26) {
-                m_colorbuffers[handle] = {std::move(cb), 1, false, 0};
+                m_colorbuffers.try_emplace(
+                    handle,
+                    ColorBufferRef{std::move(cb), 1, false, 0});
 
                 RenderThreadInfo* tInfo = RenderThreadInfo::get();
                 uint64_t puid = tInfo->m_puid;
@@ -1287,7 +1282,9 @@
                 }
 
             } else {
-                m_colorbuffers[handle] = {std::move(cb), 0, false, 0};
+                m_colorbuffers.try_emplace(
+                    handle,
+                    ColorBufferRef{std::move(cb), 0, false, 0});
             }
         }
     } else {
@@ -2246,7 +2243,11 @@
 }
 
 ColorBufferPtr FrameBuffer::getColorBuffer_locked(HandleType p_colorBuffer) {
-    return android::base::findOrDefault(m_colorbuffers, p_colorBuffer).cb;
+    auto i = m_colorbuffers.find(p_colorBuffer);
+    if (i == m_colorbuffers.end()) {
+        return nullptr;
+    }
+    return i->second.cb;
 }
 
 WindowSurfacePtr FrameBuffer::getWindowSurface_locked(HandleType p_windowsurface) {
@@ -2466,14 +2467,14 @@
     bool ret = false;
     ColorBufferMap::iterator c;
 
-    // TODO(kaiyili, b/179481815): make DisplayVk::post asynchronous.
-    if (m_displayVk != nullptr) {
-        m_displayVk->post(p_colorbuffer);
+    c = m_colorbuffers.find(p_colorbuffer);
+    if (c == m_colorbuffers.end()) {
         goto EXIT;
     }
 
-    c = m_colorbuffers.find(p_colorbuffer);
-    if (c == m_colorbuffers.end()) {
+    // TODO(kaiyili, b/179481815): make DisplayVk::post asynchronous.
+    if (m_displayVk != nullptr) {
+        m_displayVk->post(c->second.cb->getDisplayBufferVk());
         goto EXIT;
     }
 
@@ -3015,7 +3016,8 @@
         if (refCount == 0) {
             m_colorBufferDelayedCloseList.push_back({closedTs, handle});
         }
-        return { handle, { std::move(cb), refCount, opened, closedTs } };
+        return {handle, ColorBufferRef{std::move(cb), refCount, opened,
+                                       closedTs}};
     });
     m_lastPostedColorBuffer = static_cast<HandleType>(stream->getBe32());
     GL_LOG("Got lasted posted color buffer from snapshot");
diff --git a/stream-servers/FrameBuffer.h b/stream-servers/FrameBuffer.h
index dfc1c5e..6fc45a1 100644
--- a/stream-servers/FrameBuffer.h
+++ b/stream-servers/FrameBuffer.h
@@ -1,18 +1,18 @@
 /*
-* Copyright (C) 2011-2015 The Android Open Source Project
-*
-* Licensed under the Apache License, Version 2.0 (the "License");
-* you may not use this file except in compliance with the License.
-* You may obtain a copy of the License at
-*
-* http://www.apache.org/licenses/LICENSE-2.0
-*
-* Unless required by applicable law or agreed to in writing, software
-* distributed under the License is distributed on an "AS IS" BASIS,
-* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-* See the License for the specific language governing permissions and
-* limitations under the License.
-*/
+ * Copyright (C) 2011-2015 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 #ifndef _LIBRENDER_FRAMEBUFFER_H
 #define _LIBRENDER_FRAMEBUFFER_H
 
@@ -21,6 +21,7 @@
 
 #include <functional>
 #include <map>
+#include <memory>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -62,7 +63,8 @@
     BufferPtr buffer;
 };
 
-typedef std::unordered_map<HandleType, std::pair<WindowSurfacePtr, HandleType> > WindowSurfaceMap;
+typedef std::unordered_map<HandleType, std::pair<WindowSurfacePtr, HandleType>>
+    WindowSurfaceMap;
 typedef std::unordered_set<HandleType> WindowSurfaceSet;
 typedef std::unordered_map<uint64_t, WindowSurfaceSet> ProcOwnedWindowSurfaces;
 
@@ -193,7 +195,7 @@
     // |version| specifies the GLES version as a GLESApi enum.
     // Return a new handle value, which will be 0 in case of error.
     HandleType createRenderContext(int p_config, HandleType p_share,
-        GLESApi version = GLESApi_CM);
+                                   GLESApi version = GLESApi_CM);
 
     // Create a new WindowSurface instance from this display instance.
     // |p_config| is the index of one of the configs returned by getConfigs().
@@ -210,16 +212,16 @@
     // list of valid values. Note that ColorBuffer instances are reference-
     // counted. Use openColorBuffer / closeColorBuffer to operate on the
     // internal count.
-    HandleType createColorBuffer(
-        int p_width, int p_height, GLenum p_internalFormat,
-        FrameworkFormat p_frameworkFormat);
+    HandleType createColorBuffer(int p_width, int p_height,
+                                 GLenum p_internalFormat,
+                                 FrameworkFormat p_frameworkFormat);
     // Variant of createColorBuffer except with a particular
     // handle already assigned. This is for use with
     // virtio-gpu's RESOURCE_CREATE ioctl.
-    void createColorBufferWithHandle(
-        int p_width, int p_height, GLenum p_internalFormat,
-        FrameworkFormat p_frameworkFormat,
-        HandleType handle);
+    void createColorBufferWithHandle(int p_width, int p_height,
+                                     GLenum p_internalFormat,
+                                     FrameworkFormat p_frameworkFormat,
+                                     HandleType handle);
 
     // Create a new data Buffer instance from this display instance.
     // The buffer will be backed by a VkBuffer and VkDeviceMemory (if Vulkan
@@ -271,9 +273,8 @@
     // of the context, the draw surface and the read surface, respectively.
     // Returns true on success, false on failure.
     // Note: if all handle values are 0, this is an unbind operation.
-    bool  bindContext(HandleType p_context,
-                      HandleType p_drawSurface,
-                      HandleType p_readSurface);
+    bool bindContext(HandleType p_context, HandleType p_drawSurface,
+                     HandleType p_readSurface);
 
     // Return a render context pointer from its handle
     RenderContextPtr getContext_locked(HandleType p_context);
@@ -289,14 +290,14 @@
     // |p_surface| is the target WindowSurface's handle value.
     // |p_colorbuffer| is the ColorBuffer handle value.
     // Returns true on success, false otherwise.
-    bool  setWindowSurfaceColorBuffer(
-            HandleType p_surface, HandleType p_colorbuffer);
+    bool setWindowSurfaceColorBuffer(HandleType p_surface,
+                                     HandleType p_colorbuffer);
 
     // Copy the content of a WindowSurface's Pbuffer to its attached
     // ColorBuffer. See the documentation for WindowSurface::flushColorBuffer()
     // |p_surface| is the target WindowSurface's handle value.
     // Returns true on success, false on failure.
-    bool  flushWindowSurfaceColorBuffer(HandleType p_surface);
+    bool flushWindowSurfaceColorBuffer(HandleType p_surface);
 
     // Retrieves the color buffer handle associated with |p_surface|.
     // Returns 0 if there is no such handle.
@@ -307,15 +308,15 @@
     // glEGLImageTargetTexture2DOES() for all GLES versions.
     // |p_colorbuffer| is the ColorBuffer's handle value.
     // Returns true on success, false on failure.
-    bool  bindColorBufferToTexture(HandleType p_colorbuffer);
-    bool  bindColorBufferToTexture2(HandleType p_colorbuffer);
+    bool bindColorBufferToTexture(HandleType p_colorbuffer);
+    bool bindColorBufferToTexture2(HandleType p_colorbuffer);
 
     // Bind the current context's EGL_RENDERBUFFER_OES render buffer to this
     // ColorBuffer's EGLImage. This is intended to implement
     // glEGLImageTargetRenderbufferStorageOES() for all GLES versions.
     // |p_colorbuffer| is the ColorBuffer's handle value.
     // Returns true on success, false on failure.
-    bool  bindColorBufferToRenderbuffer(HandleType p_colorbuffer);
+    bool bindColorBufferToRenderbuffer(HandleType p_colorbuffer);
 
     // Read the content of a given ColorBuffer into client memory.
     // |p_colorbuffer| is the ColorBuffer's handle value. Similar
@@ -326,9 +327,8 @@
     // |type| is the type of pixel data, e.g. GL_UNSIGNED_BYTE.
     // |pixels| is the address of a caller-provided buffer that will be filled
     // with the pixel data.
-    void  readColorBuffer(HandleType p_colorbuffer,
-                           int x, int y, int width, int height,
-                           GLenum format, GLenum type, void *pixels);
+    void readColorBuffer(HandleType p_colorbuffer, int x, int y, int width,
+                         int height, GLenum format, GLenum type, void* pixels);
 
     // Read the content of a given YUV420_888 ColorBuffer into client memory.
     // |p_colorbuffer| is the ColorBuffer's handle value. Similar
@@ -338,29 +338,19 @@
     // |pixels| is the address of a caller-provided buffer that will be filled
     // with the pixel data.
     // |pixles_size| is the size of buffer
-    void  readColorBufferYUV(HandleType p_colorbuffer,
-                             int x, int y, int width, int height,
-                             void *pixels, uint32_t pixels_size);
+    void readColorBufferYUV(HandleType p_colorbuffer, int x, int y, int width,
+                            int height, void* pixels, uint32_t pixels_size);
 
     // create a Y texture and a UV texture with width and height, the created
     // texture ids are stored in textures respectively
-    void createYUVTextures(uint32_t type,
-                           uint32_t count,
-                           int width,
-                           int height,
+    void createYUVTextures(uint32_t type, uint32_t count, int width, int height,
                            uint32_t* output);
     void destroyYUVTextures(uint32_t type, uint32_t count, uint32_t* textures);
-    void updateYUVTextures(uint32_t type,
-                           uint32_t* textures,
-                           void* privData,
+    void updateYUVTextures(uint32_t type, uint32_t* textures, void* privData,
                            void* func);
-    void swapTexturesAndUpdateColorBuffer(uint32_t colorbufferhandle,
-                                          int x,
-                                          int y,
-                                          int width,
-                                          int height,
-                                          uint32_t format,
-                                          uint32_t type,
+    void swapTexturesAndUpdateColorBuffer(uint32_t colorbufferhandle, int x,
+                                          int y, int width, int height,
+                                          uint32_t format, uint32_t type,
                                           uint32_t texture_type,
                                           uint32_t* textures);
 
@@ -369,23 +359,21 @@
     // to glReadPixels(), this can be a slow operation.
     // |x|, |y|, |width| and |height| are the position and dimensions of
     // a rectangle whose pixel values will be transfered to the GPU
-    // |format| indicates the format of the OpenGL buffer, e.g. GL_RGB or GL_RGBA.
-    // |frameworkFormat| indicates the format of the pixel data; if
+    // |format| indicates the format of the OpenGL buffer, e.g. GL_RGB or
+    // GL_RGBA. |frameworkFormat| indicates the format of the pixel data; if
     // FRAMEWORK_FORMAT_GL_COMPATIBLE, |format| (OpenGL format) is used.
     // Otherwise, explicit conversion to |format| is needed.
     // |type| is the type of pixel data, e.g. GL_UNSIGNED_BYTE.
     // |pixels| is the address of a buffer containing the new pixel data.
     // Returns true on success, false otherwise.
-    bool updateColorBuffer(HandleType p_colorbuffer,
-                           int x, int y, int width, int height,
-                           GLenum format, GLenum type, void *pixels);
+    bool updateColorBuffer(HandleType p_colorbuffer, int x, int y, int width,
+                           int height, GLenum format, GLenum type,
+                           void* pixels);
     // Replaces contents completely using the color buffer's current format,
     // with row length equal to width of a row in bytes.
     // The number of bytes is passed as a check.
-    bool replaceColorBufferContents(
-        HandleType p_colorbuffer,
-        const void *pixels,
-        size_t numBytes);
+    bool replaceColorBufferContents(HandleType p_colorbuffer,
+                                    const void* pixels, size_t numBytes);
     // Reads back the raw color buffer to |pixels|
     // if |pixels| is not null.
     // Always returns in |numBytes| how many bytes were
@@ -394,14 +382,10 @@
     // fewer or more bytes cannot be specified.
     // If the framework format is YUV, it will read
     // back as raw YUV data.
-    bool readColorBufferContents(
-        HandleType p_colorbuffer,
-        size_t* numBytes,
-        void *pixels);
+    bool readColorBufferContents(HandleType p_colorbuffer, size_t* numBytes,
+                                 void* pixels);
 
-    bool getColorBufferInfo(HandleType p_colorbuffer,
-                            int* width,
-                            int* height,
+    bool getColorBufferInfo(HandleType p_colorbuffer, int* width, int* height,
                             GLint* internalformat,
                             FrameworkFormat* frameworkFormat = nullptr);
     bool getBufferInfo(HandleType p_buffer, int* size);
@@ -469,7 +453,8 @@
 
     // Create an eglImage and return its handle.  Reference:
     // https://www.khronos.org/registry/egl/extensions/KHR/EGL_KHR_image_base.txt
-    HandleType createClientImage(HandleType context, EGLenum target, GLuint buffer);
+    HandleType createClientImage(HandleType context, EGLenum target,
+                                 GLuint buffer);
     // Call the implementation of eglDestroyImageKHR, return if succeeds or
     // not. Reference:
     // https://www.khronos.org/registry/egl/extensions/KHR/EGL_KHR_image_base.txt
@@ -487,8 +472,7 @@
     // GL state.
     // It can be unsafe / leaky to change the structure of contexts
     // outside the facilities the FrameBuffer class provides.
-    void createTrivialContext(HandleType shared,
-                              HandleType* contextOut,
+    void createTrivialContext(HandleType shared, HandleType* contextOut,
                               HandleType* surfOut);
     // createAndBindTrivialSharedContext(), but with a m_pbufContext
     // as shared, and not adding itself to the context map at all.
@@ -543,35 +527,29 @@
     // nChannels should be 3 (RGB) or 4 (RGBA).
     // Note: swiftshader_indirect does not work with 3 channels
     void getScreenshot(unsigned int nChannels, unsigned int* width,
-            unsigned int* height, std::vector<unsigned char>& pixels,
-            int displayId, int desiredWidth, int desiredHeight,
-            int desiredRotation);
+                       unsigned int* height, std::vector<unsigned char>& pixels,
+                       int displayId, int desiredWidth, int desiredHeight,
+                       int desiredRotation);
     void onLastColorBufferRef(uint32_t handle);
     ColorBuffer::Helper* getColorBufferHelper() { return m_colorBufferHelper; }
     ColorBufferPtr findColorBuffer(HandleType p_colorbuffer);
 
-    void registerProcessCleanupCallback(void* key, std::function<void()> callback);
+    void registerProcessCleanupCallback(void* key,
+                                        std::function<void()> callback);
     void unregisterProcessCleanupCallback(void* key);
 
     void registerProcessSequenceNumberForPuid(uint64_t puid);
     uint32_t* getProcessSequenceNumberPtr(uint64_t puid);
 
-    int createDisplay(uint32_t *displayId);
+    int createDisplay(uint32_t* displayId);
     int destroyDisplay(uint32_t displayId);
     int setDisplayColorBuffer(uint32_t displayId, uint32_t colorBuffer);
     int getDisplayColorBuffer(uint32_t displayId, uint32_t* colorBuffer);
     int getColorBufferDisplay(uint32_t colorBuffer, uint32_t* displayId);
-    int getDisplayPose(uint32_t displayId,
-                       int32_t* x,
-                       int32_t* y,
-                       uint32_t* w,
+    int getDisplayPose(uint32_t displayId, int32_t* x, int32_t* y, uint32_t* w,
                        uint32_t* h);
-    int setDisplayPose(uint32_t displayId,
-                       int32_t x,
-                       int32_t y,
-                       uint32_t w,
-                       uint32_t h,
-                       uint32_t dpi = 0);
+    int setDisplayPose(uint32_t displayId, int32_t x, int32_t y, uint32_t w,
+                       uint32_t h, uint32_t dpi = 0);
     void getCombinedDisplaySize(int* w, int* h);
     struct DisplayInfo {
         uint32_t cb;
@@ -580,8 +558,10 @@
         uint32_t width;
         uint32_t height;
         uint32_t dpi;
-        DisplayInfo() : cb(0), pos_x(0), pos_y(0), width(0), height(0), dpi(0) {};
-        DisplayInfo(uint32_t cb, int32_t x, int32_t y, uint32_t w, uint32_t h, uint32_t d)
+        DisplayInfo()
+            : cb(0), pos_x(0), pos_y(0), width(0), height(0), dpi(0){};
+        DisplayInfo(uint32_t cb, int32_t x, int32_t y, uint32_t w, uint32_t h,
+                    uint32_t d)
             : cb(cb), pos_x(x), pos_y(y), width(w), height(h), dpi(d) {}
     };
     // Inline with MultiDisplay::s_invalidIdMultiDisplay
@@ -605,7 +585,8 @@
     bool bindFakeWindow_locked();
     bool removeSubWindow_locked();
     // Returns the set of ColorBuffers destroyed (for further cleanup)
-    std::vector<HandleType> cleanupProcGLObjects_locked(uint64_t puid, bool forced = false);
+    std::vector<HandleType> cleanupProcGLObjects_locked(uint64_t puid,
+                                                        bool forced = false);
 
     void markOpened(ColorBufferRef* cbRef);
     // Returns true if the color buffer was erased.
@@ -615,21 +596,17 @@
     // Close all expired color buffers for real.
     // Treat all delayed color buffers as expired if forced=true
     void performDelayedColorBufferCloseLocked(bool forced = false);
-    void eraseDelayedCloseColorBufferLocked(
-            HandleType cb, uint64_t ts);
+    void eraseDelayedCloseColorBufferLocked(HandleType cb, uint64_t ts);
 
-    bool postImpl(HandleType p_colorbuffer, bool needLockAndBind = true, bool repaint = false);
+    bool postImpl(HandleType p_colorbuffer, bool needLockAndBind = true,
+                  bool repaint = false);
     void setGuestPostedAFrame() { m_guestPostedAFrame = true; }
-    HandleType createColorBufferLocked(int p_width,
-                                       int p_height,
+    HandleType createColorBufferLocked(int p_width, int p_height,
                                        GLenum p_internalFormat,
                                        FrameworkFormat p_frameworkFormat);
     HandleType createColorBufferWithHandleLocked(
-        int p_width,
-        int p_height,
-        GLenum p_internalFormat,
-        FrameworkFormat p_frameworkFormat,
-        HandleType handle);
+        int p_width, int p_height, GLenum p_internalFormat,
+        FrameworkFormat p_frameworkFormat, HandleType handle);
     HandleType createBufferLocked(int p_size);
     HandleType createBufferWithHandleLocked(int p_size, HandleType handle);
 
@@ -637,8 +614,8 @@
     void setDisplayPoseInSkinUI(int totalHeight);
     void sweepColorBuffersLocked();
 
-private:
-    static FrameBuffer *s_theFrameBuffer;
+   private:
+    static FrameBuffer* s_theFrameBuffer;
     static HandleType s_nextHandle;
     int m_x = 0;
     int m_y = 0;
@@ -680,8 +657,8 @@
     //
     // Note: this collection is ordered by |ts| field.
     struct ColorBufferCloseInfo {
-        uint64_t ts; // when we got the close request.
-        HandleType cbHandle;    // 0 == already closed, do nothing
+        uint64_t ts;          // when we got the close request.
+        HandleType cbHandle;  // 0 == already closed, do nothing
     };
     using ColorBufferDelayedClose = std::vector<ColorBufferCloseInfo>;
     ColorBufferDelayedClose m_colorBufferDelayedCloseList;
@@ -701,11 +678,11 @@
     EGLSurface m_prevDrawSurf = EGL_NO_SURFACE;
     EGLNativeWindowType m_subWin = {};
     TextureDraw* m_textureDraw = nullptr;
-    EGLConfig  m_eglConfig = nullptr;
+    EGLConfig m_eglConfig = nullptr;
     HandleType m_lastPostedColorBuffer = 0;
-    float      m_zRot = 0;
-    float      m_px = 0;
-    float      m_py = 0;
+    float m_zRot = 0;
+    float m_px = 0;
+    float m_py = 0;
 
     // Async readback
     enum class ReadbackCmd {
@@ -724,7 +701,8 @@
         uint32_t width;
         uint32_t height;
     };
-    android::base::WorkerProcessingResult sendReadbackWorkerCmd(const Readback& readback);
+    android::base::WorkerProcessingResult sendReadbackWorkerCmd(
+        const Readback& readback);
     bool m_asyncReadbackSupported = true;
     bool m_guestPostedAFrame = false;
 
diff --git a/stream-servers/tests/DisplayVk_unittest.cpp b/stream-servers/tests/DisplayVk_unittest.cpp
index 28ff2b9..fee1747 100644
--- a/stream-servers/tests/DisplayVk_unittest.cpp
+++ b/stream-servers/tests/DisplayVk_unittest.cpp
@@ -227,10 +227,10 @@
                                          textureWidth, textureHeight);
     std::vector<uint32_t> pixels(textureWidth * textureHeight, 0);
     ASSERT_TRUE(texture->write(pixels));
-    displayVk.importVkImage(0, texture->m_vkImage, RenderTexture::k_vkFormat,
-                            textureWidth, textureHeight);
-    displayVk.post(0);
-    displayVk.releaseImport(0);
+    auto cbvk = displayVk.createDisplayBuffer(texture->m_vkImage,
+                                              RenderTexture::k_vkFormat,
+                                              textureWidth, textureHeight);
+    displayVk.post(cbvk);
 }
 
 TEST_F(DisplayVkTest, SimplePost) {
@@ -251,12 +251,12 @@
         }
     }
     ASSERT_TRUE(texture->write(pixels));
-    m_displayVk->importVkImage(0, texture->m_vkImage, texture->k_vkFormat,
-                               texture->m_width, texture->m_height);
+    auto cbvk = m_displayVk->createDisplayBuffer(
+        texture->m_vkImage, texture->k_vkFormat, texture->m_width,
+        texture->m_height);
     for (uint32_t i = 0; i < 10; i++) {
-        m_displayVk->post(0);
+        m_displayVk->post(cbvk);
     }
-    m_displayVk->releaseImport(0);
 }
 
 TEST_F(DisplayVkTest, PostTwoColorBuffers) {
@@ -274,15 +274,14 @@
     std::vector<uint32_t> greenPixels(textureWidth * textureHeight, green);
     ASSERT_TRUE(redTexture->write(redPixels));
     ASSERT_TRUE(greenTexture->write(greenPixels));
-    m_displayVk->importVkImage(0, redTexture->m_vkImage, redTexture->k_vkFormat,
-                               redTexture->m_width, redTexture->m_height);
-    m_displayVk->importVkImage(1, greenTexture->m_vkImage,
-                               greenTexture->k_vkFormat, greenTexture->m_width,
-                               greenTexture->m_height);
+    auto redCbvk = m_displayVk->createDisplayBuffer(
+        redTexture->m_vkImage, redTexture->k_vkFormat, redTexture->m_width,
+        redTexture->m_height);
+    auto greenCbvk = m_displayVk->createDisplayBuffer(
+        greenTexture->m_vkImage, greenTexture->k_vkFormat,
+        greenTexture->m_width, greenTexture->m_height);
     for (uint32_t i = 0; i < 10; i++) {
-        m_displayVk->post(0);
-        m_displayVk->post(1);
+        m_displayVk->post(redCbvk);
+        m_displayVk->post(greenCbvk);
     }
-    m_displayVk->releaseImport(0);
-    m_displayVk->releaseImport(1);
 }
\ No newline at end of file