1. add DeviceMemory pointer to CommandBufferStagingStream am: 036bd1d742

Original change: https://android-review.googlesource.com/c/device/generic/goldfish-opengl/+/2403131

Change-Id: I95accc9a62ee4d28881ec7d74fd2e46579fd37e5
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/system/vulkan_enc/CommandBufferStagingStream.cpp b/system/vulkan_enc/CommandBufferStagingStream.cpp
index 0890185..faee5bf 100644
--- a/system/vulkan_enc/CommandBufferStagingStream.cpp
+++ b/system/vulkan_enc/CommandBufferStagingStream.cpp
@@ -33,127 +33,130 @@
 static const size_t kReadSize = 512 * 1024;
 static const size_t kWriteOffset = kReadSize;
 
-CommandBufferStagingStream::CommandBufferStagingStream(Alloc&& allocFn, Free&& freeFn)
-    : IOStream(1048576),
-      m_buf(nullptr),
-      m_size(0),
-      m_writePos(0),
-      m_customAlloc(allocFn),
-      m_customFree(freeFn) {
-    // custom allocator/free
-    if (allocFn && freeFn) {
-        m_usingCustomAlloc = true;
-        // for custom allocation, allocate metadata memory at the beginning.
-        // m_alloc, m_free and m_realloc wraps sync data logic
-
-        // \param size to allocate
-        // \return ptr starting at data
-        m_alloc = [this](size_t size) -> void* {
-            // allocation requested size + sync data size
-
-            // <---sync bytes--><----Data--->
-            // |———————————————|————————————|
-            // |0|1|2|3|4|5|6|7|............|
-            // |———————————————|————————————|
-            // ꜛ               ꜛ
-            // allocated ptr   ptr to data [dataPtr]
-            const size_t totalSize = size + kSyncDataSize;
-
-            unsigned char* dataPtr = static_cast<unsigned char*>(m_customAlloc(totalSize));
-            if (!dataPtr) {
-                ALOGE("Custom allocation (%zu bytes) failed\n", size);
-                return nullptr;
-            }
-
-            // set DWORD sync data to 0
-            *(reinterpret_cast<uint32_t*>(dataPtr)) = kSyncDataReadComplete;
-
-            // pointer for data starts after sync data
-            dataPtr += kSyncDataSize;
-
-            return dataPtr;
+CommandBufferStagingStream::CommandBufferStagingStream()
+    : IOStream(1048576), m_size(0), m_writePos(0) {
+    // use default allocators
+    m_alloc = [](size_t size) -> Memory {
+        return {
+            .deviceMemory = nullptr,  // no device memory for malloc
+            .ptr = malloc(size),
         };
+    };
+    m_free = [](const Memory& mem) { free(mem.ptr); };
+    m_realloc = [](const Memory& mem, size_t size) -> Memory {
+        return {.deviceMemory = nullptr, .ptr = realloc(mem.ptr, size)};
+    };
+}
 
-        // Free freeMemory(freeFn);
-        //  \param dataPtr to free
-        m_free = [this](void* dataPtr) {
-            // for custom allocation/free, memory holding metadata must be freed
-            // <---sync byte---><----Data--->
-            // |———————————————|————————————|
-            // |0|1|2|3|4|5|6|7|............|
-            // |———————————————|————————————|
-            // ꜛ               ꜛ
-            // ptr to free     ptr to data [dataPtr]
-            unsigned char* toFreePtr = static_cast<unsigned char*>(dataPtr);
-            toFreePtr -= kSyncDataSize;
-            m_customFree(toFreePtr);
-        };
+CommandBufferStagingStream::CommandBufferStagingStream(const Alloc& allocFn, const Free& freeFn)
+    : CommandBufferStagingStream() {
+    m_usingCustomAlloc = true;
+    // for custom allocation, allocate metadata memory at the beginning.
+    // m_alloc, m_free and m_realloc wraps sync data logic
 
-        // \param ptr is the data pointer currently allocated
-        // \return dataPtr starting at data
-        m_realloc = [this](void* ptr, size_t size) -> void* {
-            // realloc requires freeing previously allocated memory
-            // read sync DWORD to ensure host is done reading this memory
-            // before releasing it.
+    // \param size to allocate
+    // \return ptr starting at data
+    m_alloc = [&allocFn, this](size_t size) -> Memory {
+        // allocation requested size + sync data size
 
-            size_t hostWaits = 0;
-            unsigned char* syncDataStart = static_cast<unsigned char*>(ptr) - kSyncDataSize;
-            uint32_t* syncDWordPtr = reinterpret_cast<uint32_t*>(syncDataStart);
+        // <---sync bytes--><----Data--->
+        // |———————————————|————————————|
+        // |0|1|2|3|4|5|6|7|............|
+        // |———————————————|————————————|
+        // ꜛ               ꜛ
+        // allocated ptr   ptr to data [dataPtr]
 
-            while (__atomic_load_n(syncDWordPtr, __ATOMIC_ACQUIRE) != kSyncDataReadComplete) {
-                hostWaits++;
-                usleep(10);
-                if (hostWaits > 1000) {
-                    ALOGD("%s: warning, stalled on host decoding on this command buffer stream\n",
-                          __func__);
-                }
+        Memory memory;
+        if (!allocFn) {
+            ALOGE("Custom allocation (%zu bytes) failed\n", size);
+            return memory;
+        }
+
+        // custom allocation/free requires metadata for sync between host/guest
+        const size_t totalSize = size + kSyncDataSize;
+        memory = allocFn(totalSize);
+        if (!memory.ptr) {
+            ALOGE("Custom allocation (%zu bytes) failed\n", size);
+            return memory;
+        }
+
+        // set sync data to read complete
+        uint32_t* syncDWordPtr = reinterpret_cast<uint32_t*>(memory.ptr);
+        __atomic_store_n(syncDWordPtr, kSyncDataReadComplete, __ATOMIC_RELEASE);
+        return memory;
+    };
+
+    m_free = [&freeFn](const Memory& mem) {
+        if (!freeFn) {
+            ALOGE("Custom free for device memory(%p) failed\n", mem.deviceMemory);
+            return;
+        }
+        freeFn(mem);
+    };
+
+    // \param ptr is the data pointer currently allocated
+    // \return dataPtr starting at data
+    m_realloc = [this](const Memory& mem, size_t size) -> Memory {
+        // realloc requires freeing previously allocated memory
+        // read sync DWORD to ensure host is done reading this memory
+        // before releasing it.
+
+        size_t hostWaits = 0;
+
+        uint32_t* syncDWordPtr = reinterpret_cast<uint32_t*>(mem.ptr);
+        while (__atomic_load_n(syncDWordPtr, __ATOMIC_ACQUIRE) != kSyncDataReadComplete) {
+            hostWaits++;
+            usleep(10);
+            if (hostWaits > 1000) {
+                ALOGD("%s: warning, stalled on host decoding on this command buffer stream\n",
+                      __func__);
             }
+        }
 
-            // for custom allocation/free, memory holding metadata must be copied
-            // along with stream data
-            // <---sync byte---><----Data--->
-            // |———————————————|————————————|
-            // |0|1|2|3|4|5|6|7|............|
-            // |———————————————|————————————|
-            // ꜛ               ꜛ
-            // [copyLocation]  ptr to data [ptr]
+        // for custom allocation/free, memory holding metadata must be copied
+        // along with stream data
+        // <---sync bytes--><----Data--->
+        // |———————————————|————————————|
+        // |0|1|2|3|4|5|6|7|............|
+        // |———————————————|————————————|
+        // ꜛ               ꜛ
+        // [copyLocation]  ptr to data [ptr]
 
-            const size_t toCopySize = m_writePos + kSyncDataSize;
-            unsigned char* copyLocation = static_cast<unsigned char*>(ptr) - kSyncDataSize;
-            std::vector<uint8_t> tmp(copyLocation, copyLocation + toCopySize);
-            m_free(ptr);
+        const size_t toCopySize = m_writePos + kSyncDataSize;
+        unsigned char* copyLocation = static_cast<unsigned char*>(mem.ptr);
+        std::vector<uint8_t> tmp(copyLocation, copyLocation + toCopySize);
+        m_free(mem);
 
-            // get new buffer and copy previous stream data to it
-            unsigned char* newBuf = static_cast<unsigned char*>(m_alloc(size));
-            if (!newBuf) {
-                ALOGE("Custom allocation (%zu bytes) failed\n", size);
-                return nullptr;
-            }
-            // custom allocator will allocate space for metadata too
-            // copy previous metadata too
-            memcpy(newBuf - kSyncDataSize, tmp.data(), toCopySize);
+        // get new buffer and copy previous stream data to it
+        Memory newMemory = m_alloc(size);
+        unsigned char* newBuf = static_cast<unsigned char*>(newMemory.ptr);
+        if (!newBuf) {
+            ALOGE("Custom allocation (%zu bytes) failed\n", size);
+            return newMemory;
+        }
+        // copy previous data
+        memcpy(newBuf, tmp.data(), toCopySize);
 
-            return newBuf;
-        };
-    } else {
-        // use default allocators
-        m_alloc = [](size_t size) { return malloc(size); };
-        m_free = [](void* ptr) { free(ptr); };
-        m_realloc = [](void* ptr, size_t size) { return realloc(ptr, size); };
-    }
+        return newMemory;
+    };
 }
 
 CommandBufferStagingStream::~CommandBufferStagingStream() {
     flush();
-    if (m_buf) m_free(m_buf);
+    if (m_mem.ptr) m_free(m_mem);
+}
+
+unsigned char* CommandBufferStagingStream::getDataPtr() {
+    if (!m_mem.ptr) return nullptr;
+    const size_t metadataSize = m_usingCustomAlloc ? kSyncDataSize : 0;
+    return static_cast<unsigned char*>(m_mem.ptr) + metadataSize;
 }
 
 void CommandBufferStagingStream::markFlushing() {
     if (!m_usingCustomAlloc) {
         return;
     }
-    // mark read of stream buffer as pending
-    uint32_t* syncDWordPtr = reinterpret_cast<uint32_t*>(m_buf - kSyncDataSize);
+    uint32_t* syncDWordPtr = reinterpret_cast<uint32_t*>(m_mem.ptr);
     __atomic_store_n(syncDWordPtr, kSyncDataReadPending, __ATOMIC_RELEASE);
 }
 
@@ -165,10 +168,10 @@
 void* CommandBufferStagingStream::allocBuffer(size_t minSize) {
     size_t allocSize = (1048576 < minSize ? minSize : 1048576);
     // Initial case: blank
-    if (!m_buf) {
-        m_buf = (unsigned char*)m_alloc(allocSize);
+    if (!m_mem.ptr) {
+        m_mem = m_alloc(allocSize);
         m_size = allocSize;
-        return (void*)m_buf;
+        return getDataPtr();
     }
 
     // Calculate remaining
@@ -177,13 +180,13 @@
     // if not, reallocate a buffer of big enough size
     if (remaining < minSize) {
         size_t newAllocSize = m_size * 2 + allocSize;
-        m_buf = (unsigned char*)m_realloc(m_buf, newAllocSize);
+        m_mem = m_realloc(m_mem, newAllocSize);
         m_size = newAllocSize;
 
-        return (void*)(m_buf + m_writePos);
+        return (void*)(getDataPtr() + m_writePos);
     }
 
-    return (void*)(m_buf + m_writePos);
+    return (void*)(getDataPtr() + m_writePos);
 }
 
 int CommandBufferStagingStream::commitBuffer(size_t size)
@@ -224,7 +227,7 @@
 }
 
 void CommandBufferStagingStream::getWritten(unsigned char** bufOut, size_t* sizeOut) {
-    *bufOut = m_buf;
+    *bufOut = getDataPtr();
     *sizeOut = m_writePos;
 }
 
diff --git a/system/vulkan_enc/CommandBufferStagingStream.h b/system/vulkan_enc/CommandBufferStagingStream.h
index 9513b45..0fb859d 100644
--- a/system/vulkan_enc/CommandBufferStagingStream.h
+++ b/system/vulkan_enc/CommandBufferStagingStream.h
@@ -31,18 +31,30 @@
  // indicates read is pending
  static constexpr uint32_t kSyncDataReadPending = 0X1;
 
+ // alias for DeviceMemory
+ using DeviceMemory = void*;
+ // \struct backing memory structure
+ struct Memory {
+     DeviceMemory deviceMemory = nullptr;  // device memory associated with allocated memory
+     void* ptr = nullptr;                  // pointer to allocated memory
+     bool operator==(const Memory& rhs) const {
+         return (deviceMemory == rhs.deviceMemory) && (ptr == rhs.ptr);
+     }
+ };
+
  // allocator
  // param size to allocate
- // return pointer to allocated memory
- using Alloc = std::function<void*(size_t)>;
+ // return allocated memory
+ using Alloc = std::function<Memory(size_t)>;
  // free function
- // param pointer to free
- using Free = std::function<void(void*)>;
+ // param memory to free
+ using Free = std::function<void(const Memory&)>;
  // constructor
- // \param allocFn is the allocation function provided. Default allocation function used if nullptr
- // \param freeFn is the free function provided. Default free function used if nullptr
- // freeFn must be provided if allocFn is provided and vice versa
- explicit CommandBufferStagingStream(Alloc&& allocFn = nullptr, Free&& freeFn = nullptr);
+ // \param allocFn is the allocation function provided.
+ // \param freeFn is the free function provided
+ explicit CommandBufferStagingStream(const Alloc& allocFn, const Free& freeFn);
+ // constructor
+ explicit CommandBufferStagingStream();
  ~CommandBufferStagingStream();
 
  virtual size_t idealAllocSize(size_t len);
@@ -64,37 +76,32 @@
  void markFlushing();
 
 private:
- // underlying buffer for data
- unsigned char* m_buf;
- // size of portion of m_buf available for data.
+ // underlying memory for data
+ Memory m_mem;
+ // size of portion of memory available for data.
  // for custom allocation, this size excludes size of sync data.
  size_t m_size;
- // current write position in m_buf
+ // current write position in data buffer
  uint32_t m_writePos;
 
+ // alloc function
  Alloc m_alloc;
+ // free function
  Free m_free;
 
- // underlying custom alloc. default is null
- Alloc m_customAlloc = nullptr;
- // underlying free alloc. default is null
- Free m_customFree = nullptr;
-
  // realloc function
  // \param size of memory to be allocated
  // \ param reference size to update with actual size allocated. This size can be < requested size
  // for custom allocation to account for sync data
- using Realloc = std::function<void*(void*, size_t)>;
+ using Realloc = std::function<Memory(const Memory&, size_t)>;
  Realloc m_realloc;
 
  // flag tracking use of custom allocation/free
  bool m_usingCustomAlloc = false;
 
- // calculates actual allocation size for data
- // \param requestedSize is the size requested for allocation
- // \return actual data size allocated for requested size. For
- // custom allocation the data size < requested size to account for sync data word
- size_t getDataAllocationSize(const size_t requestedSize);
+ // adjusted memory location to point to start of data after accounting for metadata
+ // \return pointer to data start
+ unsigned char* getDataPtr();
 };
 
 #endif
diff --git a/system/vulkan_enc_unit_tests/CommandBufferStagingStream_test.cpp b/system/vulkan_enc_unit_tests/CommandBufferStagingStream_test.cpp
index b1df082..eb34cfb 100644
--- a/system/vulkan_enc_unit_tests/CommandBufferStagingStream_test.cpp
+++ b/system/vulkan_enc_unit_tests/CommandBufferStagingStream_test.cpp
@@ -201,30 +201,37 @@
         << "commitBufferAndReadFully should not be supported";
 }
 
-// CommandBufferStagingStreamCustomAllocationTest tests tests behavior of CommandBufferStagingStream
+using MockAlloc = MockFunction<CommandBufferStagingStream::Memory(size_t)>;
+using MockFree = MockFunction<void(const CommandBufferStagingStream::Memory&)>;
+// default empty implementation of free
+static std::function<void(const CommandBufferStagingStream::Memory&)> EmptyFree =
+    [](const CommandBufferStagingStream::Memory&) {};
+// CommandBufferStagingStreamCustomAllocationTest tests behavior of CommandBufferStagingStream
 // when initialized with custom allocator/free function.
 // These tests test the same outcome as CommandBufferStagingStreamTest tests
 
 // tests allocBuffer can successfully allocate a buffer of given size
-
 TEST(CommandBufferStagingStreamCustomAllocationTest, AllocateBufferTest) {
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
 
-    MockFunction<void*(size_t)> allocFn;
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+
+    MockAlloc allocFn;
 
     // alloc function should be called once
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
     // free function should be called once
-    MockFunction<void(void*)> freeFn;
-    EXPECT_CALL(freeFn, Call(Eq(static_cast<void*>(memorySrc.data())))).Times(1);
+    MockFree freeFn;
+    EXPECT_CALL(freeFn, Call(Eq(memory))).Times(1);
 
     // scope: CommandBufferStagingStream_Creation
     {
-        CommandBufferStagingStream stream(allocFn.AsStdFunction(), freeFn.AsStdFunction());
+        auto allocStdFn = allocFn.AsStdFunction();
+        auto freeStdFn = freeFn.AsStdFunction();
+        CommandBufferStagingStream stream(allocStdFn, freeStdFn);
         uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
         EXPECT_THAT(buffer, NotNull());
     }
@@ -235,24 +242,84 @@
     // memory source for initial allocation
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
 
-    MockFunction<void*(size_t)> allocFn;
+    CommandBufferStagingStream::Memory memory{
+        .deviceMemory = nullptr,  // not needed for this test
+        .ptr = nullptr,           // to test alloc call failing
+    };
+
+    MockAlloc allocFn;
 
     // alloc function should be called once
-    // return nullptr to test alloc call failing
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(nullptr));
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
     // free function should not be called if allocation fails
-    MockFunction<void(void*)> freeFn;
+    MockFree freeFn;
     EXPECT_CALL(freeFn, Call).Times(0);
 
     // scope: CommandBufferStagingStream_Creation
     {
-        CommandBufferStagingStream stream(allocFn.AsStdFunction(), freeFn.AsStdFunction());
+        auto allocStdFn = allocFn.AsStdFunction();
+        auto freeStdFn = freeFn.AsStdFunction();
+        CommandBufferStagingStream stream(allocStdFn, freeStdFn);
         void* buffer = stream.allocBuffer(kTestBufferSize);
         EXPECT_THAT(buffer, IsNull());
     }
 }
 
