[vulkan] Suballocate host visible memory

bug: 111137294
bug: 121420031

Change-Id: Ied783769ac1a35ed16dfb60e44857daf22d17233
diff --git a/android-emu/android/base/SubAllocator.cpp b/android-emu/android/base/SubAllocator.cpp
index 6c3085f..0d835d6 100644
--- a/android-emu/android/base/SubAllocator.cpp
+++ b/android-emu/android/base/SubAllocator.cpp
@@ -131,5 +131,9 @@
     mImpl->freeAll();
 }
 
+uint64_t SubAllocator::getOffset(void* ptr) {
+    return mImpl->getOffset(ptr);
+}
+
 } // namespace base
 } // namespace android
diff --git a/android-emu/android/base/SubAllocator.h b/android-emu/android/base/SubAllocator.h
index 894b205..b2363bd 100644
--- a/android-emu/android/base/SubAllocator.h
+++ b/android-emu/android/base/SubAllocator.h
@@ -40,6 +40,7 @@
     void* alloc(size_t wantedSize);
     void free(void* ptr);
     void freeAll();
+    uint64_t getOffset(void* ptr);
 
     // Convenience function to allocate an array
     // of objects of type T.
diff --git a/system/vulkan_enc/HostVisibleMemoryVirtualization.cpp b/system/vulkan_enc/HostVisibleMemoryVirtualization.cpp
index 88fbc6e..c1495fb 100644
--- a/system/vulkan_enc/HostVisibleMemoryVirtualization.cpp
+++ b/system/vulkan_enc/HostVisibleMemoryVirtualization.cpp
@@ -14,10 +14,17 @@
 // limitations under the License.
 #include "HostVisibleMemoryVirtualization.h"
 
+#include "android/base/SubAllocator.h"
+
+#include "Resources.h"
+#include "VkEncoder.h"
+
 #include <log/log.h>
 
 #include <set>
 
+using android::base::SubAllocator;
+
 namespace goldfish_vk {
 
 bool canFitVirtualHostVisibleMemoryInfo(
@@ -40,21 +47,7 @@
     }
 
     uint32_t numFreeMemoryTypes = VK_MAX_MEMORY_TYPES - typeCount;
-    uint32_t numFreeMemoryHeaps = VK_MAX_MEMORY_HEAPS - heapCount;
-
     uint32_t hostVisibleMemoryTypeCount = 0;
-    uint32_t hostVisibleMemoryHeapCount = 0;
-    std::set<uint32_t> hostVisibleMemoryHeaps;
-
-    for (uint32_t i = 0; i < typeCount; ++i) {
-        const auto& type = memoryProperties->memoryTypes[i];
-        if (type.propertyFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) {
-            ++hostVisibleMemoryTypeCount;
-            hostVisibleMemoryHeaps.insert(type.heapIndex);
-        }
-    }
-    hostVisibleMemoryHeapCount =
-        (uint32_t)hostVisibleMemoryHeaps.size();
 
     if (hostVisibleMemoryTypeCount > numFreeMemoryTypes) {
         ALOGE("Underlying device has too many host visible memory types (%u)"
@@ -63,13 +56,6 @@
         canFit = false;
     }
 
-    if (hostVisibleMemoryHeapCount > numFreeMemoryHeaps) {
-        ALOGE("Underlying device has too many host visible memory types (%u)"
-              "and not enough free types (%u)",
-              hostVisibleMemoryHeapCount, numFreeMemoryHeaps);
-        canFit = false;
-    }
-
     return canFit;
 }
 
@@ -161,7 +147,7 @@
                 ~(VK_MEMORY_HEAP_DEVICE_LOCAL_BIT);
 
             // TODO: Figure out how to support bigger sizes
-            newVirtualMemoryHeap.size = 512ULL * 1048576ULL; // 512 MB
+            newVirtualMemoryHeap.size = VIRTUAL_HOST_VISIBLE_HEAP_SIZE;
 
             info_out->memoryTypeIndexMappingToHost[firstFreeTypeIndex] = i;
             info_out->memoryHeapIndexMappingToHost[firstFreeHeapIndex] = i;
@@ -175,12 +161,14 @@
                 type.propertyFlags & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
 
             ++firstFreeTypeIndex;
-            ++firstFreeHeapIndex;
+
+            // Explicitly only create one new heap.
+            // ++firstFreeHeapIndex;
         }
     }
 
     info_out->guestMemoryProperties.memoryTypeCount = firstFreeTypeIndex;
