From e1342ab836a0a581d3dd41807edcba5bffeadd9f Mon Sep 17 00:00:00 2001
From: kearwood <kearwood@kearwood.com>
Date: Wed, 20 Jul 2022 21:35:20 -0700
Subject: [PATCH] Added KRDevice::checkFlushStreamBuffer.

---
 kraken/KRDevice.cpp | 16 ++++++++++++++++
 kraken/KRDevice.h   |  1 +
 2 files changed, 17 insertions(+)

diff --git a/kraken/KRDevice.cpp b/kraken/KRDevice.cpp
index 0560acc..90901f4 100644
--- a/kraken/KRDevice.cpp
+++ b/kraken/KRDevice.cpp
@@ -417,6 +417,7 @@ bool KRDevice::initStagingBuffer(VkDeviceSize size, StagingBufferInfo* info
   if (vmaMapMemory(m_allocator, info->allocation, &info->data) != VK_SUCCESS) {
     return false;
   }
+  info->size = size;
   return true;
 }
 
@@ -645,8 +646,21 @@ void KRDevice::streamUpload(KRDataBlock& data, VkBuffer destination)
   data.unlock();
 }
 
+void KRDevice::checkFlushStreamBuffer(size_t size)
+{
+  // Flush the buffers if we would run out of space
+  if (m_streamingStagingBuffer.usage + size > m_streamingStagingBuffer.size) {
+    // If we hit this often, then we need a larger staging buffer.
+    // TODO - Dynamically allocate more/larger staging buffers.
+    assert(size < m_streamingStagingBuffer.size);
+    streamEnd();
+    streamStart();
+  }
+}
+
 void KRDevice::streamUpload(void* data, size_t size, VkBuffer destination)
 {
+  checkFlushStreamBuffer(size);
   memcpy((uint8_t*)m_streamingStagingBuffer.data + m_streamingStagingBuffer.usage, data, size);
 
   // TODO - Beneficial to batch many regions in a single call?
@@ -662,6 +676,7 @@ void KRDevice::streamUpload(void* data, size_t size, VkBuffer destination)
 
 void KRDevice::streamUpload(void* data, size_t size, Vector2i dimensions, VkImage destination)
 {
+  checkFlushStreamBuffer(size);
 
   memcpy((uint8_t*)m_streamingStagingBuffer.data + m_streamingStagingBuffer.usage, data, size);
 
@@ -749,6 +764,7 @@ void KRDevice::streamUpload(void* data, size_t size, Vector2i dimensions, VkImag
 void KRDevice::streamEnd()
 {
   vkEndCommandBuffer(m_transferCommandBuffers[0]);
+  m_streamingStagingBuffer.usage = 0;
 
   // TODO - Should double buffer and use a fence rather than block the thread
   VkSubmitInfo submitInfo{};
diff --git a/kraken/KRDevice.h b/kraken/KRDevice.h
index b748db8..4862490 100644
--- a/kraken/KRDevice.h
+++ b/kraken/KRDevice.h
@@ -106,6 +106,7 @@ public:
   // TODO - We should allocate at least two of these and double-buffer for increased CPU-GPU concurrency
   StagingBufferInfo m_graphicsStagingBuffer;
 private:
+  void checkFlushStreamBuffer(size_t size);
 
   void getQueueFamiliesForSharing(uint32_t* queueFamilyIndices, uint32_t* familyCount);