+TEST(CommandBufferStagingStreamCustomAllocationTest, DeviceMemoryPointerIsPassedDuringFree) {
+    // memory source
+    std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
+
+    // device memory for test purposes. The test just needs a pointer
+    uint32_t deviceMem = 0;
+    uint32_t* deviceMemPtr = &deviceMem;
+
+    CommandBufferStagingStream::Memory memory{.deviceMemory = deviceMemPtr,
+                                              .ptr = memorySrc.data()};
+
+    MockAlloc allocFn;
+
+    // alloc function should be called once
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
+
+    // free function should be called once
+    MockFree freeFn;
+    EXPECT_CALL(freeFn, Call(Eq(memory))).Times(1);
+
+    // scope: CommandBufferStagingStream_Creation
+    {
+        auto allocStdFn = allocFn.AsStdFunction();
+        auto freeStdFn = freeFn.AsStdFunction();
+        CommandBufferStagingStream stream(allocStdFn, freeStdFn);
+        uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
+        EXPECT_THAT(buffer, NotNull());
+    }
+}
+
+// test verifies that there are no crashes if alloc/free function reference becomes null
+TEST(CommandBufferStagingStreamCustomAllocationTest, AllocFreeInvalidReference) {
+    MockAlloc allocFn;
+    // alloc shouldn't be called if reference is invalidated
+    EXPECT_CALL(allocFn, Call).Times(0);
+
+    MockFree freeFn;
+    // free shouldn't be called if reference is invalidated
+    EXPECT_CALL(freeFn, Call).Times(0);
+
+    auto allocStdFn = allocFn.AsStdFunction();
+    auto freeStdFn = freeFn.AsStdFunction();
+    // scope: CommandBufferStagingStream_Creation
+    {
+        CommandBufferStagingStream stream(allocStdFn, freeStdFn);
+        // invalidate alloc/free functions
+        allocStdFn = nullptr;
+        freeStdFn = nullptr;
+        stream.allocBuffer(kTestBufferSize);
+        uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
+        EXPECT_THAT(buffer, IsNull());
+    }
+}
+
 // test reallocate buffer remembers previously committed buffers
 TEST(CommandBufferStagingStreamCustomAllocationTest, ReallocateBuffer) {
     // memory source for initial allocation
@@ -260,35 +327,42 @@
     // memory source after reallocation
     std::vector<uint8_t> reallocatedMemorySrc(kTestBufferSize * 3);
 
-    MockFunction<void*(size_t)> allocFn;
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+
+    CommandBufferStagingStream::Memory reallocatedMemory{
+        .deviceMemory = nullptr,  // not needed for this test
+        .ptr = reallocatedMemorySrc.data()};
+
+    MockAlloc allocFn;
 
     // alloc function should be called twice
     {
         InSequence seq;
 
         // expect initial allocation call with allocation size == kTestBufferSize;
-        EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-            .Times(1)
-            .WillRepeatedly(Return(memorySrc.data()));
+        EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
         // expect reallocation call with allocation size > kTestBufferSize.
         EXPECT_CALL(allocFn, Call(testing::Ge(kTestBufferSize)))
             .Times(1)
-            .WillRepeatedly(Return(reallocatedMemorySrc.data()));
+            .WillRepeatedly(Return(reallocatedMemory));
     }
 
-    MockFunction<void(void*)> freeFn;
+    MockFree freeFn;
     {
         InSequence seq;
         // free function should be called when reallocation happens
-        EXPECT_CALL(freeFn, Call(Eq(memorySrc.data()))).Times(1);
+        EXPECT_CALL(freeFn, Call(Eq(memory))).Times(1);
         // free function should be called when stream goes out of scope
-        EXPECT_CALL(freeFn, Call(Eq(reallocatedMemorySrc.data()))).Times(1);
+        EXPECT_CALL(freeFn, Call(Eq(reallocatedMemory))).Times(1);
     }
 
     // scope: CommandBufferStagingStream_Creation
     {
-        CommandBufferStagingStream stream(allocFn.AsStdFunction(), freeFn.AsStdFunction());
+        auto allocStdFn = allocFn.AsStdFunction();
+        auto freeStdFn = freeFn.AsStdFunction();
+        CommandBufferStagingStream stream(allocStdFn, freeStdFn);
         uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
         EXPECT_THAT(buffer, NotNull());
 
@@ -318,14 +392,16 @@
 TEST(CommandBufferStagingStreamCustomAllocationTest, CommitBuffer) {
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
-    MockFunction<void*(size_t)> allocFn;
+
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+
+    MockAlloc allocFn;
 
     // alloc function should be called once
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
-
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), [](void*) {});
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
     EXPECT_THAT(buffer, NotNull());
 
@@ -351,17 +427,15 @@
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
 
-    MockFunction<void*(size_t)> allocFn;
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+
+    MockAlloc allocFn;
 
     // alloc function should be called once
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
-
-    // free function expectation not needed in this test
-    MockFunction<void(void*)> freeFn;
-
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), freeFn.AsStdFunction());
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
 
     // write some arbitrary data