-    info_out->guestMemoryProperties.memoryHeapCount = firstFreeHeapIndex;
+    info_out->guestMemoryProperties.memoryHeapCount = firstFreeHeapIndex + 1;
 
     for (uint32_t i = info_out->guestMemoryProperties.memoryTypeCount; i < VK_MAX_MEMORY_TYPES; ++i) {
         memset(&info_out->guestMemoryProperties.memoryTypes[i],
@@ -188,4 +176,81 @@
     }
 }
 
+bool isHostVisibleMemoryTypeIndexForGuest(
+    const HostVisibleMemoryVirtualizationInfo* info,
+    uint32_t index) {
+    return info->guestMemoryProperties.memoryTypes[index].propertyFlags &
+           VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
+}
+
+VkResult finishHostMemAllocInit(
+    VkEncoder* enc,
+    VkDevice device,
+    uint32_t memoryTypeIndex,
+    VkDeviceSize nonCoherentAtomSize,
+    VkDeviceSize allocSize,
+    VkDeviceSize mappedSize,
+    uint8_t* mappedPtr,
+    HostMemAlloc* out) {
+
+    out->enc = enc;
+    out->device = device;
+    out->memoryTypeIndex = memoryTypeIndex;
+    out->nonCoherentAtomSize = nonCoherentAtomSize;
+    out->allocSize = allocSize;
+    out->mappedSize = mappedSize;
+    out->mappedPtr = mappedPtr;
+
+    out->subAlloc = new
+        SubAllocator(
+            out->mappedPtr,
+            out->mappedSize,
+            out->nonCoherentAtomSize);
+
+    out->initialized = true;
+    out->initResult = VK_SUCCESS;
+    return VK_SUCCESS;
+}
+
+void destroyHostMemAlloc(
+    VkDevice device,
+    HostMemAlloc* toDestroy) {
+
+    if (toDestroy->initResult != VK_SUCCESS) return;
+    if (!toDestroy->initialized) return;
+
+    toDestroy->enc->vkFreeMemory(device, toDestroy->memory, nullptr);
+    delete toDestroy->subAlloc;
+}
+
+void subAllocHostMemory(
+    HostMemAlloc* alloc,
+    const VkMemoryAllocateInfo* pAllocateInfo,
+    SubAlloc* out) {
+
+    VkDeviceSize mappedSize =
+        alloc->nonCoherentAtomSize * (
+            (pAllocateInfo->allocationSize +
+             alloc->nonCoherentAtomSize - 1) /
+            alloc->nonCoherentAtomSize);
+
+    void* subMapped = alloc->subAlloc->alloc(mappedSize);
+    out->mappedPtr = (uint8_t*)subMapped;
+
+    out->subAllocSize = pAllocateInfo->allocationSize;
+    out->subMappedSize = mappedSize;
+
+    out->baseMemory = alloc->memory;
+    out->baseOffset = alloc->subAlloc->getOffset(subMapped);
+
+    out->subMemory = new_from_host_VkDeviceMemory(VK_NULL_HANDLE);
+    out->subAlloc = alloc->subAlloc;
+}
+
+void subFreeHostMemory(SubAlloc* toFree) {
+    delete_goldfish_VkDeviceMemory(toFree->subMemory);
+    toFree->subAlloc->free(toFree->mappedPtr);
+    memset(toFree, 0x0, sizeof(SubAlloc));
+}
+
 } // namespace goldfish_vk
diff --git a/system/vulkan_enc/HostVisibleMemoryVirtualization.h b/system/vulkan_enc/HostVisibleMemoryVirtualization.h
index b1a7b87..a2396f8 100644
--- a/system/vulkan_enc/HostVisibleMemoryVirtualization.h
+++ b/system/vulkan_enc/HostVisibleMemoryVirtualization.h
@@ -16,8 +16,20 @@
 
 #include <vulkan/vulkan.h>
 
