Made worker threads in ezCompressedStreamWriterZstd configurable

luoxz-ai · Jul 5, 2023 · a93eb42 · a93eb42
1 parent 30c391d
commit a93eb42
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 14 deletions.
diff --git a/Code/Engine/Foundation/IO/Archive/Implementation/ArchiveBuilder.cpp b/Code/Engine/Foundation/IO/Archive/Implementation/ArchiveBuilder.cpp
@@ -44,23 +44,23 @@ void ezArchiveBuilder::AddFolder(ezStringView sAbsFolderPath, ezArchiveCompressi
 
           case InclusionMode::Compress_zstd_fastest:
             compression = ezArchiveCompressionMode::Compressed_zstd;
-            iCompressionLevel = ezCompressedStreamWriterZstd::Compression::Fastest;
+            iCompressionLevel = static_cast<ezInt32>(ezCompressedStreamWriterZstd::Compression::Fastest);
             break;
           case InclusionMode::Compress_zstd_fast:
             compression = ezArchiveCompressionMode::Compressed_zstd;
-            iCompressionLevel = ezCompressedStreamWriterZstd::Compression::Fast;
+            iCompressionLevel = static_cast<ezInt32>(ezCompressedStreamWriterZstd::Compression::Fast);
             break;
           case InclusionMode::Compress_zstd_average:
             compression = ezArchiveCompressionMode::Compressed_zstd;
-            iCompressionLevel = ezCompressedStreamWriterZstd::Compression::Average;
+            iCompressionLevel = static_cast<ezInt32>(ezCompressedStreamWriterZstd::Compression::Average);
             break;
           case InclusionMode::Compress_zstd_high:
             compression = ezArchiveCompressionMode::Compressed_zstd;
-            iCompressionLevel = ezCompressedStreamWriterZstd::Compression::High;
+            iCompressionLevel = static_cast<ezInt32>(ezCompressedStreamWriterZstd::Compression::High);
             break;
           case InclusionMode::Compress_zstd_highest:
             compression = ezArchiveCompressionMode::Compressed_zstd;
-            iCompressionLevel = ezCompressedStreamWriterZstd::Compression::Highest;
+            iCompressionLevel = static_cast<ezInt32>(ezCompressedStreamWriterZstd::Compression::Highest);
             break;
         }
       }