@@ -389,18 +463,16 @@
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
 
-    MockFunction<void*(size_t)> allocFn;
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+
+    MockAlloc allocFn;
 
     // alloc function should be called once, no reallocation
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
-    // free function expectation not needed in this test
-    MockFunction<void(void*)> freeFn;
-
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), freeFn.AsStdFunction());
-
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
     EXPECT_THAT(buffer, NotNull());
 
@@ -413,31 +485,32 @@
 TEST(CommandBufferStagingStreamCustomAllocationTest, ReallocationBoundary) {
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 3);
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
 
-    MockFunction<void*(size_t)> allocFn;
+    MockAlloc allocFn;
 
     // alloc function should be called twice
     {
         InSequence seq;
 
         // expect initial allocation call with allocation size >= kTestBufferSize;
-        EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-            .Times(1)
-            .WillRepeatedly(Return(memorySrc.data()));
+        EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
         // expect reallocation call with allocation size > kTestBufferSize.
         EXPECT_CALL(allocFn, Call(testing::Ge(kTestBufferSize)))
             .Times(1)
-            .WillRepeatedly(Return(memorySrc.data()));
+            .WillRepeatedly(Return(memory));
     }
 
     // free function should be called once during reallocation,
     // once when stream goes out of scope
-    MockFunction<void(void*)> freeFn;
+    MockFree freeFn;
 