+#define VIRTUAL_HOST_VISIBLE_HEAP_SIZE 128ULL * (1048576ULL)
+
+namespace android {
+namespace base {
+
+class SubAllocator;
+
+} // namespace base
+} // namespace android
+
 namespace goldfish_vk {
 
+class VkEncoder;
+
 struct HostVisibleMemoryVirtualizationInfo {
     bool initialized = false;
     bool memoryPropertiesSupported;
@@ -47,4 +59,54 @@
     bool directMemSupported,
     HostVisibleMemoryVirtualizationInfo* info_out);
 
-} // namespace goldfish_vk
\ No newline at end of file
+bool isHostVisibleMemoryTypeIndexForGuest(
+    const HostVisibleMemoryVirtualizationInfo* info,
+    uint32_t index);
+
+struct HostMemAlloc {
+    bool initialized = false;
+    VkEncoder* enc;
+    VkResult initResult = VK_SUCCESS;
+    VkDevice device = nullptr;
+    uint32_t memoryTypeIndex = 0;
+    VkDeviceSize nonCoherentAtomSize = 0;
+    VkDeviceMemory memory = VK_NULL_HANDLE;
+    VkDeviceSize allocSize = 0;
+    VkDeviceSize mappedSize = 0;
+    uint8_t* mappedPtr = nullptr;
+    android::base::SubAllocator* subAlloc = nullptr;
+};
+
+VkResult finishHostMemAllocInit(
+    VkEncoder* enc,
+    VkDevice device,
+    uint32_t memoryTypeIndex,
+    VkDeviceSize nonCoherentAtomSize,
+    VkDeviceSize allocSize,
+    VkDeviceSize mappedSize,
+    uint8_t* mappedPtr,
+    HostMemAlloc* out);
+
+void destroyHostMemAlloc(
+    VkDevice device,
+    HostMemAlloc* toDestroy);
+
+struct SubAlloc {
+    uint8_t* mappedPtr = nullptr;
+    VkDeviceSize subAllocSize = 0;
+    VkDeviceSize subMappedSize = 0;
+
+    VkDeviceMemory baseMemory = VK_NULL_HANDLE;
+    VkDeviceSize baseOffset = 0;
+    android::base::SubAllocator* subAlloc = nullptr;
+    VkDeviceMemory subMemory = VK_NULL_HANDLE;
+};
+
+void subAllocHostMemory(
+    HostMemAlloc* alloc,
+    const VkMemoryAllocateInfo* pAllocateInfo,
+    SubAlloc* out);
+
+void subFreeHostMemory(SubAlloc* toFree);
+
+} // namespace goldfish_vk
diff --git a/system/vulkan_enc/ResourceTracker.cpp b/system/vulkan_enc/ResourceTracker.cpp
index aac0fc6..23fa666 100644
--- a/system/vulkan_enc/ResourceTracker.cpp
+++ b/system/vulkan_enc/ResourceTracker.cpp
@@ -115,6 +115,7 @@
         VkPhysicalDevice physdev;
         VkPhysicalDeviceProperties props;
         VkPhysicalDeviceMemoryProperties memProps;
+        HostMemAlloc hostMemAllocs[VK_MAX_MEMORY_TYPES] = {};
     };
 
     struct VkDeviceMemory_Info {
@@ -122,9 +123,11 @@
         VkDeviceSize mappedSize = 0;
         uint8_t* mappedPtr = nullptr;
         uint32_t memoryTypeIndex = 0;
+        bool virtualHostVisibleBacking = false;
         bool directMapped = false;
-        std::unique_ptr<GoldfishAddressSpaceBlock>
-            goldfishAddressSpaceBlock = {};
+        GoldfishAddressSpaceBlock*
+            goldfishAddressSpaceBlock = nullptr;
+        SubAlloc subAlloc;
     };
 
 #define HANDLE_REGISTER_IMPL_IMPL(type) \