diff --git a/Code/Engine/Foundation/IO/Archive/Implementation/ArchiveUtils.cpp b/Code/Engine/Foundation/IO/Archive/Implementation/ArchiveUtils.cpp
@@ -114,9 +114,12 @@ ezResult ezArchiveUtils::WriteEntry(
 
 #ifdef BUILDSYSTEM_ENABLE_ZSTD_SUPPORT
     case ezArchiveCompressionMode::Compressed_zstd:
-      zstdWriter.SetOutputStream(&inout_stream, (ezCompressedStreamWriterZstd::Compression)iCompressionLevel);
+    {
+      constexpr ezUInt32 uiMaxNumWorkerThreads = 12u;
+      zstdWriter.SetOutputStream(&inout_stream, uiMaxNumWorkerThreads, (ezCompressedStreamWriterZstd::Compression)iCompressionLevel);
       pWriter = &zstdWriter;
-      break;
+    }
+    break;
 #endif
 
     default:

diff --git a/Code/Engine/Foundation/IO/CompressedStreamZstd.h b/Code/Engine/Foundation/IO/CompressedStreamZstd.h
@@ -62,7 +62,7 @@ class EZ_FOUNDATION_DLL ezCompressedStreamWriterZstd final : public ezStreamWrit
 {
 public:
   /// \brief Specifies the compression level of the stream.
-  enum Compression
+  enum class Compression
   {
     Fastest = 1,
     Fast = 5,
@@ -76,7 +76,7 @@ class EZ_FOUNDATION_DLL ezCompressedStreamWriterZstd final : public ezStreamWrit
   ezCompressedStreamWriterZstd();
 
   /// \brief The constructor takes another stream writer to pass the output into, and a compression level.
-  ezCompressedStreamWriterZstd(ezStreamWriter* pOutputStream, Compression ratio = Compression::Default); // [tested]
+  ezCompressedStreamWriterZstd(ezStreamWriter* pOutputStream, ezUInt32 uiMaxNumWorkerThreads, Compression ratio = Compression::Default, ezUInt32 uiCompressionCacheSizeKB = 4); // [tested]
 
   /// \brief Calls FinishCompressedStream() internally.
   ~ezCompressedStreamWriterZstd(); // [tested]
@@ -92,7 +92,7 @@ class EZ_FOUNDATION_DLL ezCompressedStreamWriterZstd final : public ezStreamWrit
   /// another stream. This can prevent internal allocations, if one wants to use compression on multiple streams consecutively. It also
   /// allows to create a compressor stream early, but decide at a later pointer whether or with which stream to use it, and it will only
   /// allocate internal structures once that final decision is made.
-  void SetOutputStream(ezStreamWriter* pOutputStream, Compression ratio = Compression::Default, ezUInt32 uiCompressionCacheSizeKB = 4); // [tested]
+  void SetOutputStream(ezStreamWriter* pOutputStream, ezUInt32 uiMaxNumWorkerThreads, Compression ratio = Compression::Default, ezUInt32 uiCompressionCacheSizeKB = 4); // [tested]
 
   /// \brief Compresses \a uiBytesToWrite from \a pWriteBuffer.
   ///

diff --git a/Code/Engine/Foundation/IO/Implementation/CompressedStreamZstd.cpp b/Code/Engine/Foundation/IO/Implementation/CompressedStreamZstd.cpp
@@ -135,9 +135,9 @@ ezResult ezCompressedStreamReaderZstd::RefillReadCache()
 
 ezCompressedStreamWriterZstd::ezCompressedStreamWriterZstd() = default;
 
-ezCompressedStreamWriterZstd::ezCompressedStreamWriterZstd(ezStreamWriter* pOutputStream, Compression ratio)
+ezCompressedStreamWriterZstd::ezCompressedStreamWriterZstd(ezStreamWriter* pOutputStream, ezUInt32 uiMaxNumWorkerThreads, Compression ratio /*= Compression::Default*/, ezUInt32 uiCompressionCacheSizeKB /*= 4*/)
 {
-  SetOutputStream(pOutputStream, ratio);
+  SetOutputStream(pOutputStream, uiMaxNumWorkerThreads, ratio, uiCompressionCacheSizeKB);
 }
 
 ezCompressedStreamWriterZstd::~ezCompressedStreamWriterZstd()
@@ -159,7 +159,7 @@ ezCompressedStreamWriterZstd::~ezCompressedStreamWriterZstd()
   }
 }
 
-void ezCompressedStreamWriterZstd::SetOutputStream(ezStreamWriter* pOutputStream, Compression ratio /*= Compression::Default*/, ezUInt32 uiCompressionCacheSizeKB /*= 4*/)
+void ezCompressedStreamWriterZstd::SetOutputStream(ezStreamWriter* pOutputStream, ezUInt32 uiMaxNumWorkerThreads, Compression ratio /*= Compression::Default*/, ezUInt32 uiCompressionCacheSizeKB /*= 4*/)
 {
   if (m_pOutputStream == pOutputStream)
     return;
@@ -183,7 +183,7 @@ void ezCompressedStreamWriterZstd::SetOutputStream(ezStreamWriter* pOutputStream
       m_pZstdCStream = ZSTD_createCStream();
     }
 
-    const ezUInt32 uiCoreCount = ezMath::Clamp(ezSystemInformation::Get().GetCPUCoreCount(), 1u, 12u);
+    const ezUInt32 uiCoreCount = (uiMaxNumWorkerThreads > 0) ? ezMath::Clamp(ezSystemInformation::Get().GetCPUCoreCount(), 1u, uiMaxNumWorkerThreads) : 0u;
 
     ZSTD_CCtx_reset(reinterpret_cast<ZSTD_CStream*>(m_pZstdCStream), ZSTD_reset_session_only);
     ZSTD_CCtx_refCDict(reinterpret_cast<ZSTD_CStream*>(m_pZstdCStream), nullptr);