-    EXPECT_CALL(freeFn, Call(Eq(memorySrc.data()))).Times(2);
-
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), freeFn.AsStdFunction());
+    EXPECT_CALL(freeFn, Call(Eq(memory))).Times(2);
+    auto allocStdFn = allocFn.AsStdFunction();
+    auto freeStdFn = freeFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, freeStdFn);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
     EXPECT_THAT(buffer, NotNull());
 
@@ -473,12 +546,14 @@
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 3);
 
-    MockFunction<void*(size_t)> allocFn;
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
 
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), [](void*) {});
+    MockAlloc allocFn;
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
+
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
     EXPECT_THAT(buffer, NotNull());
 
@@ -495,30 +570,36 @@
 TEST(CommandBufferStagingStreamCustomAllocationTest, MetadataCheck) {
     // memory source
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+
     // CommandBufferStagingStream allocates metadata when using custom allocators
     static const size_t expectedMetadataSize = 8;
-    MockFunction<void*(size_t)> allocFn;
+    MockAlloc allocFn;
 
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), [](void*) {});
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
-    EXPECT_THAT(buffer, NotNull());
+    // data should start after metadata
+    EXPECT_THAT(buffer, memorySrc.data() + expectedMetadataSize);
+    // metadata should be initialized to read complete
+    uint32_t* metadataPtr = reinterpret_cast<uint32_t*>(memorySrc.data());
+    EXPECT_THAT(*metadataPtr, CommandBufferStagingStream::kSyncDataReadComplete);
 }
 
