Change stream I/O behaviour around sockets

* Previously each read and write would 1:1 become a send or recv call.
* Now we buffer writes and send in batch (or when a chunk finishes),
  and every time we read we try to non-blocking read more data to fill
  the buffer, to allow batching reads where possible without blocking
  on data that will never come.

renderdoc/serialise/serialiser.cpp

@@ -404,6 +404,8 @@ void Serialiser<SerialiserMode::Writing>::EndChunk()
   m_Write->AlignTo<ChunkAlignment>();
 
   m_ChunkMetadata = SDChunkMetaData();
+
+  m_Write->Flush();
 }
 
 template <>

renderdoc/serialise/streamio.cpp

@@ -24,6 +24,7 @@
 
 #include "streamio.h"
 #include <errno.h>
+#include "common/timing.h"
 
 Compressor::~Compressor()
 {
@@ -76,11 +77,12 @@ StreamReader::StreamReader(Network::Socket *sock, Ownership own)
 {
   m_Sock = sock;
 
-  m_InputSize = m_BufferSize = initialBufferSize;
+  m_BufferSize = initialBufferSize;
   m_BufferBase = AllocAlignedBuffer(m_BufferSize);
-  // place head at the *end* of the buffer, because for sockets we pretend the buffer is constantly
-  // exhausted, and just do a read of the minimum number of bytes each time to satisfy each read
-  m_BufferHead = m_BufferBase + m_BufferSize;
+  m_BufferHead = m_BufferBase;
+
+  // for sockets we use m_InputSize to indicate how much data has been read into the buffer.
+  m_InputSize = 0;
 
   m_Ownership = own;
 }
@@ -189,25 +191,7 @@ void StreamReader::SetOffset(uint64_t offs)
 
 bool StreamReader::Reserve(uint64_t numBytes)
 {
-  if(m_Sock)
-  {
-    // if we're reading more bytes than our buffer size, resize up
-    if(numBytes > m_BufferSize)
-    {
-      FreeAlignedBuffer(m_BufferBase);
-      m_InputSize = m_BufferSize = AlignUp<uint64_t>(numBytes, 256ULL);
-      m_BufferBase = AllocAlignedBuffer(m_BufferSize);
-      m_BufferHead = m_BufferBase + m_BufferSize;
-    }
-
-    m_ReadOffset += numBytes;
-
-    // read into the end of the buffer, so that the subsequent read re-exhausts the buffer
-    m_BufferHead = m_BufferBase + m_BufferSize - numBytes;
-    return ReadFromExternal(m_BufferSize - numBytes, numBytes);
-  }
-
-  RDCASSERT(m_File || m_Decompressor);
+  RDCASSERT(m_Sock || m_File || m_Decompressor);
 
   // store old buffer and the read data, so we can move it into the new buffer
   byte *oldBuffer = m_BufferBase;
@@ -218,6 +202,9 @@ bool StreamReader::Reserve(uint64_t numBytes)
   byte *currentData = m_BufferHead - backwardsWindow;
   uint64_t currentDataSize = m_BufferSize - (m_BufferHead - m_BufferBase) + backwardsWindow;
 
+  if(m_Sock)
+    currentDataSize = m_InputSize - (m_BufferHead - m_BufferBase) + backwardsWindow;
+
   uint64_t BufferOffset = m_BufferHead - m_BufferBase;
 
   // if we are reading more than our current buffer size, expand the buffer size
@@ -243,11 +230,21 @@ bool StreamReader::Reserve(uint64_t numBytes)
     m_BufferHead = m_BufferBase + BufferOffset;
   }
 
+  if(m_Sock)
+    m_InputSize = currentDataSize;
+
   // if there's anything left of the file to read in, do so now
   bool ret = false;
 
-  ret = ReadFromExternal(currentDataSize, RDCMIN(m_BufferSize - currentDataSize,
-                                                 m_InputSize - m_ReadOffset - currentDataSize));
+  uint64_t readSize =
+      RDCMIN(m_BufferSize - currentDataSize, m_InputSize - m_ReadOffset - currentDataSize);
+
+  // we'll read as much as possible anyway using m_BufferSize, but we need to know how much we
+  // *must* read.
+  if(m_Sock)
+    readSize = numBytes - Available();
+
+  ret = ReadFromExternal(currentDataSize, readSize);
 
   if(oldBuffer != m_BufferBase && m_BufferBase)
     FreeAlignedBuffer(oldBuffer);
@@ -276,7 +273,24 @@ bool StreamReader::ReadFromExternal(uint64_t bufferOffs, uint64_t length)
     }
     else
     {
-      success = m_Sock->RecvDataBlocking(m_BufferBase + bufferOffs, (uint32_t)length);
+      // first get the required data blocking (this will sleep the thread until it comes in).
+      byte *readDest = m_BufferBase + bufferOffs;
+
+      success = m_Sock->RecvDataBlocking(readDest, (uint32_t)length);
+
+      if(success)
+      {
+        m_InputSize += length;
+        readDest += length;
+
+        uint32_t bufSize = uint32_t(m_BufferSize - m_InputSize);
+
+        // now read more, as much as possible, to try and batch future reads
+        success = m_Sock->RecvDataNonBlocking(readDest, bufSize);
+
+        if(success)
+          m_InputSize += bufSize;
+      }
     }
   }
   else