@@ -145,7 +148,12 @@
 
     void unregister_VkDevice(VkDevice device) {
         AutoLock lock(mLock);
+
+        auto it = info_VkDevice.find(device);
+        if (it == info_VkDevice.end()) return;
+        auto info = it->second;
         info_VkDevice.erase(device);
+        lock.unlock();
     }
 
     void unregister_VkDeviceMemory(VkDeviceMemory mem) {
@@ -156,12 +164,18 @@
 
         auto& memInfo = it->second;
 
-        if (memInfo.mappedPtr && !memInfo.directMapped) {
+        if (memInfo.mappedPtr &&
+            !memInfo.virtualHostVisibleBacking &&
+            !memInfo.directMapped) {
             aligned_buf_free(memInfo.mappedPtr);
         }
 
-        // Direct mapping is erased by GoldfishAddressSpaceBlock's
-        // dtor
+        if (memInfo.directMapped) {
+            subFreeHostMemory(&memInfo.subAlloc);
+        }
+
+        delete memInfo.goldfishAddressSpaceBlock;
+
         info_VkDeviceMemory.erase(mem);
     }
 
@@ -281,16 +295,42 @@
         (void)memoryCount;
         (void)offsetCount;
         (void)sizeCount;
+
         const auto& hostVirt =
             mHostVisibleMemoryVirtInfo;
 
         if (!hostVirt.virtualizationSupported) return;
 
-        for (uint32_t i = 0; i < memoryCount; ++i) {
-            // TODO
-            (void)memory;
-            (void)offset;
-            (void)size;
+        if (memory) {
+            AutoLock lock (mLock);
+
+            for (uint32_t i = 0; i < memoryCount; ++i) {
+                VkDeviceMemory mem = memory[i];
+
+                auto it = info_VkDeviceMemory.find(mem);
+                if (it == info_VkDeviceMemory.end()) return;
+
+                const auto& info = it->second;
+
+                if (!info.directMapped) continue;
+
+                memory[i] = info.subAlloc.baseMemory;
+
+                if (offset) {
+                    offset[i] = info.subAlloc.baseOffset + offset[i];
+                }
+
+                if (size) {
+                    if (size[i] == VK_WHOLE_SIZE) {
+                        size[i] = info.subAlloc.subMappedSize;
+                    }
+                }
+
+                // TODO
+                (void)memory;
+                (void)offset;
+                (void)size;
+            }
         }
 
         for (uint32_t i = 0; i < typeIndexCount; ++i) {
@@ -320,12 +360,14 @@
         (void)memoryCount;
         (void)offsetCount;
         (void)sizeCount;
-        
+
         const auto& hostVirt =
             mHostVisibleMemoryVirtInfo;
 
         if (!hostVirt.virtualizationSupported) return;
 
+        AutoLock lock (mLock);
+
         for (uint32_t i = 0; i < memoryCount; ++i) {
             // TODO
             (void)memory;
@@ -386,13 +428,13 @@
         void*,
         VkPhysicalDevice physdev,
         VkPhysicalDeviceMemoryProperties* out) {
-        
+
         initHostVisibleMemoryVirtualizationInfo(
             physdev,
             out,
             mFeatureInfo->hasDirectMem,
             &mHostVisibleMemoryVirtInfo);
-        
+
         if (mHostVisibleMemoryVirtInfo.virtualizationSupported) {
             *out = mHostVisibleMemoryVirtInfo.guestMemoryProperties;
         }
@@ -420,6 +462,24 @@
         return input_result;
     }
 
+    void on_vkDestroyDevice_pre(
+        void*,
+        VkDevice device,
+        const VkAllocationCallbacks*) {
+
+        AutoLock lock(mLock);
+
+        auto it = info_VkDevice.find(device);
+        if (it == info_VkDevice.end()) return;
+        auto info = it->second;
+
+        lock.unlock();
+
+        for (uint32_t i = 0; i < VK_MAX_MEMORY_TYPES; ++i) {
+            destroyHostMemAlloc(device, &info.hostMemAllocs[i]);
+        }
+    }
+
     VkResult on_vkAllocateMemory(
         void* context,
         VkResult input_result,
@@ -432,62 +492,144 @@
 
         VkEncoder* enc = (VkEncoder*)context;
 
-        input_result =
-            enc->vkAllocateMemory(device, pAllocateInfo, pAllocator, pMemory);
-        
-        if (input_result != VK_SUCCESS) return input_result;
+        // Device local memory: pass through
+        if (!isHostVisibleMemoryTypeIndexForGuest(
+                &mHostVisibleMemoryVirtInfo,
+                pAllocateInfo->memoryTypeIndex)) {
 
-        VkDeviceSize allocationSize = pAllocateInfo->allocationSize;
-        VkDeviceSize mappedSize = getNonCoherentExtendedSize(device, allocationSize);
-        uint8_t* mappedPtr = nullptr;
-        bool hostVisible =
-            isMemoryTypeHostVisible(device, pAllocateInfo->memoryTypeIndex);
+            input_result =
+                enc->vkAllocateMemory(
+                    device, pAllocateInfo, pAllocator, pMemory);
+
+            if (input_result != VK_SUCCESS) return input_result;
+
+            VkDeviceSize allocationSize = pAllocateInfo->allocationSize;
+            setDeviceMemoryInfo(
+                device, *pMemory,
+                pAllocateInfo->allocationSize,
+                0, nullptr,
+                pAllocateInfo->memoryTypeIndex);
+
+            return VK_SUCCESS;
+        }
+
+        // Host visible memory with no direct mapping support
         bool directMappingSupported = usingDirectMapping();
-        if (hostVisible && !directMappingSupported) {
-            mappedPtr = (uint8_t*)aligned_buf_alloc(4096, mappedSize);
+        if (!directMappingSupported) {
+            VkDeviceSize mappedSize =
+                getNonCoherentExtendedSize(device,
+                    pAllocateInfo->allocationSize);
+            uint8_t* mappedPtr = (uint8_t*)aligned_buf_alloc(4096, mappedSize);
             D("host visible alloc (non-direct): "
               "size 0x%llx host ptr %p mapped size 0x%llx",
               (unsigned long long)allocationSize, mappedPtr,
               (unsigned long long)mappedSize);
+            setDeviceMemoryInfo(
+                device, *pMemory,
+                pAllocateInfo->allocationSize,
+                mappedSize, mappedPtr,
+                pAllocateInfo->memoryTypeIndex);
+            return VK_SUCCESS;
         }
 
-        setDeviceMemoryInfo(
-            device, *pMemory, allocationSize, mappedSize, mappedPtr,
-            pAllocateInfo->memoryTypeIndex);
+        // Host visible memory with direct mapping
+        AutoLock lock(mLock);
 
-        bool doDirectMap =
-            hostVisible && directMappingSupported;
+        auto it = info_VkDevice.find(device);
+        if (it == info_VkDevice.end()) return VK_ERROR_DEVICE_LOST;
+        auto& deviceInfo = it->second;
 
-        if (doDirectMap) {
+        HostMemAlloc* hostMemAlloc =
+            &deviceInfo.hostMemAllocs[pAllocateInfo->memoryTypeIndex];
+
+        if (!hostMemAlloc->initialized) {
+            VkMemoryAllocateInfo allocInfoForHost = *pAllocateInfo;
+            allocInfoForHost.allocationSize = VIRTUAL_HOST_VISIBLE_HEAP_SIZE;
+            // TODO: Support dedicated allocation
+            allocInfoForHost.pNext = nullptr;
+
+            lock.unlock();
+            VkResult host_res =
+                enc->vkAllocateMemory(
+                    device,
+                    &allocInfoForHost,
+                    nullptr,
+                    &hostMemAlloc->memory);
+            lock.lock();
+
+            if (host_res != VK_SUCCESS) {
+                ALOGE("Could not allocate backing for virtual host visible memory: %d",
+                      host_res);
+                hostMemAlloc->initialized = true;
+                hostMemAlloc->initResult = host_res;
+                return host_res;
+            }
+
+            auto& hostMemInfo = info_VkDeviceMemory[hostMemAlloc->memory];
+            hostMemInfo.allocationSize = allocInfoForHost.allocationSize;
+            VkDeviceSize nonCoherentAtomSize =
+                deviceInfo.props.limits.nonCoherentAtomSize;
+            hostMemInfo.mappedSize = hostMemInfo.allocationSize;
+            hostMemInfo.memoryTypeIndex =
+                pAllocateInfo->memoryTypeIndex;
 
             uint64_t directMappedAddr = 0;
-
+            lock.unlock();
             VkResult directMapResult =
                 enc->vkMapMemoryIntoAddressSpaceGOOGLE(
-                    device, *pMemory, &directMappedAddr);
+                    device, hostMemAlloc->memory, &directMappedAddr);
+            lock.lock();
 
             if (directMapResult != VK_SUCCESS) {
+                hostMemAlloc->initialized = true;
+                hostMemAlloc->initResult = directMapResult;
                 return directMapResult;
             }
 
-            AutoLock lock(mLock);
+            hostMemInfo.mappedPtr =
+                (uint8_t*)(uintptr_t)directMappedAddr;
+            hostMemInfo.virtualHostVisibleBacking = true;
 
-            auto it = info_VkDeviceMemory.find(*pMemory);
-            if (it == info_VkDeviceMemory.end()) {
-                return VK_ERROR_INITIALIZATION_FAILED;
+            VkResult hostMemAllocRes =
+                finishHostMemAllocInit(
+                    enc,
+                    device,
+                    pAllocateInfo->memoryTypeIndex,
+                    nonCoherentAtomSize,
+                    hostMemInfo.allocationSize,
+                    hostMemInfo.mappedSize,
+                    hostMemInfo.mappedPtr,
+                    hostMemAlloc);
+
+            if (hostMemAllocRes != VK_SUCCESS) {
+                return hostMemAllocRes;
             }
-
-            auto& info = it->second;
-            info.mappedPtr = (uint8_t*)(uintptr_t)directMappedAddr;
-            info.directMapped = true;
-
-            D("host visible alloc (direct): "
-              "size 0x%llx host ptr %p mapped size 0x%llx",
-              (unsigned long long)allocationSize, info.mappedPtr,
-              (unsigned long long)mappedSize);
         }
 
-        return input_result;
+        VkDeviceMemory_Info virtualMemInfo;
+
+        subAllocHostMemory(
+            hostMemAlloc,
+            pAllocateInfo,
+            &virtualMemInfo.subAlloc);
+
+        virtualMemInfo.allocationSize = virtualMemInfo.subAlloc.subAllocSize;
+        virtualMemInfo.mappedSize = virtualMemInfo.subAlloc.subMappedSize;
+        virtualMemInfo.mappedPtr = virtualMemInfo.subAlloc.mappedPtr;
+        virtualMemInfo.memoryTypeIndex = pAllocateInfo->memoryTypeIndex;
+        virtualMemInfo.directMapped = true;
+
+        D("host visible alloc (direct, suballoc): "
+          "size 0x%llx ptr %p mapped size 0x%llx",
+          (unsigned long long)virtualMemInfo.allocationSize, virtualMemInfo.mappedPtr,
+          (unsigned long long)virtualMemInfo.mappedSisze);
+
+        info_VkDeviceMemory[
+            virtualMemInfo.subAlloc.subMemory] = virtualMemInfo;
+
+        *pMemory = virtualMemInfo.subAlloc.subMemory;
+
+        return VK_SUCCESS;
     }
 
     void on_vkFreeMemory(
@@ -496,8 +638,20 @@
         VkDeviceMemory memory,
         const VkAllocationCallbacks* pAllocateInfo) {
 
-        VkEncoder* enc = (VkEncoder*)context;
-        enc->vkFreeMemory(device, memory, pAllocateInfo);
+        AutoLock lock(mLock);
+
+        auto it = info_VkDeviceMemory.find(memory);
+        if (it == info_VkDeviceMemory.end()) return;
+        auto& info = it->second;
+
+        if (!info.directMapped) {
+            lock.unlock();
+            VkEncoder* enc = (VkEncoder*)context;
+            enc->vkFreeMemory(device, memory, pAllocateInfo);
+            return;
+        }
+
+        subFreeHostMemory(&info.subAlloc);
     }
 
     VkResult on_vkMapMemory(
@@ -603,8 +757,8 @@
         }
 
         auto& memInfo = it->second;
-        memInfo.goldfishAddressSpaceBlock.reset(
-            new GoldfishAddressSpaceBlock);
+        memInfo.goldfishAddressSpaceBlock =
+            new GoldfishAddressSpaceBlock;
         auto& block = *(memInfo.goldfishAddressSpaceBlock);
 
         block.allocate(
@@ -796,6 +950,13 @@
         context, input_result, physicalDevice, pCreateInfo, pAllocator, pDevice);
 }
 
+void ResourceTracker::on_vkDestroyDevice_pre(
+    void* context,
+    VkDevice device,
+    const VkAllocationCallbacks* pAllocator) {
+    mImpl->on_vkDestroyDevice_pre(context, device, pAllocator);
+}
+
 VkResult ResourceTracker::on_vkAllocateMemory(
     void* context,
     VkResult input_result,
diff --git a/system/vulkan_enc/ResourceTracker.h b/system/vulkan_enc/ResourceTracker.h
index a36a5d5..6ad0ef1 100644
--- a/system/vulkan_enc/ResourceTracker.h
+++ b/system/vulkan_enc/ResourceTracker.h
@@ -67,6 +67,10 @@
         const VkDeviceCreateInfo* pCreateInfo,
         const VkAllocationCallbacks* pAllocator,
         VkDevice* pDevice);
+    void on_vkDestroyDevice_pre(
+        void* context,
+        VkDevice device,
+        const VkAllocationCallbacks* pAllocator);
 
     VkResult on_vkAllocateMemory(
         void* context,
diff --git a/system/vulkan_enc/VkEncoder.cpp b/system/vulkan_enc/VkEncoder.cpp
index 3822de8..7b9e989 100644
--- a/system/vulkan_enc/VkEncoder.cpp
+++ b/system/vulkan_enc/VkEncoder.cpp
@@ -792,6 +792,7 @@
     VkDevice device,
     const VkAllocationCallbacks* pAllocator)
 {
+    mImpl->resources()->on_vkDestroyDevice_pre(this, device, pAllocator);
     auto stream = mImpl->stream();
     auto countingStream = mImpl->countingStream();
     auto resources = mImpl->resources();
@@ -840,6 +841,7 @@
         marshal_VkAllocationCallbacks(stream, (VkAllocationCallbacks*)(local_pAllocator));
     }
     resources->destroyMapping()->mapHandles_VkDevice((VkDevice*)&device);
+    stream->flush();
 }
 
 VkResult VkEncoder::vkEnumerateInstanceExtensionProperties(
diff --git a/system/vulkan_enc/VulkanStream.cpp b/system/vulkan_enc/VulkanStream.cpp
index f862bf7..5b2e48d 100644
--- a/system/vulkan_enc/VulkanStream.cpp
+++ b/system/vulkan_enc/VulkanStream.cpp
@@ -70,6 +70,10 @@
         return mCurrentHandleMapping;
     }
 
+    void flush() {
+        commitWrite();
+    }
+
 private:
     size_t oustandingWriteBuffer() const {
         return mWritePos;
@@ -179,6 +183,10 @@
     return mImpl->handleMapping();
 }
 
+void VulkanStream::flush() {
+    mImpl->flush();
+}
+
 VulkanCountingStream::VulkanCountingStream() : VulkanStream(nullptr) { }
 VulkanCountingStream::~VulkanCountingStream() = default;
 
diff --git a/system/vulkan_enc/VulkanStream.h b/system/vulkan_enc/VulkanStream.h
index b763f32..3d17d8f 100644
--- a/system/vulkan_enc/VulkanStream.h
+++ b/system/vulkan_enc/VulkanStream.h
@@ -50,6 +50,7 @@
     void unsetHandleMapping();
     VulkanHandleMapping* handleMapping() const;
 
+    void flush();
 private:
     class Impl;
     std::unique_ptr<Impl> mImpl;
@@ -72,4 +73,4 @@
     size_t m_read = 0;
 };
 
-} // namespace goldfish_vk
\ No newline at end of file
+} // namespace goldfish_vk