-TEST(CommandBufferStagingStreamCustomAllocationTest, MarkFlushingTest) {
+TEST(CommandBufferStagingStreamCustomAllocationTest, MarkFlushing) {
     // memory source for  allocation
     std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
+    MockAlloc allocFn;
 
-    MockFunction<void*(size_t)> allocFn;
+    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
-
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), [](void*) {});
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);
     uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
 
     // write some data
@@ -536,47 +617,18 @@
     EXPECT_EQ(*readPtr, CommandBufferStagingStream::kSyncDataReadPending);
 }
 
-TEST(CommandBufferStagingStreamCustomAllocationTest, MarkFlushing) {
-    // memory source for  allocation
-    std::vector<uint8_t> memorySrc(kTestBufferSize * 2);
-
-    MockFunction<void*(size_t)> allocFn;
-
-    EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-        .Times(1)
-        .WillRepeatedly(Return(memorySrc.data()));
-
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), [](void*) {});
-    uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
-
-    // write some data
-    const std::string_view commandData{"some command"};
-    const size_t dataSize = commandData.size();
-    memcpy(buffer, commandData.data(), dataSize);
-
-    // commit data
-    stream.commitBuffer(dataSize);
-
-    // will set metadata of the stream buffer to pending read
-    stream.markFlushing();
-
-    uint32_t* readPtr = reinterpret_cast<uint32_t*>(memorySrc.data());
-    EXPECT_EQ(*readPtr, 0x01);
-}
-
 // this test verifies that realloc waits till consumer of memory has completed read
 TEST(CommandBufferStagingStreamCustomAllocationTest, ReallocNotCalledTillBufferIsRead) {
     // memory source for  allocation
     // allocate a big enough buffer to avoid resizes in test
     std::vector<uint8_t> memorySrc(kTestBufferSize * 3);
-    auto ptr = memorySrc.data();
+    CommandBufferStagingStream::Memory memory{.deviceMemory = nullptr,  // not needed for this test
+                                              .ptr = memorySrc.data()};
 
     std::condition_variable memoryFlushedCondition;
     std::mutex mutex;
 
-    // track the number of times allocFn is called
-
-    MockFunction<void*(size_t)> allocFn;
+    MockAlloc allocFn;
 
     // mock function to notify read is complete
     // this will be used to set up the expectation that realloc should
@@ -589,7 +641,8 @@
         std::unique_lock readLock(mutex);
         memoryFlushedCondition.wait(readLock, [&]() {
             // wait till memorySrc is ready for read
-            return *ptr == CommandBufferStagingStream::kSyncDataReadPending;
+            uint32_t* syncData = static_cast<uint32_t*>(memory.ptr);
+            return *syncData = CommandBufferStagingStream::kSyncDataReadPending;
         });
 
         readLock.unlock();
@@ -601,14 +654,12 @@
         fn();
     });
 
-    CommandBufferStagingStream stream(allocFn.AsStdFunction(), [](void*) {});
-    // scope for writeLock
+    auto allocStdFn = allocFn.AsStdFunction();
+    CommandBufferStagingStream stream(allocStdFn, EmptyFree);  // scope for writeLock
     {
         std::lock_guard writeLock(mutex);
 
-        EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize)))
-            .Times(1)
-            .WillRepeatedly(Return(memorySrc.data()));
+        EXPECT_CALL(allocFn, Call(Ge(kTestBufferSize))).Times(1).WillRepeatedly(Return(memory));
 
         uint8_t* buffer = static_cast<uint8_t*>(stream.allocBuffer(kTestBufferSize));
 
@@ -630,7 +681,7 @@
     EXPECT_CALL(allocFn, Call(testing::Ge(kTestBufferSize)))
         .Times(1)
         .After(readCompleteExpectation)
-        .WillRepeatedly(Return(memorySrc.data()));
+        .WillRepeatedly(Return(memory));
 
     // realloc will be blocked till buffer read is complete by reader
     (void)stream.allocBuffer(kTestBufferSize);