@@ -344,7 +358,8 @@ StreamWriter::StreamWriter(StreamInvalidType)
 
 StreamWriter::StreamWriter(Network::Socket *sock, Ownership own)
 {
-  m_BufferBase = m_BufferHead = m_BufferEnd = NULL;
+  m_BufferBase = m_BufferHead = AllocAlignedBuffer(initialBufferSize);
+  m_BufferEnd = m_BufferBase + initialBufferSize;
 
   m_Sock = sock;
 
@@ -389,6 +404,56 @@ StreamWriter::~StreamWriter()
   }
 }
 
+bool StreamWriter::SendSocketData(const void *data, uint64_t numBytes)
+{
+  // try to coalesce small writes without doing blocking sends, at least until we're flushed.
+  // if the buffer is already full, flush it.
+  if(m_BufferHead + numBytes >= m_BufferEnd)
+  {
+    bool success = FlushSocketData();
+    if(!success)
+    {
+      HandleError();
+      return false;
+    }
+  }
+
+  // if it's larger than our buffer (even after flushing) just write directly
+  if(m_BufferHead + numBytes >= m_BufferEnd)
+  {
+    bool success = m_Sock->SendDataBlocking(data, (uint32_t)numBytes);
+    if(!success)
+    {
+      HandleError();
+      return false;
+    }
+  }
+  else
+  {
+    // otherwise, write it into the in-memory buffer
+    memcpy(m_BufferHead, data, (size_t)numBytes);
+    m_BufferHead += numBytes;
+  }
+
+  return true;
+}
+
+bool StreamWriter::FlushSocketData()
+{
+  // send out what we have buffered up
+  bool success = m_Sock->SendDataBlocking(m_BufferBase, uint32_t(m_BufferHead - m_BufferBase));
+  if(!success)
+  {
+    HandleError();
+    return false;
+  }
+
+  // reset buffer to the start
+  m_BufferHead = m_BufferBase;
+
+  return true;
+}
+
 void StreamWriter::HandleError()
 {
   if(m_File)

renderdoc/serialise/streamio.h

@@ -91,7 +91,12 @@ public:
 
   inline uint64_t GetOffset() { return m_BufferHead - m_BufferBase + m_ReadOffset; }
   inline uint64_t GetSize() { return m_InputSize; }
-  inline bool AtEnd() { return GetOffset() >= GetSize(); }
+  inline bool AtEnd()
+  {
+    if(m_Sock)
+      return Available() == 0;
+    return GetOffset() >= GetSize();
+  }
   template <uint64_t alignment>
   bool AlignTo()
   {
@@ -189,7 +194,7 @@ private:
   inline uint64_t Available()
   {
     if(m_Sock)
-      return 0;
+      return m_InputSize - (m_BufferHead - m_BufferBase);
     return m_BufferSize - (m_BufferHead - m_BufferBase);
   }
   bool Reserve(uint64_t numBytes);
@@ -327,14 +332,7 @@ public:
     }
     else if(m_Sock)
     {
-      bool success = m_Sock->SendDataBlocking(data, (uint32_t)numBytes);
-      if(!success)
-      {
-        HandleError();
-        return false;
-      }
-
-      return true;
+      return SendSocketData(data, numBytes);
     }
     else
     {
@@ -396,6 +394,18 @@ public:
     return false;
   }
 
+  bool Flush()
+  {
+    if(m_Compressor)
+      return true;
+    else if(m_File)
+      return FileIO::fflush(m_File);
+    else if(m_Sock)
+      return FlushSocketData();
+
+    return true;
+  }
+
   bool Finish()
   {
     if(m_Compressor)
@@ -437,6 +447,9 @@ private:
 
   void HandleError();
 
+  bool SendSocketData(const void *data, uint64_t numBytes);
+  bool FlushSocketData();
+
   // used for aligned writes
   static const byte empty[128];